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

class GeminiTelegramInferenceBot(OpenAICompatibleInferenceBot):
    def __init__(self):
        super().__init__()
        self.client = OpenAI(api_key=os.environ.get("GEMINI_API_KEY"), base_url=os.environ.get("GEMINI_API_BASE_URL"))
        
        self._configure_model_and_tokens(
            os.environ.get("GEMINI_SMALL_MODEL", "gemini-pro"),
            os.environ.get("GEMINI_SMALL_MODEL_MAX_TOKENS")
        )

    async def switch_model(self):
        current_small_model = os.environ.get("GEMINI_SMALL_MODEL", "gemini-pro")
        current_large_model = os.environ.get("GEMINI_LARGE_MODEL", "gemini-1.5-pro-latest")

        if self.model == current_large_model or self.model != current_small_model :
            target_model = current_small_model
            target_max_tokens = os.environ.get("GEMINI_SMALL_MODEL_MAX_TOKENS")
        else:
            target_model = current_large_model
            target_max_tokens = os.environ.get("GEMINI_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("GEMINI_API_KEY"):
        logging.error("FATAL: GEMINI_API_KEY environment variable not set.")
        return

    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    
    bot = GeminiTelegramInferenceBot()
    telegram_helper = TelegramHelper(bot) 
    telegram_helper.run()

if __name__ == '__main__':
    main()