import os
from typing import Optional

from smolagents import LiteLLMModel, CodeAgent, DuckDuckGoSearchTool, PythonInterpreterTool, VisitWebpageTool

from agentcourse_unit4.tools.audio_transcriber import AudioTranscriberTool
from agentcourse_unit4.tools.chess_board_recognizer import ChessBoardRecognizerTool
from agentcourse_unit4.tools.chess_predictor import ChessPredictorTool
from agentcourse_unit4.tools.csv_reader import CsvReaderTool
from agentcourse_unit4.tools.excel_reader import ExcelReaderTool
from agentcourse_unit4.tools.file_downloader import FileDownloaderTool
from agentcourse_unit4.tools.image_describer import ImageDescriberTool
from agentcourse_unit4.tools.image_text_extractor import ImageTextExtractorTool
from agentcourse_unit4.tools.pdf_reader import PdfReaderTool
from agentcourse_unit4.tools.py_code_interpreter import PyCodeInterpreterTool
from agentcourse_unit4.tools.youtube_transcriber import YoutubeTranscriberTool


class BasicAgent:
    def __init__(self):
        self.model = LiteLLMModel(
           # model_id="gemini/gemini-2.0-flash",
            model_id="gemini/gemini-2.5-flash-preview-04-17",
            api_key=os.getenv("GOOGLE_API_KEY"),
            max_tokens=8196,
            temperature=0.9
        )
        self.basic_agent = CodeAgent(
            name="basic_agent",
            description="Basic code agent",
            tools=[
                PythonInterpreterTool(),
                DuckDuckGoSearchTool(max_results=5),
                VisitWebpageTool(max_output_length=1_000_000),
                FileDownloaderTool(),
                ExcelReaderTool(),
                CsvReaderTool(),
                PdfReaderTool(),
                PyCodeInterpreterTool(),
                YoutubeTranscriberTool(),
                AudioTranscriberTool(),
                ChessBoardRecognizerTool(),
                ChessPredictorTool(),
                ImageDescriberTool(),
                ImageTextExtractorTool()
            ],
            model=self.model,
            add_base_tools=False,
            additional_authorized_imports=["pandas", "numpy", "datetime", "json", "csv"],
            planning_interval=None,
            verbosity_level=1,
            max_steps=5,
            max_print_outputs_length=1_000_000
        )
        print("==> Agent initialized.")

    def run(self, question: str, file_path: Optional[str] = None) -> str:
        """
        Process the incoming question and then return the answer.

        Args:
            question: The question or task
            file_path: Optional path to a file associated with the question or task

        Returns:
            The final answer to the question
        """

        associated_file_prompt = f"\nAssociated file at path: {file_path}\n\nReading file content by proper tool is mandatory." if file_path else ''

        prompt = f"""
Question: 
\"{question}\"
{associated_file_prompt}

Remember that you answer to the question from GAIA benchmark. It requires short, exact and precise answer.
Don't include: thinking, explanations, steps, reasoning, intermediate or additional text.

Finish your answer with a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. 
If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.

For instance, if question is "What is the capital of Spain?", respond with "Madrid". 
It is exact and expected answer.
"""
        return self.basic_agent.run(prompt)