#!/usr/bin/env python3
"""
Test script for the LangGraph multi-agent system with LangChain tools
"""

import asyncio
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv("env.local")

async def test_langgraph_system():
    """Test the LangGraph system with a simple question"""
    
    print("🔧 Testing LangGraph System with LangChain Tools")
    print("=" * 60)
    
    try:
        # Import the main system
        from langgraph_agent_system import run_agent_system
        
        # Test with a simple computational question
        test_query = "What is 25 + 17?"
        print(f"📝 Test Query: {test_query}")
        print("-" * 40)
        
        # Run the agent system
        result = await run_agent_system(
            query=test_query,
            user_id="test_user", 
            session_id="test_session",
            max_iterations=2
        )
        
        print("\n📊 Final Result:")
        print(result)
        print("\n✅ Test completed successfully!")
        
    except Exception as e:
        print(f"❌ Test failed: {e}")
        import traceback
        traceback.print_exc()


async def test_research_tools():
    """Test the research tools separately"""
    
    print("\n🔍 Testing Research Tools")
    print("=" * 40)
    
    try:
        from langgraph_tools import get_research_tools
        
        # Get the tools
        tools = get_research_tools()
        print(f"✅ Loaded {len(tools)} research tools:")
        for tool in tools:
            print(f"  - {tool.name}: {tool.description}")
        
        # Test Wikipedia tool (if available)
        wiki_tool = next((t for t in tools if t.name == "wikipedia_search"), None)
        if wiki_tool:
            print("\n📚 Testing Wikipedia search...")
            result = wiki_tool.invoke({"query": "Python programming"})
            print(f"Wikipedia result length: {len(str(result))} characters")
            print(f"Preview: {str(result)[:200]}...")
        
    except Exception as e:
        print(f"❌ Research tools test failed: {e}")


async def test_code_tools():
    """Test the code tools separately"""
    
    print("\n🧮 Testing Code Tools")
    print("=" * 40)
    
    try:
        from langgraph_tools import get_code_tools
        
        # Get the tools
        tools = get_code_tools()
        print(f"✅ Loaded {len(tools)} code tools:")
        for tool in tools:
            print(f"  - {tool.name}: {tool.description}")
        
        # Test add tool
        add_tool = next((t for t in tools if t.name == "add"), None)
        if add_tool:
            print("\n➕ Testing addition...")
            result = add_tool.invoke({"a": 25, "b": 17})
            print(f"Addition result: {result}")
        
    except Exception as e:
        print(f"❌ Code tools test failed: {e}")


if __name__ == "__main__":
    async def main():
        await test_research_tools()
        await test_code_tools()
        await test_langgraph_system()
    
    asyncio.run(main())