"""
cortex/core/context_calculator.py — Context Window Calculator (Issue #329).

Real-time token tracking across all memory tiers, giving operators and the
agent visibility into how the context window is being used.

Tracks token usage by category (boot, working_memory, episodic, semantic,
procedural, system, user, reserve) and provides both budget allocation
(how much each category *should* use) and usage recording (how much it
*actually* uses).

Design decisions:
    - Complementary to token_budget.py (does NOT replace it). TokenBudgetManager
      manages slot-level memory allocation; ContextWindowCalculator tracks the
      full context window across ALL tiers including system/user/reserve.
    - Allocation strategy "proportional" distributes (total - reserve) across
      categories using configurable weights. Weights must sum to ~1.0.
    - category_weights in config exclude reserve (reserve is held back first).
    - Metrics integration: record_usage() emits to MetricsCollector if wired in.
    - Thread-safe within asyncio (uses no shared mutable state between instances).
    - reset() is intended for new-turn housekeeping.

Dependencies: cortex.config.CortexConfig, cortex.core.metrics.MetricsCollector
Dependents: cortex.api.routes.context (HTTP endpoint)
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from cortex.config import CortexConfig
from cortex.core.metrics import MetricsCollector

import logging

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Known categories (ordered for display)
# ---------------------------------------------------------------------------

CATEGORIES = [
    "boot",
    "working_memory",
    "episodic",
    "semantic",
    "procedural",
    "system",
    "user",
    "reserve",
]

# Default allocation weights (excluding reserve — reserve is controlled by
# context_calculator.reserve_fraction and subtracted before distribution).
DEFAULT_CATEGORY_WEIGHTS: Dict[str, float] = {
    "boot": 0.10,
    "system": 0.10,
    "user": 0.25,
    "episodic": 0.15,
    "semantic": 0.15,
    "working_memory": 0.05,
    "procedural": 0.05,
}
# Note: weights above sum to 0.85, which matches 1.0 − 0.15 reserve default.


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

@dataclass
class CategoryUsage:
    """Token usage for a single category."""

    budget: int
    """Allocated token budget for this category."""

    used: int
    """Actual tokens recorded for this category."""

    remaining: int
    """Budget minus used (may be negative if over budget)."""

    item_count: int
    """Number of items recorded in this category."""

    utilization: float
    """used / budget (0.0 when budget is 0 to avoid division by zero)."""


@dataclass
class ContextAllocation:
    """Result of a budget allocation calculation."""

    total_tokens: int
    """Total context window size passed to allocate()."""

    allocations: Dict[str, int]
    """Per-category token budget."""

    reserve_tokens: int
    """Tokens held back for model output (not in any category budget)."""

    allocation_strategy: str
    """Name of the strategy used to compute allocations."""


@dataclass
class ContextUsageReport:
    """Current token usage snapshot across all categories."""

    total_budget: int
    """Total context window size (same as ContextAllocation.total_tokens)."""

    total_used: int
    """Sum of tokens recorded across all non-reserve categories."""

    total_remaining: int
    """total_budget − total_used."""

    utilization: float
    """total_used / total_budget (0.0 when budget is 0)."""

    by_category: Dict[str, CategoryUsage]
    """Per-category usage breakdown."""

    warnings: List[str]
    """Budget violation messages (one per over-budget category)."""


# ---------------------------------------------------------------------------
# Internal usage record
# ---------------------------------------------------------------------------

@dataclass
class _UsageEntry:
    tokens: int = 0
    item_count: int = 0


# ---------------------------------------------------------------------------
# ContextWindowCalculator
# ---------------------------------------------------------------------------

class ContextWindowCalculator:
    """
    Real-time context window token tracker across all memory tiers.

    Usage::

        calc = ContextWindowCalculator(config)
        allocation = calc.allocate(total_context=128_000)

        calc.record_usage("boot", tokens=1200, item_count=3)
        calc.record_usage("semantic", tokens=800, item_count=5)

        report = calc.get_usage()
        print(f"Utilization: {report.utilization:.1%}")
        print(f"Over budget: {calc.is_over_budget('boot')}")
    """

    def __init__(
        self,
        config: CortexConfig,
        metrics: Optional[MetricsCollector] = None,
    ) -> None:
        self._config = config
        self._metrics = metrics

        # Current allocation (set by allocate())
        self._allocation: Optional[ContextAllocation] = None

        # Per-category usage
        self._usage: Dict[str, _UsageEntry] = {cat: _UsageEntry() for cat in CATEGORIES}

        # Timestamp of last reset (for snapshot)
        self._reset_at: float = time.time()

    # ------------------------------------------------------------------ #
    # Allocation
    # ------------------------------------------------------------------ #

    def allocate(self, total_context: int) -> ContextAllocation:
        """Compute per-category budget allocation.

        Args:
            total_context: Total context window size in tokens.

        Returns:
            ContextAllocation with per-category budgets and reserve.
        """
        strategy = getattr(
            self._config, "context_calculator_allocation_strategy", "proportional"
        )
        reserve_fraction = getattr(
            self._config, "context_calculator_reserve_fraction", 0.15
        )
        category_weights: Dict[str, float] = dict(
            getattr(self._config, "context_calculator_category_weights", None)
            or DEFAULT_CATEGORY_WEIGHTS
        )

        reserve_tokens = max(0, int(total_context * reserve_fraction))
        distributable = max(0, total_context - reserve_tokens)

        allocations: Dict[str, int] = {}

        if strategy == "proportional" and category_weights and distributable > 0:
            # Normalise weights (defensive: ensure they sum to 1.0)
            total_weight = sum(category_weights.values())
            if total_weight <= 0:
                total_weight = 1.0  # fallback: avoid divide-by-zero

            # First pass: floor allocations
            running = 0
            floored: Dict[str, int] = {}
            for cat, w in category_weights.items():
                floored[cat] = int(distributable * (w / total_weight))
                running += floored[cat]

            # Distribute any leftover tokens (due to int truncation) to
            # categories in weight-descending order
            leftover = distributable - running
            if leftover > 0:
                sorted_cats = sorted(category_weights, key=lambda c: category_weights[c], reverse=True)
                for cat in sorted_cats:
                    if leftover <= 0:
                        break
                    floored[cat] += 1
                    leftover -= 1

            allocations = floored
        else:
            # Even split among non-reserve categories (fallback)
            cats = [c for c in CATEGORIES if c != "reserve"]
            per_cat = distributable // len(cats) if cats else 0
            allocations = {c: per_cat for c in cats}

        # Ensure reserve is in allocations (read-only tracking bucket)
        allocations["reserve"] = reserve_tokens

        # Zero-out any unknown categories (safety)
        for cat in CATEGORIES:
            allocations.setdefault(cat, 0)

        self._allocation = ContextAllocation(
            total_tokens=total_context,
            allocations=allocations,
            reserve_tokens=reserve_tokens,
            allocation_strategy=strategy,
        )
        return self._allocation

    # ------------------------------------------------------------------ #
    # Usage recording
    # ------------------------------------------------------------------ #

    def record_usage(self, category: str, tokens: int, item_count: int = 1) -> None:
        """Record actual token usage for a category.

        Args:
            category: One of the known category names.
            tokens: Number of tokens used.
            item_count: Number of items contributing to this usage.
        """
        if category not in self._usage:
            # Auto-register unknown category (defensive)
            self._usage[category] = _UsageEntry()

        entry = self._usage[category]
        entry.tokens += tokens
        entry.item_count += item_count

        # Emit to MetricsCollector if wired in
        if self._metrics is not None and self._metrics.enabled:
            import asyncio
            try:
                # Use get_running_loop() (Python 3.10+) instead of deprecated
                # get_event_loop(). Raises RuntimeError when no loop is running,
                # which we catch below to gracefully skip metrics in sync contexts.
                loop = asyncio.get_running_loop()
                loop.create_task(
                    self._metrics.increment(
                        "context_tokens_recorded",
                        labels={"category": category},
                        amount=float(tokens),
                    )
                )
            except RuntimeError:
                # No running event loop (sync context) — skip metrics emission
                pass

    # ------------------------------------------------------------------ #
    # Query methods
    # ------------------------------------------------------------------ #

    def get_usage(self) -> ContextUsageReport:
        """Return the current usage breakdown.

        Returns:
            ContextUsageReport with per-category breakdown and warnings.
        """
        allocations = (
            self._allocation.allocations if self._allocation is not None else {}
        )
        total_budget = self._allocation.total_tokens if self._allocation is not None else 0

        by_category: Dict[str, CategoryUsage] = {}
        warnings: List[str] = []

        total_used = 0

        for cat in CATEGORIES:
            budget = allocations.get(cat, 0)
            entry = self._usage.get(cat, _UsageEntry())
            used = entry.tokens
            remaining = budget - used
            utilization = (used / budget) if budget > 0 else 0.0

            by_category[cat] = CategoryUsage(
                budget=budget,
                used=used,
                remaining=remaining,
                item_count=entry.item_count,
                utilization=utilization,
            )

            if cat != "reserve":
                total_used += used

            if used > budget and budget > 0:
                warnings.append(
                    f"Category '{cat}' over budget: used {used} / budget {budget} tokens "
                    f"({utilization:.1%} utilization)"
                )

        total_remaining = total_budget - total_used
        overall_utilization = (total_used / total_budget) if total_budget > 0 else 0.0

        return ContextUsageReport(
            total_budget=total_budget,
            total_used=total_used,
            total_remaining=total_remaining,
            utilization=overall_utilization,
            by_category=by_category,
            warnings=warnings,
        )

    def get_remaining(self, category: str) -> int:
        """Return remaining token budget for a category.

        Args:
            category: Category name.

        Returns:
            Remaining tokens (may be negative if over budget). Returns 0 if
            no allocation has been computed yet.
        """
        if self._allocation is None:
            return 0
        budget = self._allocation.allocations.get(category, 0)
        used = self._usage.get(category, _UsageEntry()).tokens
        return budget - used

    def is_over_budget(self, category: str) -> bool:
        """Check whether a category has exceeded its allocation.

        Args:
            category: Category name.

        Returns:
            True if recorded usage exceeds allocated budget.
        """
        if self._allocation is None:
            return False
        budget = self._allocation.allocations.get(category, 0)
        if budget == 0:
            return False
        used = self._usage.get(category, _UsageEntry()).tokens
        return used > budget

    def utilization(self) -> float:
        """Overall context window utilization (0.0–1.0).

        Returns:
            Fraction of total_budget consumed by recorded usage.
            Returns 0.0 if no allocation has been computed.
        """
        if self._allocation is None or self._allocation.total_tokens == 0:
            return 0.0
        report = self.get_usage()
        return report.utilization

    def reset(self) -> None:
        """Clear all recorded usage (call at start of each new turn)."""
        for entry in self._usage.values():
            entry.tokens = 0
            entry.item_count = 0
        self._reset_at = time.time()

    def snapshot(self) -> Dict[str, Any]:
        """Return a serialisable state dict for logging/metrics.

        Returns:
            Dict containing allocation, usage, utilization, and metadata.
        """
        report = self.get_usage()

        allocation_dict: Dict[str, Any] = {}
        if self._allocation is not None:
            allocation_dict = {
                "total_tokens": self._allocation.total_tokens,
                "reserve_tokens": self._allocation.reserve_tokens,
                "allocation_strategy": self._allocation.allocation_strategy,
                "allocations": dict(self._allocation.allocations),
            }

        by_category_dict: Dict[str, Any] = {}
        for cat, cu in report.by_category.items():
            by_category_dict[cat] = {
                "budget": cu.budget,
                "used": cu.used,
                "remaining": cu.remaining,
                "item_count": cu.item_count,
                "utilization": round(cu.utilization, 4),
            }

        return {
            "allocation": allocation_dict,
            "usage": {
                "total_budget": report.total_budget,
                "total_used": report.total_used,
                "total_remaining": report.total_remaining,
                "utilization": round(report.utilization, 4),
                "warnings": list(report.warnings),
                "by_category": by_category_dict,
            },
            "reset_at": self._reset_at,
        }
