#!/usr/bin/env python3
"""
Test the complete compactor layer agents system.

Tests each agent individually then the full pipeline:
1. IterationSummarizerAgent
2. AggregationSummarizerAgent  
3. ReasoningAgent
4. ConceptExtractorAgent
5. Complete auto_summarize pipeline
"""

import asyncio
import os
import sys
from datetime import datetime
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

# Add heaven-base to path
sys.path.insert(0, '/home/GOD/heaven-base')

from heaven_base.memory.history import History
from heaven_base.utils.auto_summarize import (
    IterationSummarizerAgent,
    AggregationSummarizerAgent, 
    ReasoningAgent,
    ConceptExtractorAgent,
    auto_summarize,
    reason_about_summary,
    get_summary_concept,
    get_summary_reasoning_and_concept
)


def create_test_history() -> History:
    """Create a realistic test history for summarization."""
    history = History(messages=[
        SystemMessage(content="You are a helpful coding assistant."),
        
        # Iteration 1: User asks for help, agent analyzes
        HumanMessage(content="I need help implementing a token counter for my chat system. It should work with different AI models."),
        AIMessage(content="I'll help you implement a token counter. Let me create a utility that works with tiktoken for OpenAI models and provides fallbacks for others."),
        
        # Iteration 2: Implementation discussion
        HumanMessage(content="Great! Can you show me how to integrate this with a context window management system?"),
        AIMessage(content="Absolutely! I'll create a ContextWindowConfig class that tracks token usage and manages buffer allocations for different purposes like workspace, summaries, and padding."),
        
        # Iteration 3: Testing and refinement
        HumanMessage(content="Perfect! Now I need to test this with actual conversations to make sure the buffer triggers work correctly."),
        AIMessage(content="Let me create a comprehensive test that fills the workspace buffer and verifies the compaction triggers fire at exactly the right token limits. This will ensure your context management works reliably."),
        
        # Iteration 4: Integration
        HumanMessage(content="Excellent! The tests pass. Now how do I integrate this with my agent that needs auto-summarization?"),
        AIMessage(content="I'll show you how to integrate this with AutoSummarizingAgent. It will automatically trigger compaction when the workspace buffer fills up, create summaries, and inject them back into the context while maintaining the engineered buffer allocations.")
    ])
    
    # Save the history
    history_id = history.save("test_compactor_agents")
    print(f"Created test history: {history_id}")
    
    return history


async def test_iteration_summarizer():
    """Test IterationSummarizerAgent with IterationSummarizerTool."""
    print("\n=== Testing IterationSummarizerAgent ===")
    
    agent = IterationSummarizerAgent()
    
    # Test iteration 1 summary
    goal = """Summarize iteration 1 using the IterationSummarizerTool:
    - Actions taken: Created token counter utility with tiktoken integration
    - Outcomes: Working token counter for OpenAI models with fallbacks
    - Challenges: Handling different model types and API variations
    - Tools used: tiktoken library, model configuration system"""
    
    prompt = f'agent goal="{goal}", iterations=3'
    
    print("Running IterationSummarizerAgent...")
    result = await agent.run(prompt)
    
    print(f"Agent run complete. History ID: {result.get('history_id')}")
    
    # Look for tool calls
    agent.look_for_particular_tool_calls()
    
    if agent.last_summary:
        print("✅ IterationSummarizerAgent working!")
        print(f"Summary preview: {agent.last_summary[:200]}...")
        return agent.last_summary
    else:
        print("❌ IterationSummarizerAgent failed - no summary generated")
        return None


async def test_aggregation_summarizer():
    """Test AggregationSummarizerAgent with AggregationSummarizerTool."""
    print("\n=== Testing AggregationSummarizerAgent ===")
    
    agent = AggregationSummarizerAgent()
    
    # Mock iteration summaries
    iteration_summaries = """
## Iteration 1 Summary
**Actions Taken:** Created token counter utility with tiktoken integration
**Outcomes:** Working token counter for OpenAI models with fallbacks  
**Challenges:** Handling different model types and API variations
**Tools Used:** tiktoken library, model configuration system

## Iteration 2 Summary  
**Actions Taken:** Designed ContextWindowConfig class with buffer allocations
**Outcomes:** Engineered buffer system for workspace, summaries, padding
**Challenges:** Calculating optimal buffer ratios
**Tools Used:** ContextWindowConfig, buffer management utilities

## Iteration 3 Summary
**Actions Taken:** Created comprehensive test suite for buffer triggers
**Outcomes:** Verified compaction triggers at exact token limits
**Challenges:** Debugging trigger timing and workspace calculations
**Tools Used:** pytest, mock histories, token calculation utilities
"""
    
    goal = f"""Create an aggregated summary using the AggregationSummaryTool from these iteration summaries:

{iteration_summaries}

Focus on the overall progress, key actions across iterations, major outcomes, and final state."""
    
    prompt = f'agent goal="{goal}", iterations=5'
    
    print("Running AggregationSummarizerAgent...")
    result = await agent.run(prompt)
    
    print(f"Agent run complete. History ID: {result.get('history_id')}")
    
    # Look for tool calls
    agent.look_for_particular_tool_calls()
    
    if agent.last_summary:
        print("✅ AggregationSummarizerAgent working!")
        print(f"Summary preview: {agent.last_summary[:200]}...")
        return agent.last_summary
    else:
        print("❌ AggregationSummarizerAgent failed - no summary generated")
        return None


