import logging
from textwrap import dedent
from agno.models.openai import OpenAILike
from agno.tools.arxiv import ArxivTools
from agno.tools.pubmed import PubmedTools
from agno.agent import Agent
import os
import random

API_KEYS = [
    os.getenv("SK1"),
    os.getenv("SK2"),
    os.getenv("SK3"),
    os.getenv("SK4"),
    os.getenv("SK5")
]
class ModelHandler:
    def __init__(self):
        """Initialize the model handler"""
        self.model = None
        self.tokenizer = None
        self.translator = None
        self.researcher = None
        self.summarizer = None
        self.presenter = None
        self._initialize_model()
    
    def _initialize_model(self):
        """Initialize model and tokenizer"""
        self.translator = Agent(
            name="Translator",
            role="You will translate the query to English",
            model=OpenAILike(
                api_key=str(random.choice(API_KEYS)),
                base_url="https://api.moonshot.cn/v1",
                id="moonshot-v1-8k",
                instructions=dedent("""\
                    Translate the query to English
                """)
            )
        ) 
        
        self.researcher = Agent(
            name="Researcher",
            role="You are a research scholar who specializes in autism research.",
            model=OpenAILike(
                api_key=str(random.choice(API_KEYS)),
                base_url="https://api.moonshot.cn/v1",
                id="moonshot-v1-8k",
                instructions=dedent("""\
                    - You have ArxivTools and PubmedTools at your disposal. Use them to find relevant papers for the question.
                    - You need to understand the context of the question to provide the best answer based on your tools.
                    - Be precise and provide just enough information to be useful.
                    - You must cite the sources used in your answer.
                    - You must create an accessible summary.
                    - The content must be for people without autism knowledge.
                    - Focus in the main findings of the paper taking in consideration the question.
                    - The answer must be brief.
                """),
                show_tool_calls=True
            ),
            tools=[ArxivTools(), PubmedTools()]
        )
        self.summarizer = Agent(
            name="Summarizer",
            role="You are a specialist in summarizing research papers for people without autism knowledge.",
            model=OpenAILike(
                api_key=str(random.choice(API_KEYS)),
                base_url="https://api.moonshot.cn/v1",
                id="moonshot-v1-8k",
                instructions=dedent("""\
                    - You must provide just enough information to be useful
                    - You must cite the sources used in your answer.
                    - You must be clear and concise.
                    - You must create an accessible summary.
                    - The content must be for people without autism knowledge.
                    - Focus in the main findings of the paper taking in consideration the question.
                    - The answer must be brief.
                    - Remove everything related to the run itself like: 'Running: transfer_' just use plain text
                    - You must use the language provided by the user to present the results.
                    - Add references to the sources used in the answer.
                    - Add emojis to make the presentation more interactive.
                    - Translate the answer to Portuguese.
                """),
            )
        )
        
        self.presenter = Agent(
            name="Presenter",
            role="You are a professional researcher who presents the results of the research.",
            model=OpenAILike(
                api_key=str(random.choice(API_KEYS)),
                base_url="https://api.moonshot.cn/v1",
                id="moonshot-v1-8k",
                instructions=dedent("""\
                    - You are multilingual
                    - You must present the results in a clear and concise manner.
                    - Clean up the presentation to make it more readable.
                    - Remove unnecessary information.
                    - Remove everything related to the run itself like: 'Running: transfer_', just use plain text
                    - You must use the language provided by the user to present the results.
                    - Add references to the sources used in the answer.
                    - Add emojis to make the presentation more interactive.
                    - Translate the answer to Portuguese.
                """),
            )
        )
        
    
    def generate_answer(self, query: str) -> str:
        try:
            translator = self.translator.run(query, stream=False)
            logging.info(f"Translated query")
            research = self.researcher.run(translator.content, stream=False)
            logging.info(f"Generated research")
            summary = self.summarizer.run(research.content, stream=False)
            logging.info(f"Generated summary")
            presentation = self.presenter.run(summary.content, stream=False)
            logging.info(f"Generated presentation")
            
            if not presentation.content:
                return self._get_fallback_response()
            return presentation.content
        except Exception as e:
            logging.error(f"Error generating answer: {str(e)}")
            return self._get_fallback_response()
    
    @staticmethod
    def _get_fallback_response() -> str:
        """Provide a friendly, helpful fallback response"""
        return """
            Peço descula, mas encontrei um erro ao gerar a resposta. Tente novamente ou refaça a sua pergunta.
        """