import os
from smolagents import InferenceClientModel, ToolCallingAgent
from smolagents import DuckDuckGoSearchTool, VisitWebpageTool
from gaia_tools import RunPythonFileTool, WikipediaSearchTool, ReverseTextTool

class GaiaAgent:
    def __init__(self):
        print("Gaia Agent Initialized")
        
        self.model = InferenceClientModel(
            model_id="meta-llama/Meta-Llama-3-70B-Instruct",
            token=os.getenv("HUGGINGFACEHUB_API_TOKEN")
        )

        self.tools = [
            DuckDuckGoSearchTool(),
            VisitWebpageTool(),
            RunPythonFileTool,
            WikipediaSearchTool,
            ReverseTextTool
        ]

        self.agent = ToolCallingAgent(
            tools = self.tools,
            model = self.model
        )

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")

        prompt = f"""
        You are a helpful agent that must provide exact answers to questions. Do not explain or format your answer in any way.

        CRITICAL: If the question starts with a period or looks backwards, use ReverseTextTool to reverse it first.

        For Wikipedia research:
        - ALWAYS search for the main Wikipedia page of the subject first
        - Use WikipediaSearchTool with the exact name (e.g., "Mercedes Sosa")
        - Look specifically in the "Discography" or "Albums" section
        - Count only items explicitly labeled as "studio albums"
        - Exclude live albums, compilation albums, or singles
        - For Featured Articles, search "Wikipedia Featured Articles [month] [year]"

        For text puzzles:
        - If reversed, use ReverseTextTool then solve the resulting question
        - Simple word/logic puzzles can be solved directly

        Question: {question}

        SEARCH CONSTRAINTS:
        - Use exact names and specific Wikipedia sections
        - Be precise about album types (studio vs. live vs. compilation)
        - For date ranges, include both start and end years
        - Always verify information from the main Wikipedia article

        Only output the final answer (number, word, or name).
        """

        try:
            result = self.agent.run(prompt)
            print(f"Raw result from agent: {result}")
            print(f"Result type: {type(result)}")

            # Handle different result types more robustly
            if isinstance(result, dict):
                if "answer" in result:
                    return str(result["answer"]).strip()
                elif "content" in result:
                    return str(result["content"]).strip()
                elif len(result) == 0:
                    # Empty dictionary - likely a tool calling issue
                    return "ERROR: Empty response from agent"
                else:
                    # Try to find any string value in the dict
                    for key, value in result.items():
                        if isinstance(value, str) and value.strip():
                            return value.strip()
                    return f"ERROR: No valid content in response: {result}"
            elif isinstance(result, str):
                return result.strip()
            elif isinstance(result, list):
                # Look for assistant messages in conversation format
                for item in reversed(result):
                    if isinstance(item, dict):
                        if item.get("role") == "assistant" and "content" in item:
                            return item["content"].strip()
                        elif "content" in item:
                            return item["content"].strip()
                # If no assistant message found, try to extract any string
                for item in reversed(result):
                    if isinstance(item, str) and item.strip():
                        return item.strip()
                return str(result)
            elif isinstance(result, (int, float)):
                return str(result)
            else:
                return str(result)
        except Exception as e:
            print(f"Exception during agent run: {e}")
            return f"AGENT ERROR: {e}"