import os
from dotenv import load_dotenv
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    PythonInterpreterTool,
    VisitWebpageTool,
    LiteLLMModel,
)
from tools import (
    multiply,
    add,
    subtract,
    divide,
    modulus,
    get_wikipedia_summary
)
# Load environment variables
load_dotenv()


def create_gaia_agent():
    """
    Create a GAIA question-answering agent using Gemini API and smolagents tools
    """

    # Initialize the Gemini model
    model = LiteLLMModel(
        model_id="gemini/gemini-2.0-flash",
        api_key=os.getenv("GEMINI_API_KEY")
    )

    # Initialize tools for the agent
    tools = [
        DuckDuckGoSearchTool(),
        PythonInterpreterTool(),
        VisitWebpageTool(),
        get_wikipedia_summary,
        multiply,
        add,
        subtract,
        divide,
        modulus
    ]

    # Create the agent with iterative decision-making capabilities
    agent = CodeAgent(
        tools=tools,
        model=model
    )

    return agent


def answer_gaia_question(agent: CodeAgent, question: str):
    """
    Answer a GAIA question using the agent

    Args:
        question (str): The GAIA question to answer

    Returns:
        str: The agent's answer
    """

    # Enhanced prompt for GAIA questions
    enhanced_prompt = f"""
    You are an expert AI assistant designed to answer complex GAIA (General AI Assistant) questions.
    These questions often require:
    - Multi-step reasoning
    - Web searches for current information
    - Mathematical calculations
    - Data analysis
    - Combining information from multiple sources

    Question: {question}

    Please approach this systematically:
    1. Break down the question into components
    2. Identify what information you need
    3. Use appropriate tools to gather information
    4. Perform any necessary calculations
    5. Synthesize your findings into a clear answer

    Make sure to:
    - Use web search when you need current or specific information
    - Use Python for any calculations or data processing
    - Visit specific webpages if you need detailed information from particular sources
    - Think step by step and show your reasoning

    Provide a final, clear answer at the end.
    """

    try:
        # Run the agent
        result = agent.run(enhanced_prompt)
        return result

    except Exception as e:
        return f"Error occurred while processing the question: {str(e)}"




# Example usage for specific GAIA questions
def example_usage():
    """
    Example of how to use the agent with specific GAIA questions
    """

    # Example GAIA questions (you can replace with actual ones)
    example_questions = [
        "What is the current population of Tokyo and how does it compare to New York City?",
    ]

    print("Running example GAIA questions:")
    print("=" * 40)

    for i, question in enumerate(example_questions, 1):
        print(f"\nExample {i}: {question}")
        answer = answer_gaia_question(question)
        print(f"Answer: {answer}")
        print("-" * 40)


if __name__ == "__main__":
    example_usage()