import os
from pathlib import Path
from typing import Optional, Union

import pandas as pd
from dotenv import load_dotenv
from smolagents import (CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool,
                        LiteLLMModel, PythonInterpreterTool,
                        WikipediaSearchTool)
from smolagents.tools import Tool
from tabulate import tabulate

# Load environment variables
load_dotenv()

# Initialize the model
model = LiteLLMModel(
    model_id=os.getenv("GEMINI_MODEL"), api_key=os.getenv("GEMINI_API_KEY")
)


class ExcelToTextTool(Tool):
    """Render an Excel worksheet as a Markdown table."""

    name = "excel_to_text"
    description = (
        "Read an Excel file and return a Markdown table of the requested sheet. "
        "Accepts either the sheet name or a zero-based index (as a string)."
    )

    inputs = {
        "excel_path": {
            "type": "string",
            "description": "Path to the Excel file (.xlsx or .xls).",
        },
        "sheet_name": {
            "type": "string",
            "description": (
                "Worksheet name or zero-based index (as a string). "
                "Optional; defaults to the first sheet."
            ),
            "nullable": True,
        },
    }

    output_type = "string"

    def forward(self, excel_path: str, sheet_name: Optional[str] = None) -> str:
        """Load the Excel file and return the sheet as a Markdown table.

        Args:
            excel_path: Path to the Excel file.
            sheet_name: Optional name or index of the sheet to read. If None, reads the first sheet.

        Returns:
            A Markdown table representing the Excel sheet, or an error message if the file is not found or cannot be read.
        """

        file_path = Path(excel_path).expanduser().resolve()
        if not file_path.is_file():
            return f"Error: Excel file not found at {file_path}"

        try:
            sheet: Union[str, int] = (
                int(sheet_name)
                if sheet_name and sheet_name.isdigit()
                else sheet_name or 0
            )

            df = pd.read_excel(file_path, sheet_name=sheet)

            if hasattr(df, "to_markdown"):
                return df.to_markdown(index=False)

            return tabulate(df, headers="keys", tablefmt="github", showindex=False)

        except Exception as e:
            return f"Error reading Excel file: {e}"


class GaiaAgent:
    """An agent capable of using tools to answer general questions."""

    def __init__(self):
        """Initializes the GaiaAgent with a set of tools."""

        print("GaiaAgent initialized with tools.")

        tools = [
            DuckDuckGoSearchTool(),
            WikipediaSearchTool(),
            ExcelToTextTool(),
            PythonInterpreterTool(),
            FinalAnswerTool(),
        ]

        self.agent = CodeAgent(
            model=model,
            tools=tools,
            add_base_tools=True,
            additional_authorized_imports=["pandas", "numpy", "csv", "subprocess"],
        )

    def __call__(self, task_id: str, question: str) -> str:
        """Processes a question using the agent and its tools.

        Args:
            task_id: A unique identifier for the task.
            question: The question to be answered.

        Returns:
            The answer generated by the agent.
        """
        print(f"Agent received task_id='{task_id}' | question='{question[:50]}...'")
        answer = self.agent.run(question)
        print(f"Agent returning answer: {answer}")
        return answer