async def test_reasoning_agent():
    """Test ReasoningAgent with ThinkTool and extraction."""
    print("\n=== Testing ReasoningAgent ===")
    
    test_summary = """# Aggregated Summary

**Total Iterations:** 3

## Overall Progress
Successfully implemented a complete token-based context window management system with engineered buffer allocations and automatic compaction triggers.

### Key Actions Across Iterations
- Created tiktoken-based token counter with multi-model support
- Designed ContextWindowConfig with 45% workspace, 15% compactor, 25% recursive buffers
- Built comprehensive test suite verifying exact trigger points
- Integrated with AutoSummarizingAgent for automatic memory management

### Major Outcomes
- Working token counter supporting OpenAI and fallback models
- Engineered buffer allocation system preventing context overflow
- Verified compaction triggers at exactly 48,960 tokens for gpt-4o-mini
- Complete integration with agent summarization workflow

## Final State
The system automatically manages conversation context by triggering compaction when workspace buffer fills, creating summaries, and maintaining optimal buffer usage across the conversation lifetime."""
    
    print("Running ReasoningAgent...")
    reasoning_result = await reason_about_summary(test_summary)
    
    if reasoning_result and reasoning_result.get('next_steps') and reasoning_result.get('strategic_analysis'):
        print("✅ ReasoningAgent working!")
        print(f"Next Steps: {reasoning_result['next_steps'][:100]}...")
        print(f"Strategic Analysis: {reasoning_result['strategic_analysis'][:100]}...")
        return reasoning_result
    else:
        print("❌ ReasoningAgent failed - no reasoning generated")
        print(f"Result: {reasoning_result}")
        return None


async def test_concept_extractor():
    """Test ConceptExtractorAgent with ConceptSummaryTool.""" 
    print("\n=== Testing ConceptExtractorAgent ===")
    
    test_summary = """Complete token-based context window management system with automatic compaction."""
    
    test_reasoning = {
        'next_steps': 'Enhance recursive summarization layer, add Neo4j semantic graph integration',
        'strategic_analysis': 'Foundation enables unlimited conversation length with intelligent memory management'
    }
    
    print("Running ConceptExtractorAgent...")
    concept_result = await get_summary_concept(test_summary, test_reasoning)
    
    if concept_result and len(concept_result) > 10:
        print("✅ ConceptExtractorAgent working!")
        print(f"Concept preview: {concept_result[:200]}...")
        return concept_result
    else:
        print("❌ ConceptExtractorAgent failed - no concept generated")
        print(f"Result: {concept_result}")
        return None

async def test_auto_summarize_pipeline():
    """Test the complete auto_summarize pipeline."""
    print("\n=== Testing Complete auto_summarize Pipeline ===")

    # Create test history
    test_history = create_test_history()

    print("Running complete auto_summarize pipeline...")
    result = await auto_summarize(test_history)

    # --- Figure out what came back -----------------------------------------
    if isinstance(result, dict):
        # Prefer the richest field that contains the full text
        summary_text = (
            result.get("total_summary")
            or result.get("aggregated_summary")
            or "\n".join(result.get("iteration_summaries", []))
        )
    else:  # old behaviour – result itself is the string we want
        summary_text = result

    # --- Evaluate ----------------------------------------------------------
    if summary_text and isinstance(summary_text, str) and len(summary_text) > 50:
        print("✅ auto_summarize pipeline working!")
        print(f"Summary preview: {summary_text[:300]}...")
        return result          # keep the original structure for later tests
    else:
        print("❌ auto_summarize pipeline failed")
        print(f"Result: {result}")
        return None


async def test_full_workflow():
    """Test the complete get_summary_reasoning_and_concept workflow."""
    print("\n=== Testing Complete Workflow ===")
    
    # Create test history
    test_history = create_test_history()
    
    print("Running complete workflow: summary + reasoning + concept...")
    complete_result = await get_summary_reasoning_and_concept(test_history)
    
    if (complete_result and 
        'summary_and_reasoning' in complete_result and 
        'concept' in complete_result):
        print("✅ Complete workflow working!")
        
        summary = complete_result['summary_and_reasoning']['summary']
        reasoning = complete_result['summary_and_reasoning']['reasoning']  
        concept = complete_result['concept']
        
        print(f"Summary length: {len(summary)} chars")
        print(f"Next steps: {reasoning.get('next_steps', 'None')[:50]}...")
        print(f"Concept length: {len(concept)} chars")
        return complete_result
    else:
        print("❌ Complete workflow failed")
        print(f"Result keys: {list(complete_result.keys()) if complete_result else 'None'}")
        return None


async def main():
    """Run all compactor agent tests."""
    print("🧪 HEAVEN Compactor Agents Test Suite")
    print("=====================================")
    
    results = {}
    
    # Test each agent individually
    results['iteration_summarizer'] = await test_iteration_summarizer()
    results['aggregation_summarizer'] = await test_aggregation_summarizer()  
    results['reasoning_agent'] = await test_reasoning_agent()
    results['concept_extractor'] = await test_concept_extractor()
    
    # Test pipelines
    results['auto_summarize'] = await test_auto_summarize_pipeline()
    results['full_workflow'] = await test_full_workflow()
    
    # Final summary
    print(f"\n=== Test Results Summary ===")
    for test_name, result in results.items():
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{test_name}: {status}")
    
    # Count successes
    successes = sum(1 for result in results.values() if result)
    total = len(results)
    
    print(f"\nOverall: {successes}/{total} tests passed")
    
    if successes == total:
        print("🎉 All compactor agents working correctly!")
    else:
        print("⚠️  Some compactor agents need fixes")
    
    return results


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