import json
import os
import logging
from base_telegram_inference_bot import BaseTelegramInferenceBot
from telegram_helper import TelegramHelper
from openai import OpenAI

class ChatGPTTelegramInferenceBot(BaseTelegramInferenceBot):
    def __init__(self):
        super().__init__()
        self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
        self.model = "gpt-4o-mini"
        self.max_tokens = 16384

    def get_chat_response(self, messages):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "system", "content": self.system_prompt}] + messages,
            functions=self.functions,
            function_call="auto",
            max_tokens=self.max_tokens
        )
        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)
        assistant_message = response.choices[0].message
        tool_calls = []
        if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None:
            tool_calls.append(assistant_message.function_call)

        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({"role": "function", "name": tool_call.name, "content": json.dumps(tool_response)})

            messages.extend(tool_use_results)

            response = self.get_chat_response(messages)
            assistant_message = response.choices[0].message
            messages.append({"role": "assistant", "content": assistant_message.content})
            tool_calls = []
            if hasattr(assistant_message, 'function_call') and assistant_message.function_call is not None:
                tool_calls.append(assistant_message.function_call)

            tool_use_count += 1

        if tool_use_count == 0:
            messages.append({"role": "assistant", "content": assistant_message.content})
        
        if len(self.conversation_history[user_id]) > 20:
            self.conversation_history[user_id] = self.conversation_history[user_id][-20:]

        return messages[-1]["content"]

    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 for user {user_id}")

    async def status(self):
        return f"Currently using: {self.model}"

    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."

    async def switch_model(self):
        if self.model == "gpt-4o-mini":
            self.model = "gpt-4o"
            self.max_tokens = 4096
        else:
            self.model = "gpt-4o-mini"
            self.max_tokens = 16384
        logging.info(f"Switched to model: {self.model}")
        return f"Switched to model: {self.model}"

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

if __name__ == '__main__':
    main()