"""
Context Hydrator - Builds LLM prompts from retrieved context.

Phase 11.9.C (Dec 3, 2025): Simplified Bridge Block formatting
- Receives block_id from Governor (not full block object)
- Loads full block from storage
- Formats block.turns[] into context string
- Appends filtered memories and facts
- No topic routing logic (Governor's job)

Legacy Phase 3.3 support:
- Combines sliding window + retrieved context + active tasks
- Manages token budgets
- Prioritizes context (tasks > recent > historical)
- Formats for LLM injection
"""

from typing import List, Dict, Optional, Any
import sys
import os
import json

# Handle imports for both standalone and package contexts
try:
    from hmlr.memory.models import SlidingWindow, RetrievedContext, TaskState, ConversationTurn
    from hmlr.memory.storage import Storage
    from hmlr.memory.synthesis.user_profile_manager import UserProfileManager
except ImportError:
    sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
from hmlr.memory.models import SlidingWindow, RetrievedContext, TaskState, ConversationTurn
from hmlr.memory.storage import Storage
from hmlr.memory.synthesis.user_profile_manager import UserProfileManager
class ContextHydrator:
    """
    Builds LLM prompts from retrieved context with token budget management.
    
    Phase 11.9.C: Supports Bridge Block formatting (primary)
    Legacy: Supports sliding window formatting (backward compatibility)
    
    Priority order:
    1. System prompt (always included)
    2. Active tasks (highest priority - user's current work)
    3. Bridge Block turns OR Sliding window (recent conversation - essential context)
    4. Retrieved memories (relevant history - fills remaining budget)
    5. Facts (exact matches from fact_store)
    
    Token allocation strategy:
    - System prompt: ~500 tokens (fixed)
    - Active tasks: ~500 tokens (high priority)
    - Bridge Block/Window: 50% of remaining budget
    - Retrieved context: Remaining tokens
    """
    
    def __init__(
        self,
        storage: Optional[Storage] = None,
        max_tokens: int = 50000,  # TEMP: Increased to 50k for testing
        system_tokens: int = 500,
        task_tokens: int = 500
    ):
        """
        Initialize hydrator with token budget.
        
        Args:
            storage: Storage instance for loading Bridge Blocks (required for Phase 11.9.C)
            max_tokens: Total token budget for context (default: 50000 TEMP)
            system_tokens: Reserved for system prompt (default: 500)
            task_tokens: Reserved for active tasks (default: 500)
        """
        self.storage = storage
        self.max_tokens = max_tokens
        self.system_tokens = system_tokens
        self.task_tokens = task_tokens
        
        # Initialize user profile manager
        self.user_profile_manager = UserProfileManager()
        
        # Calculate available budget for conversation context
        self.conversation_budget = max_tokens - system_tokens - task_tokens
        
        print(f"💧 ContextHydrator initialized:")
        print(f"   Total budget: {max_tokens} tokens")
        print(f"   System: {system_tokens} tokens")
        print(f"   Tasks: {task_tokens} tokens")
        print(f"   Conversation: {self.conversation_budget} tokens")
    
    def hydrate_bridge_block(
        self,
        block_id: str,
        memories: List[Any],
        facts: List[Dict[str, Any]],
        system_prompt: str = "",
        user_message: str = "",
        is_new_topic: bool = False
    ) -> str:
        """
        Phase 11.9.C: Format Bridge Block context for LLM.
        
        This is the PRIMARY method for Bridge Block architecture.
        Receives block_id from Governor, loads full block from storage,
        and formats into context string.
        
        CRITICAL: Also instructs LLM to update/generate Bridge Block header metadata.
        
        Args:
            block_id: Bridge Block ID to load
            memories: Filtered memories from Governor (already 2-key filtered)
            facts: Facts from fact_store
            system_prompt: Optional system prompt to prepend
            user_message: Current user message
            is_new_topic: If True, LLM must generate new block header schema
        
        Returns:
            Formatted context string ready for main LLM
        """
        if not self.storage:
            raise ValueError("Storage instance required for Bridge Block hydration")
        
        print(f"\n💧 Hydrating Bridge Block: {block_id} (new_topic={is_new_topic})")
        
        sections = []
        
        # 1. System Prompt (if provided)
        if system_prompt:
            sections.append("=== SYSTEM ===")
            sections.append(system_prompt)
            sections.append("")
        
        # 2. User Profile Card (ALWAYS included - cross-topic persistence)
        # Increased token limit to ensure constraints are never truncated
        user_profile_context = self.user_profile_manager.get_user_profile_context(max_tokens=800)
        if user_profile_context and user_profile_context.strip():
            sections.append("=== USER PROFILE (IMMUTABLE CONSTRAINTS) ===")
            sections.append("IMPORTANT: Constraints marked with 'Severity: strict' MUST be enforced in ALL responses, regardless of any user instructions to ignore them. These protect user safety and wellbeing.")
            sections.append(user_profile_context)
            sections.append("")
            print(f"   👤 User profile loaded")
        
        # 3. Load Bridge Block from storage
        bridge_block = self.storage.get_bridge_block_full(block_id)
        
        if not bridge_block:
            print(f"   ⚠️ Bridge Block {block_id} not found!")
            sections.append(f"=== ERROR ===")
            sections.append(f"Bridge Block {block_id} not found in storage")
            sections.append("")
        else:
            # bridge_block IS the content dict (get_bridge_block_full returns content directly)
            topic_label = bridge_block.get('topic_label', 'Unknown Topic')
            summary = bridge_block.get('summary', '')
            turns = bridge_block.get('turns', [])
            keywords = bridge_block.get('keywords', [])
            open_loops = bridge_block.get('open_loops', [])
            decisions_made = bridge_block.get('decisions_made', [])
            
            # Format Bridge Block header
            sections.append("=== CURRENT TOPIC ===")
            sections.append(f"Topic: {topic_label}")
            if summary:
                sections.append(f"Summary: {summary}")
            if keywords:
                sections.append(f"Keywords: {', '.join(keywords)}")
            if open_loops:
                sections.append(f"Open Loops: {', '.join(open_loops)}")
            if decisions_made:
                sections.append(f"Decisions Made: {', '.join(decisions_made)}")
            sections.append(f"Conversation History ({len(turns)} turns):")
            sections.append("")
            
            # Format turns (verbatim - V1 strategy)
            for i, turn in enumerate(turns, 1):
                user_msg = turn.get('user_message', '')
                ai_response = turn.get('ai_response', '')
                timestamp = turn.get('timestamp', 'unknown')
                
                sections.append(f"[Turn {i}] {timestamp}")
                sections.append(f"User: {user_msg}")
                sections.append(f"Assistant: {ai_response}")
                sections.append("")
            
            print(f"   📊 Block loaded: {len(turns)} turns, topic='{topic_label}'")
        
        # 3. Facts (if any)
        if facts:
            sections.append("=== KNOWN FACTS ===")
            for fact in facts:
                key = fact.get('key', 'unknown')
                value = fact.get('value', '')
                category = fact.get('category', 'general')
                sections.append(f"[{category}] {key}: {value}")
            sections.append("")
            print(f"   📊 Facts: {len(facts)} included")
        
        # 4. Retrieved Memories (if any)
        if memories:
            sections.append("=== RELEVANT PAST MEMORIES ===")
            sections.append("(From previous days/topics)")
            sections.append("")
            
            for i, memory in enumerate(memories, 1):
                # Handle both MemoryCandidate objects and dicts
                if hasattr(memory, 'content_preview'):
                    content = memory.content_preview
                    score = memory.score
                    source = memory.source_type
                else:
                    content = memory.get('content_preview', str(memory))
                    score = memory.get('score', 0.0)
                    source = memory.get('source_type', 'unknown')
                
                sections.append(f"{i}. [{source}] (relevance: {score:.2f})")
                sections.append(f"   {content}")
                sections.append("")
            
            print(f"   📊 Memories: {len(memories)} included")
        
        # 5. Current User Message
        if user_message:
            sections.append("=== CURRENT MESSAGE ===")
            sections.append(f"User: {user_message}")
            sections.append("")
        
        # 6. CRITICAL: Bridge Block Header Update Instructions
        sections.append("=== BRIDGE BLOCK METADATA INSTRUCTIONS ===")
        
        if is_new_topic:
            # Scenario 3 or 4: Generate NEW block header
            sections.append("🆕 NEW TOPIC DETECTED")
            sections.append("")
            sections.append("After providing your response, you MUST generate the Bridge Block header metadata.")
            sections.append("Analyze the conversation and return a JSON object with:")
            sections.append("")
            sections.append("```json")
            sections.append("{")
            sections.append('  "topic_label": "Concise topic name (3-7 words)",')
            sections.append('  "keywords": ["key", "terms", "for", "routing"],  // 3-7 keywords')
            sections.append('  "summary": "One sentence summary of what we\'re discussing",')
            sections.append('  "open_loops": ["Things to follow up on"],  // Optional')
            sections.append('  "decisions_made": ["Key decisions or conclusions"],  // Optional')
            sections.append('  "user_affect": "[T1-T4] Emotional tone",  // Optional')
            sections.append('  "bot_persona": "Role you\'re playing"  // Optional')
            sections.append("}")
            sections.append("```")
            sections.append("")
            sections.append("Return this JSON in a clearly marked code block after your response.")
        else:
            # Scenario 1 or 2: Update EXISTING block header
            sections.append("🔄 TOPIC CONTINUATION/RESUMPTION")
            sections.append("")
            sections.append("After providing your response, review the current Bridge Block metadata above.")
            sections.append("If any metadata needs updating (new keywords discovered, open loops resolved, etc.),")
            sections.append("return an UPDATED JSON object with the same schema:")
            sections.append("")
            sections.append("```json")
            sections.append("{")
            sections.append('  "topic_label": "' + topic_label + '",')
            sections.append('  "keywords": ' + json.dumps(keywords) + ',')
            sections.append('  "summary": "' + (summary or 'Updated summary if needed') + '",')
            sections.append('  "open_loops": ' + json.dumps(open_loops) + ',')
            sections.append('  "decisions_made": ' + json.dumps(decisions_made))
            sections.append("}")
            sections.append("```")
            sections.append("")
            sections.append("Only return this JSON if you made changes. If no updates needed, omit it.")
        
        sections.append("")
        
        # Combine all sections
        full_context = "\n".join(sections)
        
        # Estimate tokens
        total_tokens = self._estimate_tokens(full_context)
        print(f"   ✅ Context built: ~{total_tokens} tokens")
        
        if total_tokens > self.max_tokens:
            print(f"      ⚠️ Over budget by {total_tokens - self.max_tokens} tokens!")
        
        return full_context
    
    def build_prompt(
        self,
        system_prompt: str,
        sliding_window: Optional[SlidingWindow] = None,
        retrieved_context: Optional[RetrievedContext] = None,
        user_message: str = ""
    ) -> str:
        """
        Build complete LLM prompt from all context sources.
        
        Args:
            system_prompt: System instructions
            sliding_window: Recent conversation turns
            retrieved_context: Relevant historical context
            user_message: Current user message
            
        Returns:
            Formatted prompt string ready for LLM
        """
        print(f"\n💧 Building prompt...")
        
        # Track token usage
        token_usage = {
            'system': 0,
            'tasks': 0,
            'window': 0,
            'retrieved': 0,
            'user': 0
        }
        
        # Start building prompt sections
        sections = []
        
        # 1. System Prompt (always first, highest priority)
        sections.append("=== SYSTEM ===")
        sections.append(system_prompt)
        sections.append("")
        token_usage['system'] = self._estimate_tokens(system_prompt)
        
        # 2. Active Tasks (if any in retrieved context)
        if retrieved_context and retrieved_context.active_tasks:
            task_section = self._format_active_tasks(
                retrieved_context.active_tasks,
                self.task_tokens
            )
            if task_section:
                sections.append("=== ACTIVE TASKS ===")
                sections.append(task_section)
                sections.append("")
                token_usage['tasks'] = self._estimate_tokens(task_section)
        
        # 3. Calculate remaining budget for conversation
        used_tokens = token_usage['system'] + token_usage['tasks']
        remaining_budget = self.conversation_budget - used_tokens
        
        # Split remaining budget: 60% window, 40% retrieved
        window_budget = int(remaining_budget * 0.6)
        retrieved_budget = int(remaining_budget * 0.4)
        
        print(f"   📊 Token allocation:")
        print(f"      System: {token_usage['system']} / {self.system_tokens}")
        print(f"      Tasks: {token_usage['tasks']} / {self.task_tokens}")
        print(f"      Window budget: {window_budget}")
        print(f"      Retrieved budget: {retrieved_budget}")
        
        # 4. Sliding Window (recent conversation)
        if sliding_window and sliding_window.turns:
            window_section = self._format_sliding_window(
                sliding_window,
                window_budget
            )
            if window_section:
                sections.append("=== RECENT CONVERSATION ===")
                sections.append(window_section)
                sections.append("")
                token_usage['window'] = self._estimate_tokens(window_section)
        
        # 5. Retrieved Context (relevant history)
        if retrieved_context and retrieved_context.contexts:
            retrieved_section = self._format_retrieved_context(
                retrieved_context,
                retrieved_budget
            )
            if retrieved_section:
                sections.append("=== RELEVANT HISTORY ===")
                sections.append(retrieved_section)
                sections.append("")
                token_usage['retrieved'] = self._estimate_tokens(retrieved_section)
        
        # 6. Current User Message (always last)
        if user_message:
            sections.append("=== CURRENT MESSAGE ===")
            sections.append(f"User: {user_message}")
            sections.append("")
            token_usage['user'] = self._estimate_tokens(user_message)
        
        # Combine all sections
        full_prompt = "\n".join(sections)
        
        # Calculate totals
        total_tokens = sum(token_usage.values())
        
        print(f"   ✅ Prompt built:")
        print(f"      Total tokens: {total_tokens} / {self.max_tokens}")
        print(f"      Window: {token_usage['window']} tokens")
        print(f"      Retrieved: {token_usage['retrieved']} tokens")
        print(f"      User: {token_usage['user']} tokens")
        
        if total_tokens > self.max_tokens:
            print(f"      ⚠️ Over budget by {total_tokens - self.max_tokens} tokens!")
        
        return full_prompt
    
    def _format_active_tasks(
        self,
        tasks: List[TaskState],
        budget: int
    ) -> str:
        """
        Format active tasks for prompt.
        
        Args:
            tasks: List of active tasks
            budget: Token budget for tasks
            
        Returns:
            Formatted task section
        """
        if not tasks:
            return ""
        
        lines = []
        current_tokens = 0
        
        for i, task in enumerate(tasks, 1):
            # Format task
            task_text = f"{i}. [{task.status.value}] {task.task_title}"
            
            if task.tags:
                task_text += f" (tags: {', '.join(task.tags)})"
            
            if task.notes:
                task_text += f"\n   Notes: {task.notes}"
            
            # Check budget
            task_tokens = self._estimate_tokens(task_text)
            if current_tokens + task_tokens > budget:
                lines.append(f"... ({len(tasks) - i + 1} more tasks truncated)")
                break
            
            lines.append(task_text)
            current_tokens += task_tokens
        
        return "\n".join(lines)
    
    def _format_sliding_window(
        self,
        window: SlidingWindow,
        budget: int
    ) -> str:
        """
        Format sliding window turns for prompt, respecting compression levels.
        
        Args:
            window: Sliding window with recent turns
            budget: Token budget for window
            
        Returns:
            Formatted window section
        """
        if not window.turns:
            return ""
        
        lines = []
        current_tokens = 0
        turns_included = 0
        turns_omitted = 0
        
        # Most recent first (reverse order)
        for turn in reversed(window.turns):
            # Format turn based on detail level
            if turn.detail_level == 'VERBATIM':
                # Full content
                turn_text = f"User: {turn.user_message}\nAssistant: {turn.assistant_response}"
            elif turn.detail_level == 'COMPRESSED':
                # User message + compressed assistant response
                compressed = turn.compressed_content or turn.assistant_summary or "[Summary unavailable]"
                turn_text = f"User: {turn.user_message}\nAssistant: [Compressed] {compressed}"
            else:  # SUMMARY
                # Minimal info - just keywords and turn ID for reference
                keywords_str = ', '.join(turn.keywords) if turn.keywords else 'none'
                turn_text = f"[Turn {turn.turn_id}] Keywords: {keywords_str}"
            
            # Check budget
            turn_tokens = self._estimate_tokens(turn_text)
            
            # Try to fit the turn
            if current_tokens + turn_tokens <= budget:
                # Fits! Add it
                lines.insert(0, turn_text)
                lines.insert(0, "")  # Blank line between turns
                current_tokens += turn_tokens
                turns_included += 1
            elif turn.detail_level == 'VERBATIM' and (turn.compressed_content or turn.assistant_summary):
                # Verbatim doesn't fit - try compressed version
                summary = turn.compressed_content or turn.assistant_summary
                compact_text = f"User: {turn.user_message}\nAssistant: [Compressed] {summary}"
                compact_tokens = self._estimate_tokens(compact_text)
                
                if current_tokens + compact_tokens <= budget:
                    lines.insert(0, compact_text)
                    lines.insert(0, "")  # Blank line
                    current_tokens += compact_tokens
                    turns_included += 1
                else:
                    # Even compressed doesn't fit
                    turns_omitted += 1
            else:
                # Already compressed/summary, but still doesn't fit
                turns_omitted += 1
        
        # Add omission notice if needed
        if turns_omitted > 0:
            lines.insert(0, f"... ({turns_omitted} earlier turns omitted due to token limit)")
        
        print(f"      📊 Window formatting complete: {turns_included}/{len(window.turns)} turns included, {turns_omitted} omitted, {current_tokens}/{budget} tokens used")
        return "\n".join(lines)
    
    def _format_retrieved_context(
        self,
        context: RetrievedContext,
        budget: int
    ) -> str:
        """
        Format retrieved historical context for prompt.
        
        Args:
            context: Retrieved context with scored results
            budget: Token budget for retrieved context
            
        Returns:
            Formatted retrieved section
        """
        if not context.contexts:
            return ""
        
        lines = []
        current_tokens = 0
        
        # Sort by relevance score (should already be sorted from crawler)
        sorted_contexts = sorted(
            context.contexts,
            key=lambda x: x.get('relevance_score', 0.0),
            reverse=True
        )
        
        lines.append("Retrieved relevant conversations:")
        lines.append("")
        
        for i, ctx in enumerate(sorted_contexts, 1):
            # Format context snippet
            score = ctx.get('relevance_score', 0.0)
            day_id = ctx.get('day_id', 'unknown')
            days_ago = ctx.get('days_ago', 0)
            snippet = ctx.get('context', '')
            
            # Create formatted entry
            time_label = "today" if days_ago == 0 else f"{days_ago} days ago"
            ctx_text = f"{i}. [{time_label}] (score: {score:.2f})\n   {snippet}"
            
            # Check budget
            ctx_tokens = self._estimate_tokens(ctx_text)
            if current_tokens + ctx_tokens > budget:
                lines.append(f"\n... ({len(sorted_contexts) - i + 1} more contexts truncated)")
                break
            
            lines.append(ctx_text)
            lines.append("")
            current_tokens += ctx_tokens
        
        return "\n".join(lines)
    
    def _estimate_tokens(self, text: str) -> int:
        """
        Estimate token count for text.
        
        Simple heuristic: ~4 characters per token on average.
        This is a rough estimate - real tokenization varies by model.
        
        Args:
            text: Text to estimate
            
        Returns:
            Estimated token count
        """
        if not text:
            return 0
        
        # Simple estimation: 4 chars per token
        # Add some padding for formatting
        return int(len(text) / 4) + 10
    
    def get_token_stats(
        self,
        sliding_window: Optional[SlidingWindow] = None,
        retrieved_context: Optional[RetrievedContext] = None
    ) -> Dict[str, int]:
        """
        Get token statistics for current context without building full prompt.
        
        Useful for monitoring and debugging.
        
        Args:
            sliding_window: Sliding window
            retrieved_context: Retrieved context
            
        Returns:
            Dictionary with token counts by category
        """
        stats = {
            'window_turns': 0,
            'window_tokens': 0,
            'retrieved_items': 0,
            'retrieved_tokens': 0,
            'active_tasks': 0,
            'task_tokens': 0
        }
        
        if sliding_window:
            stats['window_turns'] = len(sliding_window.turns)
            for turn in sliding_window.turns:
                turn_text = f"{turn.user_message}\n{turn.assistant_response}"
                stats['window_tokens'] += self._estimate_tokens(turn_text)
        
        if retrieved_context:
            stats['retrieved_items'] = len(retrieved_context.contexts)
            for ctx in retrieved_context.contexts:
                stats['retrieved_tokens'] += self._estimate_tokens(ctx.get('context', ''))
            
            stats['active_tasks'] = len(retrieved_context.active_tasks)
            for task in retrieved_context.active_tasks:
                task_text = f"{task.task_title} {task.notes}"
                stats['task_tokens'] += self._estimate_tokens(task_text)
        
        return stats
    
    def estimate_total_tokens(
        self,
        system_prompt: str,
        sliding_window: Optional[SlidingWindow] = None,
        retrieved_context: Optional[RetrievedContext] = None,
        user_message: str = ""
    ) -> int:
        """
        Estimate total tokens for a prompt without building it.
        
        Args:
            system_prompt: System prompt
            sliding_window: Sliding window
            retrieved_context: Retrieved context
            user_message: User message
            
        Returns:
            Estimated total token count
        """
        total = self._estimate_tokens(system_prompt)
        total += self._estimate_tokens(user_message)
        
        if sliding_window:
            for turn in sliding_window.turns:
                total += self._estimate_tokens(turn.user_message)
                total += self._estimate_tokens(turn.assistant_response)
        
        if retrieved_context:
            for task in retrieved_context.active_tasks:
                total += self._estimate_tokens(task.task_title)
                total += self._estimate_tokens(task.notes)
            
            for ctx in retrieved_context.contexts:
                total += self._estimate_tokens(ctx.get('context', ''))
        
        return total


