"""Context Window Monitor - Tracks token usage and warns before limits."""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path


@dataclass
class ContextStatus:
    """Status of context window usage."""

    status: str  # ok, warning, urgent, critical
    percent: int
    current_tokens: int
    max_tokens: int
    message: str | None
    actions: list[str] = field(default_factory=list)


class ContextMonitor:
    """Monitors context window usage and provides warnings."""

    DEFAULT_CONFIG = {
        "warn_threshold": 0.70,
        "urgent_threshold": 0.85,
        "critical_threshold": 0.95,
        "default_max_tokens": 128000,
    }

    def __init__(
        self,
        log_file: str | None = None,
        warn_threshold: float | None = None,
        urgent_threshold: float | None = None,
        critical_threshold: float | None = None,
    ) -> None:
        from oclawma.config import get_workspace_dir

        self.log_file = Path(log_file or str(get_workspace_dir() / "memory" / "context-usage.json"))
        self.warn_threshold = (
            warn_threshold if warn_threshold is not None else self.DEFAULT_CONFIG["warn_threshold"]
        )
        self.urgent_threshold = (
            urgent_threshold
            if urgent_threshold is not None
            else self.DEFAULT_CONFIG["urgent_threshold"]
        )
        self.critical_threshold = (
            critical_threshold
            if critical_threshold is not None
            else self.DEFAULT_CONFIG["critical_threshold"]
        )

    @staticmethod
    def estimate_tokens(text: str | dict | list) -> int:
        """Estimate tokens (rough approximation: 1 token ≈ 4 chars)."""
        if isinstance(text, (dict, list)):
            text = json.dumps(text)
        return max(1, len(str(text)) // 4)

    def _load_history(self) -> dict:
        """Load usage history."""
        try:
            with open(self.log_file, encoding="utf-8") as f:
                return json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            return {"sessions": [], "alerts": [], "totalTokens": 0}

    def _save_history(self, data: dict) -> None:
        """Save usage history."""
        self.log_file.parent.mkdir(parents=True, exist_ok=True)
        with open(self.log_file, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

    def check_context(self, current_tokens: int, max_tokens: int | None = None) -> ContextStatus:
        """Check context usage and return status."""
        max_toks = max_tokens or self.DEFAULT_CONFIG["default_max_tokens"]
        ratio = current_tokens / max_toks
        percent = int(ratio * 100)

        status = "ok"
        message = None
        actions: list[str] = []

        if ratio >= self.critical_threshold:
            status = "critical"
            message = f"⛔ CRITICAL: Context at {percent}% ({current_tokens:,}/{max_toks:,} tokens)"
            actions = [
                "IMMEDIATE: Archive this conversation",
                "Summarize completed work to MEMORY.md",
                "Start fresh session for next task",
            ]
        elif ratio >= self.urgent_threshold:
            status = "urgent"
            message = f"🚨 URGENT: Context at {percent}% ({current_tokens:,}/{max_toks:,} tokens)"
            actions = [
                "Drop old file references from context",
                "Summarize conversation history",
                "Consider wrapping up current task",
            ]
        elif ratio >= self.warn_threshold:
            status = "warning"
            message = f"⚠️ WARNING: Context at {percent}% ({current_tokens:,}/{max_toks:,} tokens)"
            actions = [
                "Monitor usage closely",
                "Avoid loading large files",
                "Be concise in responses",
            ]

        return ContextStatus(
            status=status,
            percent=percent,
            current_tokens=current_tokens,
            max_tokens=max_toks,
            message=message,
            actions=actions,
        )

    def log_usage(self, session_id: str, tokens: int, status: str) -> None:
        """Log usage for a session."""
        history = self._load_history()
        history["sessions"].append(
            {
                "timestamp": datetime.now().isoformat(),
                "sessionId": session_id,
                "tokens": tokens,
                "status": status,
            }
        )
        history["totalTokens"] = history.get("totalTokens", 0) + tokens
        self._save_history(history)

    def get_compression_tips(self, context_files: dict[str, dict] | None = None) -> list[str]:
        """Get recommendations for compression."""
        tips = []

        if context_files:
            # Sort files by size
            sorted_files = sorted(
                context_files.items(),
                key=lambda x: x[1].get("tokens", 0),
                reverse=True,
            )

            for file, info in sorted_files[:5]:
                if info.get("tokens", 0) > 1000:
                    tips.append(f"Consider removing: {file} ({info['tokens']} tokens)")

        return tips

    def get_history(self) -> dict:
        """Get usage history summary."""
        history = self._load_history()
        return {
            "total_sessions": len(history.get("sessions", [])),
            "total_tokens": history.get("totalTokens", 0),
            "recent_sessions": history.get("sessions", [])[-10:],
        }

    def print_history(self) -> None:
        """Print usage history."""
        history = self._load_history()

        print("\n📊 Context Usage History")
        print(f"Total sessions: {len(history.get('sessions', []))}")
        print(f"Total tokens: {history.get('totalTokens', 0):,}")
        print("\nRecent sessions:")

        for session in history.get("sessions", [])[-10:]:
            date = session["timestamp"].split("T")[0]
            print(f"  {date}: {session['tokens']:,} tokens ({session['status']})")


def main() -> None:
    """CLI entry point."""
    import sys

    monitor = ContextMonitor()
    command = sys.argv[1] if len(sys.argv) > 1 else "check"
    arg = sys.argv[2] if len(sys.argv) > 2 else None

    if command == "check":
        tokens = int(arg) if arg else 50000
        result = monitor.check_context(tokens)
        print(result.message or f"✅ Context healthy: {result.percent}%")
        if result.actions:
            print("\nRecommended actions:")
            for action in result.actions:
                print(f"  • {action}")

    elif command == "history":
        monitor.print_history()

    elif command == "estimate":
        text = arg or sys.stdin.read()
        tokens = monitor.estimate_tokens(text)
        print(f"Estimated tokens: {tokens:,}")

    else:
        print("Context Window Monitor")
        print("Usage:")
        print("  check [tokens]  - Check context status")
        print("  history         - Show usage history")
        print("  estimate [text] - Estimate token count")


if __name__ == "__main__":
    main()
