import os
import json
import logging
from anthropic import Anthropic
from base_telegram_inference_bot import BaseTelegramInferenceBot

class AnthropicTelegramInferenceBot(BaseTelegramInferenceBot):
    def __init__(self):
        super().__init__()
        self.anthropic_client = Anthropic(
            api_key=os.environ.get("ANTHROPIC_API_KEY"),
            default_headers={"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"}
        )

    def get_chat_response(self, messages):
        anthropic_tools = [
            {
                "name": function['name'],
                "description": function['description'],
                "input_schema": function['parameters'] if function['parameters'] not in [None, {}] else {"type": "object", "properties": {"param1": {"type": "string", "description": "Unnecessary"}}, "required": []}
            }
            for function in self.functions
        ]
        try:
            response = self.anthropic_client.messages.create(
                model="claude-3-5-sonnet-20240620",
                system=self.system_prompt,
                messages=messages,
                max_tokens=8192,
                tools=anthropic_tools
            )
        except Exception as e:
            logging.error(f"An error occurred: {str(e)}")
            return None
            
        return response

    async def handle_message(self, user_id, user_message):
        if user_id not in self.conversation_history:
            self.conversation_history[user_id] = []

        self.conversation_history[user_id].append({"role": "user", "content": user_message})
        messages = self.conversation_history[user_id]

        response = self.get_chat_response(messages)
        tool_calls = []
        full_message = []
        for message_part in response.content:
            full_message.append(message_part)
            if message_part.type == "tool_use":
                tool_calls.append(message_part)
        messages.append({"role": "assistant", "content": full_message})

        tool_use_count = 0
        while len(tool_calls) > 0 and tool_use_count < 50:
            tool_use_results = []
            for tool_call in tool_calls:
                tool_response = self.call_tool(tool_call)
                tool_use_results.append({"type": "tool_result", "tool_use_id": tool_call.id, "content": json.dumps(tool_response)})

            messages.append({"role": "user", "content": tool_use_results})

            response = self.get_chat_response(messages)
            full_message = []
            tool_calls = []
            for message_part in response.content:
                full_message.append(message_part)
                if message_part.type == "tool_use":
                    tool_calls.append(message_part)
            messages.append({"role": "assistant", "content": full_message})

            tool_use_count += 1

        if len(self.conversation_history[user_id]) > 20:
            self.conversation_history[user_id] = self.conversation_history[user_id][-20:]

        return messages[-1]["content"][0].text

    async def start(self):
        logging.info("Bot started")

    async def clear(self, user_id):
        super().clear_conversation(user_id)
        logging.info(f"Cleared conversation history and image for user {user_id}")

    async def status(self):
        return "Currently using claude-3-5-sonnet-20240620"

    async def abort_processing(self, user_id):
        if user_id in self.processing_status:
            self.processing_status[user_id]["processing"] = False
            await self.clear(user_id)
            return "Processing aborted."
        else:
            return "No active processing to abort."

def main():
    bot = AnthropicTelegramInferenceBot()
    telegram_helper = TelegramHelper(bot)
    telegram_helper.run()

if __name__ == '__main__':
    main()