"""
cortex/hippo/budget.py — Cost tracking + circuit breaker for compaction.

Provides per-session and global token budget enforcement with sliding-window
tracking and an exponential-backoff circuit breaker for low-effectiveness
compaction runs.

Issue: #1178
"""

from __future__ import annotations

import logging
import threading
import time
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Optional

logger = logging.getLogger(__name__)


# ── Data types ─────────────────────────────────────────────────────────

@dataclass
class BudgetEvent:
    """Event emitted when a budget boundary is crossed."""
    type: str  # "warning", "exceeded", "breaker_tripped", "breaker_reset"
    session_id: str
    stats: dict[str, Any] = field(default_factory=dict)


@dataclass
class BudgetDecision:
    """Result of a budget check."""
    allowed: bool
    reason: str  # "ok", "budget_exceeded", "global_budget_exceeded", "circuit_breaker_tripped"
    remaining_tokens: int = 0
    cooldown_until: Optional[datetime] = None


@dataclass
class SessionBudgetStats:
    """Per-session budget statistics."""
    session_id: str
    tokens_used_last_hour: int = 0
    budget_limit: int = 0
    breaker_tripped: bool = False
    breaker_trip_count: int = 0
    cooldown_until: Optional[datetime] = None
    consecutive_low_reductions: int = 0
    last_reduction_ratios: list[float] = field(default_factory=list)


@dataclass
class BudgetStats:
    """Global + per-session budget snapshot."""
    global_tokens_used_last_hour: int = 0
    global_budget_limit: int = 0
    sessions: dict[str, SessionBudgetStats] = field(default_factory=dict)


# ── Internal: sliding window bucket ──────────────────────────────────

def _current_bucket() -> int:
    """Return the current 1-minute bucket index (Unix minutes)."""
    return int(time.time()) // 60


def _bucket_cutoff() -> int:
    """Bucket index for 1 hour ago."""
    return _current_bucket() - 60


# ── Session state ────────────────────────────────────────────────────

class _SessionState:
    """Per-session mutable state (guarded by parent lock)."""

    __slots__ = (
        "token_buckets",
        "reduction_history",
        "consecutive_low",
        "breaker_tripped",
        "cooldown_until",
        "trip_count",
    )

    def __init__(self) -> None:
        # minute-bucket → total tokens
        self.token_buckets: dict[int, int] = defaultdict(int)
        self.reduction_history: list[float] = []
        self.consecutive_low: int = 0
        self.breaker_tripped: bool = False
        self.cooldown_until: Optional[float] = None  # Unix timestamp
        self.trip_count: int = 0

    def tokens_last_hour(self) -> int:
        cutoff = _bucket_cutoff()
        return sum(v for k, v in self.token_buckets.items() if k >= cutoff)

    def prune(self) -> None:
        cutoff = _bucket_cutoff()
        self.token_buckets = defaultdict(
            int,
            {k: v for k, v in self.token_buckets.items() if k >= cutoff},
        )


# ── CompactionBudget ─────────────────────────────────────────────────

