#!/usr/bin/env python3
"""
Test token-based auto-summarization with exact token control.

Creates a synthetic history with precise token counts to trigger summarization.
"""

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

from heaven_base.auto_summarizing_agent import AutoSummarizingAgent
from heaven_base.baseheavenagent import HeavenAgentConfig
from heaven_base.memory.history import History
from heaven_base.unified_chat import UnifiedChat
from heaven_base.utils.token_counter import count_tokens


def create_exact_token_message(target_tokens: int, model: str = "gpt-4o-mini") -> str:
    """Create a message with exactly the target number of tokens."""
    # Start with a base message
    base_msg = "This is a test message for token counting. "
    
    # Repeat until we get close to target
    current_msg = base_msg
    while count_tokens(current_msg, model) < target_tokens - 50:
        current_msg += base_msg
    
    # Fine-tune to exact count
    words_to_add = ["token", "count", "test", "message", "example", "data", "content", "text"]
    word_idx = 0
    
    while count_tokens(current_msg, model) < target_tokens:
        current_msg += f" {words_to_add[word_idx % len(words_to_add)]}"
        word_idx += 1
    
    # Trim if we went over
    while count_tokens(current_msg, model) > target_tokens:
        # Remove last word
        current_msg = " ".join(current_msg.split()[:-1])
    
    return current_msg


def create_1000_token_message_pair(model: str = "gpt-4o-mini"):
    """Create a user-AI message pair that totals exactly 1000 tokens."""
    # Create 500-token user message
    user_msg = create_exact_token_message(500, model)
    
    # Create 500-token AI message
    ai_msg = create_exact_token_message(500, model)
    
    # Verify total
    total_tokens = count_tokens(user_msg, model) + count_tokens(ai_msg, model)
    print(f"Created message pair: {count_tokens(user_msg, model)} + {count_tokens(ai_msg, model)} = {total_tokens} tokens")
    
    return user_msg, ai_msg


async def test_token_based_summarization():
    """Test that token-based summarization triggers at the correct threshold."""
    print("=== Testing Token-Based Auto-Summarization ===\n")
    
    model = "gpt-4o-mini"
    
    # Create 1000-token message pair
    user_msg, ai_msg = create_1000_token_message_pair(model)
    
    # Create agent config with a small context window for testing
    config = HeavenAgentConfig(
        name="test_summarizing_agent",
        system_prompt="You are a test agent for summarization testing.",
        model=model
    )
    
    # Mock unified chat (we won't actually call API)
    class MockUnifiedChat:
        def create(self, **kwargs):
            return self
        
        def invoke_uni_api(self, **kwargs):
            return {
                "choices": [{"message": {"content": "Mock response"}}],
                "usage": {"total_tokens": 1000, "prompt_tokens": 800, "completion_tokens": 200}
            }
    
    unified_chat = MockUnifiedChat()
    
    # Create synthetic history
    history = History(messages=[
        SystemMessage(content="System prompt for testing.")
    ])
    
    # Callbacks to track summarization
    summarization_events = []
    
    def on_start(token_count):
        summarization_events.append(f"START: {token_count} tokens")
        print(f"🔥 Summarization started at {token_count} tokens")
    
    def on_complete(summary_result):
        summarization_events.append(f"COMPLETE: {len(str(summary_result))} chars")
        print(f"✅ Summarization completed: {len(str(summary_result))} char summary")
    
    # Create AutoSummarizingAgent with low token threshold
    agent = AutoSummarizingAgent(
        config=config,
        unified_chat=unified_chat,
        history=history,
        token_threshold=0.7,  # Trigger at 70% of context window
        on_summarization_start=on_start,
        on_summarization_complete=on_complete
    )
    
    # Get the context window limit for this model
    context_limit = agent.context_window_config.usable_window
    trigger_point = int(context_limit * 0.7)
    print(f"Context window: {context_limit} tokens")
    print(f"Trigger point: {trigger_point} tokens (70%)\n")
    
    # Add message pairs until we hit the threshold
    message_pairs_added = 0
    total_tokens = count_tokens(history.messages[0].content, model)  # System message
    
    print("Adding 1000-token message pairs...")
    
    while total_tokens < trigger_point:
        # Add user message
        history.messages.append(HumanMessage(content=user_msg))
        total_tokens += count_tokens(user_msg, model)
        
        # Add AI response
        history.messages.append(AIMessage(content=ai_msg))
        total_tokens += count_tokens(ai_msg, model)
        
        message_pairs_added += 1
        
        # Update agent's context window with current tokens
        agent.context_window_config.current_tokens = total_tokens
        
        print(f"  Pair {message_pairs_added}: {total_tokens} tokens")
        
        # Check if we should summarize
        if total_tokens >= trigger_point:
            print(f"\n🎯 Reached trigger point! ({total_tokens} >= {trigger_point})")
            break
    
    print(f"\nAdded {message_pairs_added} message pairs ({total_tokens} total tokens)")
    print(f"Messages in history: {len(history.messages)}")
    
    # Now test the summarization check
    print("\n--- Testing Summarization Trigger ---")
    
    needs_summarization = agent.context_window_config.should_summarize(0.7)
    print(f"should_summarize(0.7): {needs_summarization}")
    
    if needs_summarization:
        print("✅ Token-based trigger working correctly!")
        
        # Test the actual summarization (this would normally call RecursiveSummarizer)
        try:
            print("\n--- Testing Summarization Process ---")
            summarized = await agent._check_and_summarize_if_needed()
            print(f"Summarization performed: {summarized}")
            
            if summarization_events:
                print("\nSummarization events:")
                for event in summarization_events:
                    print(f"  {event}")
            else:
                print("⚠️  No summarization events fired (callbacks not working)")
                
        except Exception as e:
            print(f"⚠️  Summarization failed (expected - no RecursiveSummarizer): {e}")
            print("   This is expected in isolated test without full summarization stack")
    else:
        print("❌ Token-based trigger not working!")
    
    # Test context window status
    print(f"\n--- Context Window Status ---")
    status = agent.context_window_config.get_status()
    for key, value in status.items():
        print(f"  {key}: {value}")
    
    print(f"\n=== Test Complete ===")
    return needs_summarization, message_pairs_added, total_tokens


if __name__ == "__main__":
    result = asyncio.run(test_token_based_summarization())
    print(f"\nTest Result: {result}")