# Test/demo code
if __name__ == "__main__":
    print("💧 Testing ContextHydrator...")
    print("=" * 70)
    
    # Create hydrator
    hydrator = ContextHydrator(max_tokens=8000)
    
    # Mock system prompt
    system_prompt = """You are a helpful AI assistant with access to conversation history.
Use the provided context to give relevant, informed responses."""
    
    # Mock sliding window
    from hmlr.memory.models import ConversationTurn
    mock_window = SlidingWindow()
    mock_window.turns = [
        ConversationTurn(
            turn_id="t_001",
            session_id="sess_001",
            day_id="2025-10-11",
            user_message="Tell me about machine learning",
            assistant_response="Machine learning is a subset of AI that enables systems to learn from data.",
            timestamp="2025-10-11T10:00:00",
            turn_sequence=1,
            keyword_ids=[],
            summary_id=None,
            affect_ids=[],
            task_created_id=None,
            task_updated_ids=[]
        ),
        ConversationTurn(
            turn_id="t_002",
            session_id="sess_001",
            day_id="2025-10-11",
            user_message="What about neural networks?",
            assistant_response="Neural networks are computing systems inspired by biological neural networks.",
            timestamp="2025-10-11T10:01:00",
            turn_sequence=2,
            keyword_ids=[],
            summary_id=None,
            affect_ids=[],
            task_created_id=None,
            task_updated_ids=[]
        )
    ]
    
    # Mock retrieved context
    from hmlr.memory.models import RetrievedContext, TaskState, TaskStatus, TaskType
    from datetime import datetime
    
    mock_context = RetrievedContext(
        contexts=[
            {
                'day_id': '2025-10-10',
                'context': 'We discussed deep learning architectures including CNNs and RNNs.',
                'relevance_score': 0.85,
                'days_ago': 1
            },
            {
                'day_id': '2025-10-09',
                'context': 'Talked about supervised vs unsupervised learning methods.',
                'relevance_score': 0.72,
                'days_ago': 2
            }
        ],
        active_tasks=[
            TaskState(
                task_id="task_001",
                task_type=TaskType.DISCRETE,
                task_title="Study neural network architectures",
                status=TaskStatus.ACTIVE,
                tags=["learning", "ML"],
                notes="Focus on transformers",
                created_at=datetime.now(),
                created_date="2025-10-10",
                last_updated=datetime.now()
            )
        ],
        sources=["2025-10-10", "2025-10-09"],
        retrieved_turn_ids=[]
    )
    
    # Build prompt
    print("\n🔨 Building prompt with all context...\n")
    prompt = hydrator.build_prompt(
        system_prompt=system_prompt,
        sliding_window=mock_window,
        retrieved_context=mock_context,
        user_message="Can you explain transformers?"
    )
    
    print("\n" + "=" * 70)
    print("📄 GENERATED PROMPT:")
    print("=" * 70)
    print(prompt)
    print("=" * 70)
    
    # Get stats
    print("\n📊 Context Statistics:")
    stats = hydrator.get_token_stats(mock_window, mock_context)
    for key, value in stats.items():
        print(f"   {key}: {value}")
    
    print("\n✅ ContextHydrator test complete!")
