#!/usr/bin/env python3
"""
VUA Context Manager - Vigilant Undivided Attention Pattern
============================================================
Monitors context window utilization and compresses state when needed.

When context > 80%, automatically:
1. Summarizes current reasoning into State Summary
2. Persists to disk for retrieval
3. Returns compressed state for fresh context

Based on VAST architecture principles for agentic memory.

Usage:
    from vua_context import VUAContextManager
    
    vua = VUAContextManager()
    
    # Check if compression needed
    if vua.should_compress(current_context_size):
        summary = vua.compress_and_save(reasoning_state)
        # Use summary instead of full state
    
    # Retrieve previous compressions
    summaries = vua.get_recent_summaries(limit=5)
"""

import json
import hashlib
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Optional
from dataclasses import dataclass, asdict


@dataclass
class StateSummary:
    """Compressed state representation."""
    id: str
    timestamp: str
    context_size_before: int
    summary_size: int
    compression_ratio: float
    key_findings: List[str]
    active_goals: List[str]
    pending_actions: List[str]
    axioms_snapshot: List[str]
    full_summary: str


class VUAContextManager:
    """
    Vigilant Undivided Attention - Context compression manager.
    
    Prevents context overflow by detecting when to compress
    and storing compressed state for continuity.
    """
    
    # Context thresholds (tokens roughly = chars / 4)
    COMPRESSION_THRESHOLD = 0.80  # 80% triggers compression
    TARGET_RATIO = 0.10  # Target 10% of original size
    MAX_SUMMARY_CHARS = 4000  # Max summary size
    
    def __init__(self, data_dir: str = None):
        if data_dir is None:
            data_dir = Path(__file__).parent.parent / "data"
        self.data_dir = Path(data_dir)
        self.summaries_dir = self.data_dir / "context_summaries"
        self.summaries_dir.mkdir(parents=True, exist_ok=True)
        self.state_path = self.data_dir / "vua_state.json"
        self._load_state()
    
    def _load_state(self):
        """Load VUA tracking state."""
        if self.state_path.exists():
            try:
                with open(self.state_path) as f:
                    self.state = json.load(f)
            except:
                self.state = self._default_state()
        else:
            self.state = self._default_state()
    
    def _default_state(self) -> Dict:
        return {
            "total_compressions": 0,
            "total_tokens_saved": 0,
            "last_compression": None,
            "summary_ids": []
        }
    
    def _save_state(self):
        """Persist VUA state."""
        with open(self.state_path, 'w') as f:
            json.dump(self.state, f, indent=2)
    
    def should_compress(
        self,
        current_tokens: int,
        max_tokens: int = 128000  # Claude's context
    ) -> bool:
        """
        Check if context compression is needed.
        
        Args:
            current_tokens: Estimated current token count
            max_tokens: Maximum context window
            
        Returns:
            True if compression recommended
        """
        utilization = current_tokens / max_tokens
        return utilization >= self.COMPRESSION_THRESHOLD
    
    def estimate_tokens(self, text: str) -> int:
        """Rough token estimation (chars / 4)."""
        return len(text) // 4
    
    def compress_and_save(
        self,
        reasoning_state: Dict[str, Any],
        key_findings: List[str] = None,
        active_goals: List[str] = None,
        pending_actions: List[str] = None,
        axioms: List[str] = None
    ) -> StateSummary:
        """
        Compress current state and save to disk.
        
        Args:
            reasoning_state: Full current state dict
            key_findings: Important discoveries this session
            active_goals: What we're trying to achieve
            pending_actions: What's next
            axioms: Current axiom statements
            
        Returns:
            StateSummary with compressed representation
        """
        timestamp = datetime.now().isoformat()
        
        # Serialize full state
        full_json = json.dumps(reasoning_state, indent=2, default=str)
        original_size = len(full_json)
        
        # Generate compressed summary
        summary_text = self._generate_summary(
            reasoning_state,
            key_findings or [],
            active_goals or [],
            pending_actions or []
        )
        
        # Create summary object
        summary_id = hashlib.sha256(
            f"{timestamp}:{summary_text[:100]}".encode()
        ).hexdigest()[:12]
        
        summary = StateSummary(
            id=summary_id,
            timestamp=timestamp,
            context_size_before=original_size,
            summary_size=len(summary_text),
            compression_ratio=len(summary_text) / max(original_size, 1),
            key_findings=key_findings or [],
            active_goals=active_goals or [],
            pending_actions=pending_actions or [],
            axioms_snapshot=axioms or [],
            full_summary=summary_text
        )
        
        # Save to disk
        summary_path = self.summaries_dir / f"summary_{summary_id}.json"
        with open(summary_path, 'w') as f:
            json.dump(asdict(summary), f, indent=2)
        
        # Also save full state for potential recovery
        full_path = self.summaries_dir / f"full_{summary_id}.json"
        with open(full_path, 'w') as f:
            json.dump({
                "timestamp": timestamp,
                "summary_id": summary_id,
                "full_state": reasoning_state
            }, f, indent=2, default=str)
        
        # Update tracking
        self.state["total_compressions"] += 1
        self.state["total_tokens_saved"] += (original_size - len(summary_text)) // 4
        self.state["last_compression"] = timestamp
        self.state["summary_ids"].append(summary_id)
        self._save_state()
        
        return summary
    
    def _generate_summary(
        self,
        state: Dict,
        findings: List[str],
        goals: List[str],
        actions: List[str]
    ) -> str:
        """
        Generate compressed summary text.
        
        This is a deterministic compression. For LLM-powered compression,
        override this method with API call.
        """
        lines = [
            "=== STATE SUMMARY ===",
            f"Captured: {datetime.now().isoformat()}",
            ""
        ]
        
        # Key findings
        if findings:
            lines.append("## Key Findings")
            for f in findings[:10]:  # Limit
                lines.append(f"- {f[:200]}")
            lines.append("")
        
        # Goals
        if goals:
            lines.append("## Active Goals")
            for g in goals[:5]:
                lines.append(f"- {g[:150]}")
            lines.append("")
        
        # Actions
        if actions:
            lines.append("## Pending Actions")
            for a in actions[:10]:
                lines.append(f"- {a[:150]}")
            lines.append("")
        
        # State metrics
        lines.append("## State Metrics")
        if isinstance(state, dict):
            for key, value in list(state.items())[:20]:
                if isinstance(value, (str, int, float, bool)):
                    lines.append(f"- {key}: {str(value)[:100]}")
                elif isinstance(value, list):
                    lines.append(f"- {key}: [{len(value)} items]")
                elif isinstance(value, dict):
                    lines.append(f"- {key}: {{...{len(value)} keys}}")
        
        summary = "\n".join(lines)
        
        # Enforce max size
        if len(summary) > self.MAX_SUMMARY_CHARS:
            summary = summary[:self.MAX_SUMMARY_CHARS - 50] + "\n\n[TRUNCATED]"
        
        return summary
    
    def get_recent_summaries(self, limit: int = 5) -> List[StateSummary]:
        """Get most recent state summaries."""
        summaries = []
        
        # Get recent IDs
        recent_ids = self.state.get("summary_ids", [])[-limit:]
        
        for sid in reversed(recent_ids):
            path = self.summaries_dir / f"summary_{sid}.json"
            if path.exists():
                with open(path) as f:
                    data = json.load(f)
                    summaries.append(StateSummary(**data))
        
        return summaries
    
    def get_full_state(self, summary_id: str) -> Optional[Dict]:
        """Retrieve full state from a summary ID."""
        path = self.summaries_dir / f"full_{summary_id}.json"
        if path.exists():
            with open(path) as f:
                return json.load(f)
        return None
    
    def get_injection_context(self, limit: int = 3) -> str:
        """
        Generate context injection string from recent summaries.
        
        Use this at start of new sessions to restore continuity.
        """
        summaries = self.get_recent_summaries(limit)
        
        if not summaries:
            return ""
        
        lines = [
            "=== PREVIOUS SESSION CONTEXT ===",
            f"Restored from {len(summaries)} state compression(s)",
            ""
        ]
        
        for s in summaries:
            lines.append(f"--- Session {s.id} ({s.timestamp}) ---")
            if s.key_findings:
                lines.append("Findings: " + "; ".join(s.key_findings[:3]))
            if s.active_goals:
                lines.append("Goals: " + "; ".join(s.active_goals[:3]))
            if s.pending_actions:
                lines.append("Pending: " + "; ".join(s.pending_actions[:3]))
            lines.append("")
        
        return "\n".join(lines)
    
    def status(self) -> Dict:
        """Get VUA status."""
        summary_count = len(list(self.summaries_dir.glob("summary_*.json")))
        return {
            "total_compressions": self.state.get("total_compressions", 0),
            "total_tokens_saved": self.state.get("total_tokens_saved", 0),
            "last_compression": self.state.get("last_compression"),
            "stored_summaries": summary_count,
            "threshold": f"{self.COMPRESSION_THRESHOLD * 100}%"
        }


