import os
import pathlib

import dotenv
from smolagents import CodeAgent, OpenAIServerModel

if pathlib.Path(".env").exists():
    dotenv.load_dotenv(".env")

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", None)

if not OPENAI_API_KEY:
    exit(1)

model = OpenAIServerModel(
    model_id="gpt-4o-mini",
    api_key=OPENAI_API_KEY,
)

def test_simple_calculation():
    from prompt import calc_agent_prompt
    
    calc_agent = CodeAgent(
        model=model,
        prompt_templates=calc_agent_prompt,
        additional_authorized_imports=["pandas", "numpy"],
        tools=[]
    )
    
    question = "What is the result of 5 + 3 + 1294.678?"
    print(f"\n=== Testing Calculation Agent ===")
    print(f"Question: {question}")
    print("Running agent...")
    
    try:
        result = calc_agent.run(question)
        print(f"Result: {result}")
        print("✅ Calculation agent test passed!")
        return True
    except Exception as e:
        print(f"❌ Calculation agent test failed: {e}")
        return False

def test_simple_search():
    from smolagents import VisitWebpageTool, WebSearchTool, WikipediaSearchTool

    from prompt import web_search_agent_prompt
    
    web_agent = CodeAgent(
        model=model,
        prompt_templates=web_search_agent_prompt,
        additional_authorized_imports=["requests", "beautifulsoup4"],
        tools=[
            WebSearchTool(),
            VisitWebpageTool(),
            WikipediaSearchTool(),
        ],
    )
    
    question = "Who is the current president of France?"
    print(f"\n=== Testing Web Search Agent ===")
    print(f"Question: {question}")
    print("Running agent...")
    
    try:
        result = web_agent.run(question)
        print(f"Result: {result}")
        print("✅ Web search agent test passed!")
        return True
    except Exception as e:
        print(f"❌ Web search agent test failed: {e}")
        return False

if __name__ == "__main__":
    print("Testing our updated agents with final_answer() support...")
    
    calc_success = test_simple_calculation()
    search_success = test_simple_search()
    
    if calc_success and search_success:
        print("\n🎉 All tests passed! Agents are working correctly with final_answer()")
    else:
        print("\n⚠️  Some tests failed. Check the error messages above.")