#!/usr/bin/env python3
"""
Test a basic agent using ALL runners and ways of calling uni-api.
"""

import sys
sys.path.insert(0, '/home/GOD/heaven-base')
import asyncio
from heaven_base.baseheavenagent import BaseHeavenAgentReplicant, HeavenAgentConfig
from heaven_base.baseheaventool import BaseHeavenTool, ToolArgsSchema
from heaven_base.unified_chat import ProviderEnum
from typing import Dict, Any

# Create a simple test tool
class TestToolArgsSchema(ToolArgsSchema):
    arguments: Dict[str, Dict[str, Any]] = {
        'message': {
            'name': 'message',
            'type': 'str',
            'description': 'A test message to echo back',
            'required': True
        }
    }

def test_tool_func(message: str) -> str:
    """Simple test tool that echoes back a message."""
    return f"Test tool received: {message}"

class TestTool(BaseHeavenTool):
    name = "TestTool"
    description = "A simple test tool that echoes back messages"
    func = test_tool_func
    args_schema = TestToolArgsSchema
    is_async = False

# Create a basic test agent
class BasicTestAgent(BaseHeavenAgentReplicant):
    """A basic test agent for testing all runners."""
    
    @classmethod
    def get_default_config(cls) -> HeavenAgentConfig:
        return HeavenAgentConfig(
            name="BasicTestAgent",
            system_prompt="You are a test agent. Use the TestTool to echo back any message you receive. IMPORTANT: Only call ONE tool at a time. Never make multiple tool calls in a single response.",
            tools=[TestTool],
            provider=ProviderEnum.OPENAI,
            model="gpt-4.1",
            temperature=0.1
        )

async def test_all_runners():
    print("=== Testing Basic Agent with ALL Runners ===")
    
    results = {}
    
    # Test 1: Basic run() method
    print("\n1. Testing basic run() method...")
    try:
        agent1 = BasicTestAgent()
        result1 = await agent1.run("Please use TestTool to echo back 'Hello from basic run'")
        results['basic_run'] = {'success': True, 'history_id': result1['history_id']}
        print(f"✅ Basic run SUCCESS - History ID: {result1['history_id']}")
    except Exception as e:
        results['basic_run'] = {'success': False, 'error': str(e)}
        print(f"❌ Basic run FAILED: {e}")
    
    # Test 2: Agent mode with iterations
    print("\n2. Testing agent mode...")
    try:
        agent2 = BasicTestAgent()
        result2 = await agent2.run('agent goal="Use TestTool with message: Hello from agent mode", iterations=2')
        results['agent_mode'] = {'success': True, 'history_id': result2['history_id']}
        print(f"✅ Agent mode SUCCESS - History ID: {result2['history_id']}")
    except Exception as e:
        results['agent_mode'] = {'success': False, 'error': str(e)}
        print(f"❌ Agent mode FAILED: {e}")
    
    # Test 3: Multiple tool calls
    print("\n3. Testing multiple tool calls...")
    try:
        agent3 = BasicTestAgent()
        result3 = await agent3.run("Use TestTool twice: first with 'First message' then with 'Second message'")
        results['multiple_tools'] = {'success': True, 'history_id': result3['history_id']}
        print(f"✅ Multiple tools SUCCESS - History ID: {result3['history_id']}")
    except Exception as e:
        results['multiple_tools'] = {'success': False, 'error': str(e)}
        print(f"❌ Multiple tools FAILED: {e}")
    
    # Test 4: Direct uni-api usage (if available)
    print("\n4. Testing direct uni-api calls...")
    try:
        agent4 = BasicTestAgent()
        # Force uni-api usage
        agent4.config.provider = ProviderEnum.OPENAI
        result4 = await agent4.run("Use TestTool to echo 'Testing direct uni-api'")
        results['direct_uni_api'] = {'success': True, 'history_id': result4['history_id']}
        print(f"✅ Direct uni-api SUCCESS - History ID: {result4['history_id']}")
    except Exception as e:
        results['direct_uni_api'] = {'success': False, 'error': str(e)}
        print(f"❌ Direct uni-api FAILED: {e}")
    
    # Test 5: Long conversation
    print("\n5. Testing long conversation...")
    try:
        agent5 = BasicTestAgent()
        # Add some conversation history first
        result5a = await agent5.run("Use TestTool to echo 'First turn'")
        result5b = await agent5.run("Use TestTool to echo 'Second turn'")
        result5c = await agent5.run("Use TestTool to echo 'Third turn'")
        results['long_conversation'] = {'success': True, 'history_id': result5c['history_id']}
        print(f"✅ Long conversation SUCCESS - History ID: {result5c['history_id']}")
    except Exception as e:
        results['long_conversation'] = {'success': False, 'error': str(e)}
        print(f"❌ Long conversation FAILED: {e}")
    
    # Test 6: Error handling
    print("\n6. Testing error handling...")
    try:
        agent6 = BasicTestAgent()
        result6 = await agent6.run("Try to use a NonExistentTool that doesn't exist")
        results['error_handling'] = {'success': True, 'history_id': result6['history_id']}
        print(f"✅ Error handling SUCCESS - History ID: {result6['history_id']}")
    except Exception as e:
        results['error_handling'] = {'success': False, 'error': str(e)}
        print(f"❌ Error handling FAILED: {e}")
    
    # Summary
    print("\n" + "="*50)
    print("SUMMARY OF ALL TESTS:")
    print("="*50)
    
    total_tests = len(results)
    passed_tests = sum(1 for r in results.values() if r['success'])
    
    for test_name, result in results.items():
        status = "✅ PASS" if result['success'] else "❌ FAIL"
        print(f"{test_name:20} {status}")
        if not result['success']:
            print(f"                     Error: {result['error'][:100]}...")
    
    print(f"\nOverall: {passed_tests}/{total_tests} tests passed")
    
    if passed_tests == total_tests:
        print("🎉 ALL TESTS PASSED!")
    else:
        print("💥 SOME TESTS FAILED!")
    
    return results

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