# CLI Interface
if __name__ == "__main__":
    import sys
    
    vua = VUAContextManager()
    
    if len(sys.argv) < 2:
        print("Usage: python vua_context.py <command>")
        print("Commands: status, recent, inject, test")
        sys.exit(1)
    
    command = sys.argv[1]
    
    if command == "status":
        print(json.dumps(vua.status(), indent=2))
    
    elif command == "recent":
        limit = int(sys.argv[2]) if len(sys.argv) > 2 else 5
        summaries = vua.get_recent_summaries(limit)
        for s in summaries:
            print(f"\n[{s.id}] {s.timestamp}")
            print(f"  Compression: {s.compression_ratio:.2%}")
            print(f"  Findings: {len(s.key_findings)}")
            print(f"  Goals: {len(s.active_goals)}")
    
    elif command == "inject":
        print(vua.get_injection_context())
    
    elif command == "test":
        # Test compression
        test_state = {
            "session": "test_123",
            "events": ["event1", "event2", "event3"] * 100,
            "metrics": {"a": 1, "b": 2, "c": 3},
            "large_data": "x" * 10000
        }
        
        summary = vua.compress_and_save(
            test_state,
            key_findings=["Found pattern X", "Discovered Y"],
            active_goals=["Complete Genesis", "Test integration"],
            pending_actions=["Run pulse", "Verify axioms"]
        )
        
        print(f"Compressed {summary.context_size_before} -> {summary.summary_size}")
        print(f"Ratio: {summary.compression_ratio:.2%}")
        print(f"ID: {summary.id}")
    
    else:
        print(f"Unknown command: {command}")