#!/usr/bin/env python3
"""
Test Script for New LangGraph Agent System
Tests the multi-agent architecture with memory, routing, and verification.
"""
import os
import sys
import time
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Add the current directory to Python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

def test_imports():
    """Test that all modules can be imported correctly"""
    print("Testing imports...")
    try:
        # Test core imports
        from src import run_agent_system, memory_manager
        from src.tracing import get_langfuse_callback_handler
        
        # Test agent imports
        from src.agents import (
            plan_node, router_node, retrieval_agent, 
            execution_agent, critic_agent, verification_node
        )
        
        print("✅ All imports successful")
        return True
    except ImportError as e:
        print(f"❌ Import error: {e}")
        return False


def test_memory_system():
    """Test the memory management system"""
    print("\nTesting memory system...")
    try:
        from src.memory import memory_manager
        
        # Test basic functionality
        test_query = "What is 2+2?"
        
        # Test similarity search (should not crash even without vector store)
        similar = memory_manager.similarity_search(test_query, k=1)
        print(f"✅ Similarity search completed: {len(similar)} results")
        
        # Test cache management
        memory_manager.clear_session_cache()
        print("✅ Memory cache cleared")
        
        return True
    except Exception as e:
        print(f"❌ Memory system error: {e}")
        return False


def test_tracing_system():
    """Test the Langfuse tracing integration"""
    print("\nTesting tracing system...")
    try:
        from src.tracing import get_langfuse_callback_handler, initialize_langfuse
        
        # Test handler creation (should not crash even without credentials)
        handler = get_langfuse_callback_handler()
        print(f"✅ Langfuse handler: {type(handler)}")
        
        return True
    except Exception as e:
        print(f"❌ Tracing system error: {e}")
        return False


def test_individual_agents():
    """Test each agent individually"""
    print("\nTesting individual agents...")
    
    # Test state structure
    test_state = {
        "messages": [],
        "plan_complete": False,
        "next_agent": "",
        "routing_decision": "",
        "routing_reason": "",
        "current_step": "testing",
        "agent_response": None,
        "needs_tools": False,
        "execution_result": "",
        "critic_assessment": "",
        "quality_pass": True,
        "quality_score": 7,
        "verification_status": "",
        "attempt_count": 1,
        "final_answer": ""
    }
    
    try:
        from langchain_core.messages import HumanMessage
        test_state["messages"] = [HumanMessage(content="Test query")]
        
        # Test plan node
        from src.agents.plan_node import plan_node
        plan_result = plan_node(test_state)
        print("✅ Plan node executed")
        
        # Test router node  
        from src.agents.router_node import router_node
        router_result = router_node(plan_result)
        print("✅ Router node executed")
        
        return True
    except Exception as e:
        print(f"❌ Agent testing error: {e}")
        return False


def test_graph_creation():
    """Test the main graph creation"""
    print("\nTesting graph creation...")
    try:
        from src.langgraph_system import create_agent_graph
        
        # Create the workflow
        workflow = create_agent_graph()
        print("✅ Graph created successfully")
        
        # Try to compile (this might fail without proper setup, but shouldn't crash)
        try:
            app = workflow.compile()
            print("✅ Graph compiled successfully")
        except Exception as e:
            print(f"⚠️ Graph compilation warning: {e}")
        
        return True
    except Exception as e:
        print(f"❌ Graph creation error: {e}")
        return False


def test_simple_query():
    """Test a simple query through the system"""
    print("\nTesting simple query...")
    try:
        from new_langraph_agent import run_agent
        
        # Simple test query
        test_query = "What is 1 + 1?"
        print(f"Query: {test_query}")
        
        start_time = time.time()
        result = run_agent(test_query)
        end_time = time.time()
        
        print(f"Result: {result}")
        print(f"Time taken: {end_time - start_time:.2f} seconds")
        print("✅ Simple query completed")
        
        return True
    except Exception as e:
        print(f"❌ Simple query error: {e}")
        return False


def main():
    """Run all tests"""
    print("LangGraph Agent System - Test Suite")
    print("=" * 50)
    
    tests = [
        test_imports,
        test_memory_system,
        test_tracing_system, 
        test_individual_agents,
        test_graph_creation,
        test_simple_query
    ]
    
    results = []
    for test_func in tests:
        try:
            result = test_func()
            results.append(result)
        except Exception as e:
            print(f"❌ Test {test_func.__name__} failed with exception: {e}")
            results.append(False)
    
    # Summary
    print("\n" + "=" * 50)
    print("Test Summary:")
    print(f"Passed: {sum(results)}/{len(results)}")
    print(f"Failed: {len(results) - sum(results)}/{len(results)}")
    
    if all(results):
        print("🎉 All tests passed!")
        return 0
    else:
        print("⚠️ Some tests failed. Check the output above for details.")
        return 1


if __name__ == "__main__":
    sys.exit(main())