class CompactionBudget:
    """Per-session + global cost tracking with circuit breaker.

    Thread-safe via threading.Lock (works in both sync and async contexts).
    """

    def __init__(
        self,
        *,
        max_summarization_tokens_per_hour: int = 500_000,
        global_max_summarization_tokens_per_hour: int = 2_000_000,
        circuit_breaker_threshold: int = 3,
        circuit_breaker_cooldown_seconds: int = 300,
        min_compaction_reduction_ratio: float = 0.10,
        budget_warning_threshold: float = 0.80,
        on_budget_event: Optional[Callable[[str, BudgetEvent], Any]] = None,
    ) -> None:
        self._lock = threading.Lock()
        self._sessions: dict[str, _SessionState] = {}
        self._global_buckets: dict[int, int] = defaultdict(int)

        # Config
        self.max_tokens_per_hour = max_summarization_tokens_per_hour
        self.global_max_tokens_per_hour = global_max_summarization_tokens_per_hour
        self.breaker_threshold = circuit_breaker_threshold
        self.cooldown_seconds = circuit_breaker_cooldown_seconds
        self.min_reduction_ratio = min_compaction_reduction_ratio
        self.warning_threshold = budget_warning_threshold
        self.on_budget_event = on_budget_event

    # ── helpers ───────────────────────────────────────────────────

    def _get_session(self, session_id: str) -> _SessionState:
        if session_id not in self._sessions:
            self._sessions[session_id] = _SessionState()
        return self._sessions[session_id]

    def _global_tokens_last_hour(self) -> int:
        cutoff = _bucket_cutoff()
        return sum(v for k, v in self._global_buckets.items() if k >= cutoff)

    def _prune_global(self) -> None:
        cutoff = _bucket_cutoff()
        self._global_buckets = defaultdict(
            int,
            {k: v for k, v in self._global_buckets.items() if k >= cutoff},
        )

    def _fire_event(self, session_id: str, event: BudgetEvent) -> None:
        if self.on_budget_event is not None:
            try:
                self.on_budget_event(session_id, event)
            except Exception:
                logger.debug("on_budget_event callback failed", exc_info=True)

    # ── public API ────────────────────────────────────────────────

    def record_summarization(
        self, session_id: str, input_tokens: int, output_tokens: int
    ) -> None:
        """Record a summarization call's token usage."""
        total = input_tokens + output_tokens
        bucket = _current_bucket()

        with self._lock:
            sess = self._get_session(session_id)
            sess.token_buckets[bucket] += total
            self._global_buckets[bucket] += total

            # Periodic prune
            sess.prune()
            self._prune_global()

            # Check warning threshold
            used = sess.tokens_last_hour()
            if used >= self.max_tokens_per_hour * self.warning_threshold:
                if used < self.max_tokens_per_hour:
                    self._fire_event(
                        session_id,
                        BudgetEvent(
                            type="warning",
                            session_id=session_id,
                            stats={
                                "used": used,
                                "limit": self.max_tokens_per_hour,
                                "pct": round(used / self.max_tokens_per_hour, 2),
                            },
                        ),
                    )

    def check_budget(self, session_id: str) -> BudgetDecision:
        """Check if this session can proceed with summarization."""
        with self._lock:
            sess = self._get_session(session_id)
            sess.prune()
            self._prune_global()
            now = time.time()

            # 1. Circuit breaker check
            if sess.breaker_tripped:
                if sess.cooldown_until is not None and now < sess.cooldown_until:
                    return BudgetDecision(
                        allowed=False,
                        reason="circuit_breaker_tripped",
                        remaining_tokens=max(
                            0, self.max_tokens_per_hour - sess.tokens_last_hour()
                        ),
                        cooldown_until=datetime.fromtimestamp(
                            sess.cooldown_until, tz=timezone.utc
                        ),
                    )
                # Cooldown expired — allow one retry (breaker stays tripped
                # until a good compaction clears it via record_compaction_result)

            # 2. Per-session budget check
            session_used = sess.tokens_last_hour()
            if session_used >= self.max_tokens_per_hour:
                self._fire_event(
                    session_id,
                    BudgetEvent(
                        type="exceeded",
                        session_id=session_id,
                        stats={
                            "used": session_used,
                            "limit": self.max_tokens_per_hour,
                        },
                    ),
                )
                return BudgetDecision(
                    allowed=False,
                    reason="budget_exceeded",
                    remaining_tokens=0,
                )

            # 3. Global budget check
            global_used = self._global_tokens_last_hour()
            if global_used >= self.global_max_tokens_per_hour:
                self._fire_event(
                    session_id,
                    BudgetEvent(
                        type="exceeded",
                        session_id=session_id,
                        stats={
                            "used": global_used,
                            "limit": self.global_max_tokens_per_hour,
                            "scope": "global",
                        },
                    ),
                )
                return BudgetDecision(
                    allowed=False,
                    reason="global_budget_exceeded",
                    remaining_tokens=max(
                        0, self.max_tokens_per_hour - session_used
                    ),
                )

            return BudgetDecision(
                allowed=True,
                reason="ok",
                remaining_tokens=self.max_tokens_per_hour - session_used,
            )

    def record_compaction_result(
        self, session_id: str, reduction_ratio: float
    ) -> None:
        """Track compaction effectiveness for backoff logic.

        reduction_ratio: fraction of tokens saved (e.g. 0.25 = 25% reduction).
        """
        with self._lock:
            sess = self._get_session(session_id)
            sess.reduction_history.append(reduction_ratio)
            # Keep only last 20 entries
            if len(sess.reduction_history) > 20:
                sess.reduction_history = sess.reduction_history[-20:]

            if reduction_ratio < self.min_reduction_ratio:
                sess.consecutive_low += 1
            else:
                sess.consecutive_low = 0
                # Good compaction clears the breaker
                if sess.breaker_tripped:
                    sess.breaker_tripped = False
                    sess.trip_count = 0
                    sess.cooldown_until = None
                    logger.info(
                        "Circuit breaker reset for session %s (good reduction %.2f)",
                        session_id,
                        reduction_ratio,
                    )
                    self._fire_event(
                        session_id,
                        BudgetEvent(
                            type="breaker_reset",
                            session_id=session_id,
                            stats={"reduction_ratio": reduction_ratio},
                        ),
                    )

            # Trip the breaker?
            if sess.consecutive_low >= self.breaker_threshold and not sess.breaker_tripped:
                sess.breaker_tripped = True
                sess.trip_count += 1
                # Exponential cooldown: 1x, 2x, 4x, 8x ...
                multiplier = 2 ** (sess.trip_count - 1)
                cooldown = self.cooldown_seconds * multiplier
                sess.cooldown_until = time.time() + cooldown
                logger.warning(
                    "Circuit breaker TRIPPED for session %s "
                    "(consecutive_low=%d, trip_count=%d, cooldown=%ds)",
                    session_id,
                    sess.consecutive_low,
                    sess.trip_count,
                    cooldown,
                )
                self._fire_event(
                    session_id,
                    BudgetEvent(
                        type="breaker_tripped",
                        session_id=session_id,
                        stats={
                            "consecutive_low": sess.consecutive_low,
                            "trip_count": sess.trip_count,
                            "cooldown_seconds": cooldown,
                        },
                    ),
                )

    def reset_session(self, session_id: str) -> None:
        """Manual reset after circuit breaker trips."""
        with self._lock:
            if session_id in self._sessions:
                sess = self._sessions[session_id]
                sess.breaker_tripped = False
                sess.cooldown_until = None
                sess.consecutive_low = 0
                sess.trip_count = 0
                logger.info("Session %s manually reset", session_id)
                self._fire_event(
                    session_id,
                    BudgetEvent(
                        type="breaker_reset",
                        session_id=session_id,
                        stats={"manual": True},
                    ),
                )

    def get_stats(self) -> BudgetStats:
        """Global + per-session budget stats for monitoring."""
        with self._lock:
            self._prune_global()
            session_stats: dict[str, SessionBudgetStats] = {}
            for sid, sess in self._sessions.items():
                sess.prune()
                cooldown_dt = None
                if sess.cooldown_until is not None:
                    cooldown_dt = datetime.fromtimestamp(
                        sess.cooldown_until, tz=timezone.utc
                    )
                session_stats[sid] = SessionBudgetStats(
                    session_id=sid,
                    tokens_used_last_hour=sess.tokens_last_hour(),
                    budget_limit=self.max_tokens_per_hour,
                    breaker_tripped=sess.breaker_tripped,
                    breaker_trip_count=sess.trip_count,
                    cooldown_until=cooldown_dt,
                    consecutive_low_reductions=sess.consecutive_low,
                    last_reduction_ratios=list(sess.reduction_history[-5:]),
                )
            return BudgetStats(
                global_tokens_used_last_hour=self._global_tokens_last_hour(),
                global_budget_limit=self.global_max_tokens_per_hour,
                sessions=session_stats,
            )
