#!/usr/bin/env python3
"""
Test HEAVEN's higher-level components: completion runners, LangGraph, and hermes utilities.
These components depend on BaseHeavenAgent working correctly.
"""

import sys
sys.path.insert(0, '/home/GOD/heaven-base')
import asyncio
from heaven_base.tool_utils.completion_runners import (
    exec_agent_run, 
    exec_completion_style
)
from heaven_base.baseheavenagent import HeavenAgentConfig
from heaven_base.unified_chat import ProviderEnum
from heaven_base.langgraph.foundation import (
    HeavenGraphBuilder,
    HeavenState,
    HeavenNodeType
)
from heaven_base.tool_utils.hermes_utils import hermes_step
from heaven_base.configs.hermes_config import HermesConfig

async def test_completion_runners():
    """Test the completion runner utilities"""
    print("=== Testing Completion Runners ===\n")
    
    results = {}
    
    # Test 1: exec_agent_run with string agent name
    print("1. Testing exec_agent_run with agent name...")
    try:
        result = await exec_agent_run(
            prompt="What is 2 + 2?",
            agent="BasicTestAgent"
        )
        results['exec_agent_run_string'] = {'success': True, 'response': str(result)[:100]}
        print(f"✅ SUCCESS - Got response: {str(result)[:100]}...")
    except Exception as e:
        results['exec_agent_run_string'] = {'success': False, 'error': str(e)}
        print(f"❌ FAILED: {e}")
    
    # Test 2: exec_completion_style
    print("\n2. Testing exec_completion_style...")
    try:
        config = HeavenAgentConfig(
            name="CompletionTestAgent",
            system_prompt="You are a helpful assistant. Answer concisely.",
            provider=ProviderEnum.OPENAI,
            model="gpt-4.1",
            temperature=0.1
        )
        
        result = await exec_completion_style(
            prompt="What is the capital of France?",
            agent=config
        )
        results['exec_completion_style'] = {'success': True, 'response': str(result)[:100]}
        print(f"✅ SUCCESS - Got response: {str(result)[:100]}...")
    except Exception as e:
        results['exec_completion_style'] = {'success': False, 'error': str(e)}
        print(f"❌ FAILED: {e}")
    
    # Test 3: exec_completion_style with history
    print("\n3. Testing exec_completion_style with history continuation...")
    try:
        # First call
        result1 = await exec_completion_style(
            prompt="Remember the number 42",
            agent=config
        )
        history_id = result1.get('history_id') if isinstance(result1, dict) else None
        
        # Continue with history
        result2 = await exec_completion_style(
            prompt="What number did I ask you to remember?",
            agent=config,
            history_id=history_id
        )
        results['exec_completion_with_history'] = {'success': True, 'has_history': history_id is not None}
        print(f"✅ SUCCESS - History continuation working")
    except Exception as e:
        results['exec_completion_with_history'] = {'success': False, 'error': str(e)}
        print(f"❌ FAILED: {e}")
    
    return results

async def test_langgraph_foundation():
    """Test LangGraph foundation components"""
    print("\n=== Testing LangGraph Foundation ===\n")
    
    results = {}
    
    # Test 1: Create a simple graph
    print("1. Testing HeavenGraphBuilder...")
    try:
        builder = HeavenGraphBuilder()
        
        # Add a simple completion node
        builder.add_node(
            name="test_completion",
            node_type=HeavenNodeType.COMPLETION,
            agent_config={
                "name": "TestCompletionAgent",
                "system_prompt": "You are a test agent. Answer briefly.",
                "provider": "openai",
                "model": "gpt-4.1"
            },
            prompt="What is 1 + 1?"
        )
        
        # Build and run
        graph = builder.build()
        initial_state = {
            "results": [],
            "context": {},
            "agents": {}
        }
        
        result = await graph.ainvoke(initial_state)
        results['graph_builder'] = {'success': True, 'num_results': len(result.get('results', []))}
        print(f"✅ SUCCESS - Graph executed with {len(result.get('results', []))} results")
    except Exception as e:
        results['graph_builder'] = {'success': False, 'error': str(e)}
        print(f"❌ FAILED: {e}")
    
    # Test 2: Test hermes_step
    print("\n2. Testing hermes_step...")
    try:
        hermes_config = HermesConfig(
            domain="Testing",
            subdomain="ComponentTest",
            prompt="Generate a test message",
            agent_config=HeavenAgentConfig(
                name="HermesTestAgent",
                system_prompt="You are a test agent for hermes.",
                provider=ProviderEnum.OPENAI,
                model="gpt-4.1"
            )
        )
        
        result = await hermes_step(hermes_config)
        results['hermes_step'] = {'success': True, 'has_result': result is not None}
        print(f"✅ SUCCESS - Hermes step completed")
    except Exception as e:
        results['hermes_step'] = {'success': False, 'error': str(e)}
        print(f"❌ FAILED: {e}")
    
    return results

async def test_auto_summarizing_agent():
    """Test the AutoSummarizingAgent which depends on base agent"""
    print("\n=== Testing AutoSummarizingAgent ===\n")
    
    results = {}
    
    try:
        from heaven_base.auto_summarizing_agent import AutoSummarizingAgent
        
        agent = AutoSummarizingAgent()
        
        # Create a long prompt that might trigger summarization
        long_prompt = "Please analyze this: " + " ".join([f"Item {i}" for i in range(100)])
        
        result = await agent.run(long_prompt)
        results['auto_summarizing'] = {
            'success': True, 
            'history_id': result.get('history_id'),
            'context_status': result.get('context_window_status', {})
        }
        print(f"✅ SUCCESS - AutoSummarizingAgent ran with history: {result.get('history_id')}")
    except Exception as e:
        results['auto_summarizing'] = {'success': False, 'error': str(e)}
        print(f"❌ FAILED: {e}")
    
    return results

async def main():
    """Run all component tests"""
    print("🧪 HEAVEN Higher-Level Components Test Suite")
    print("=" * 50)
    print("Testing components that depend on BaseHeavenAgent...\n")
    
    all_results = {}
    
    # Test completion runners
    runner_results = await test_completion_runners()
    all_results['completion_runners'] = runner_results
    
    # Test LangGraph foundation
    graph_results = await test_langgraph_foundation()
    all_results['langgraph'] = graph_results
    
    # Test AutoSummarizingAgent
    auto_results = await test_auto_summarizing_agent()
    all_results['auto_summarizing'] = auto_results
    
    # Summary
    print("\n" + "=" * 50)
    print("SUMMARY OF ALL COMPONENT TESTS:")
    print("=" * 50)
    
    total_tests = 0
    passed_tests = 0
    
    for category, results in all_results.items():
        print(f"\n{category.upper()}:")
        for test_name, result in results.items():
            total_tests += 1
            status = "✅ PASS" if result.get('success', False) else "❌ FAIL"
            if result.get('success', False):
                passed_tests += 1
            print(f"  {test_name:30} {status}")
            if not result.get('success', False):
                print(f"  {'':30} Error: {result.get('error', 'Unknown')[:70]}...")
    
    print(f"\nOverall: {passed_tests}/{total_tests} tests passed")
    
    if passed_tests == total_tests:
        print("🎉 ALL COMPONENT TESTS PASSED!")
    else:
        print("💥 SOME COMPONENT TESTS FAILED!")
    
    return all_results

if __name__ == "__main__":
    results = asyncio.run(main())