#!/usr/bin/env python3
"""
Test script for GAIA Agent
Run this to verify your agent works before deploying
"""

import os
import sys
from pathlib import Path

# Add current directory to path
sys.path.append(str(Path(__file__).parent))

def test_environment():
    """Test environment variables and dependencies"""
    print("🧪 Testing Environment Setup")
    print("-" * 40)
    
    # Check environment variables
    serper_key = os.getenv("SERPER_API_KEY")
    hf_token = os.getenv("HUGGINGFACE_INFERENCE_TOKEN")
    
    print(f"SERPER_API_KEY: {'✅ Found' if serper_key else '❌ Missing'}")
    print(f"HF_TOKEN: {'✅ Found' if hf_token else '❌ Missing'}")
    
    # Test imports
    try:
        import gradio as gr
        print("Gradio: ✅ Imported")
    except ImportError as e:
        print(f"Gradio: ❌ Import failed - {e}")
    
    try:
        import smolagents
        print("SmolagentS: ✅ Imported")
    except ImportError as e:
        print(f"SmolagentS: ❌ Import failed - {e}")
    
    try:
        import pandas as pd
        print("Pandas: ✅ Imported")
    except ImportError as e:
        print(f"Pandas: ❌ Import failed - {e}")
    
    try:
        import requests
        print("Requests: ✅ Imported")
    except ImportError as e:
        print(f"Requests: ❌ Import failed - {e}")

def test_agent_basic():
    """Test basic agent functionality"""
    print("\n🤖 Testing Agent Initialization")
    print("-" * 40)
    
    try:
        # Import the agent
        from app import GAIAAgent
        
        # Initialize agent
        agent = GAIAAgent()
        
        if agent.agent is None:
            print("❌ Agent initialization failed")
            return False
        
        print("✅ Agent initialized successfully")
        
        # Test with simple questions
        test_questions = [
            "What is 2 + 2?",
            "What is the capital of France?",
            "Calculate the square root of 16"
        ]
        
        for i, question in enumerate(test_questions, 1):
            print(f"\n📝 Test Question {i}: {question}")
            try:
                answer = agent(question)
                print(f"✅ Answer: {answer[:100]}...")
            except Exception as e:
                print(f"❌ Error: {e}")
        
        return True
        
    except Exception as e:
        print(f"❌ Agent test failed: {e}")
        return False

def test_tools():
    """Test individual tools"""
    print("\n🛠️ Testing Individual Tools")
    print("-" * 40)
    
    try:
        from app import SerperSearchTool, MathCalculatorTool
        
        # Test search tool
        search_tool = SerperSearchTool()
        try:
            result = search_tool("Python programming")
            print(f"✅ Search Tool: {result[:100]}...")
        except Exception as e:
            print(f"❌ Search Tool Error: {e}")
        
        # Test math tool
        math_tool = MathCalculatorTool()
        try:
            result = math_tool("2 + 2")
            print(f"✅ Math Tool: {result}")
        except Exception as e:
            print(f"❌ Math Tool Error: {e}")
        
        # Test math tool with complex expression
        try:
            result = math_tool("sqrt(16) + 3 * 2")
            print(f"✅ Math Complex: {result}")
        except Exception as e:
            print(f"❌ Math Complex Error: {e}")
            
    except Exception as e:
        print(f"❌ Tools test failed: {e}")

def main():
    """Run all tests"""
    print("🚀 GAIA Agent Test Suite")
    print("=" * 50)
    
    # Test environment
    test_environment()
    
    # Test tools
    test_tools()
    
    # Test agent
    success = test_agent_basic()
    
    print("\n" + "=" * 50)
    if success:
        print("✅ All tests passed! Your agent is ready for deployment.")
    else:
        print("❌ Some tests failed. Please check the errors above.")
    print("=" * 50)

if __name__ == "__main__":
    main()