import os
import logging
from openai import OpenAI
from openai_compatible_inference_bot import OpenAICompatibleInferenceBot
from telegram_helper import TelegramHelper 

class ChatGPTTelegramInferenceBot(OpenAICompatibleInferenceBot):
    def __init__(self):
        super().__init__()
        self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
        
        self._configure_model_and_tokens(
            os.environ.get("OPENAI_SMALL_MODEL", "gpt-3.5-turbo"),
            os.environ.get("OPENAI_SMALL_MODEL_MAX_TOKENS") 
        )

    async def switch_model(self):
        current_small_model = os.environ.get("OPENAI_SMALL_MODEL", "gpt-3.5-turbo")
        current_large_model = os.environ.get("OPENAI_LARGE_MODEL", "gpt-4")

        if self.model == current_large_model or self.model != current_small_model:
            target_model = current_small_model
            target_max_tokens = os.environ.get("OPENAI_SMALL_MODEL_MAX_TOKENS")
        else:
            target_model = current_large_model
            target_max_tokens = os.environ.get("OPENAI_LARGE_MODEL_MAX_TOKENS")
        
        self._configure_model_and_tokens(target_model, target_max_tokens)
        logging.info(f"Switched to model: {self.model}")
        return f"Switched to model: {self.model}"

def main():
    if not os.environ.get("OPENAI_API_KEY"):
        logging.error("FATAL: OPENAI_API_KEY environment variable not set.")
        return
    
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    
    bot = ChatGPTTelegramInferenceBot()
    telegram_helper = TelegramHelper(bot) 
    telegram_helper.run()

if __name__ == '__main__':
    main()