"""
cortex/core/context_compiler.py — Context Compiler (Issue #363).

Assembles retrieval results into injection-ready bundles with strict token
budgeting. The compiler classifies memory items into bundle types, allocates
token budget across bundles using priority-weighted percentages, compresses
items within each bundle to fit allocation, and enforces a hard total token cap.

Bundle Types (ordered by injection priority, highest first):
    boot           — session handoff, identity, continuity
    open_loops     — unresolved commitments, pending tasks
    task           — current task-relevant memories
    procedures     — how-to knowledge, runbooks
    recent_history — recent conversation context
    defaults       — stable facts, preferences, personality
    evidence       — supporting details, full claim text (lowest priority)

Architecture:
    ContextCompiler receives items (dicts with memory metadata), classifies
    each into a BundleType, allocates budget across bundles, compresses items
    to fit, and returns a CompiledContext ready for injection slot assignment.

Dependencies:
    - cortex.core.token_budget.TokenBudgetManager (token counting + compression)
    - cortex.core.task_mode.TaskMode (optional task mode for classification)
    - cortex.types.CompressionLevel (compression levels)

Dependents:
    - cortex.core.injection_slots.InjectionSlotManager (consumes CompiledContext)
    - cortex.api.plugin.CortexPlugin (orchestrates compile → inject pipeline)
"""

from __future__ import annotations

import logging
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from enum import Enum
from typing import Any, Dict, List, Optional

from cortex.core.task_mode import TaskMode
from cortex.core.token_budget import TokenBudgetManager
from cortex.types import CompressionLevel

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Bundle Types
# ---------------------------------------------------------------------------

class BundleType(str, Enum):
    """Memory bundle types ordered by injection priority (highest first).

    Priority order determines:
    1. Budget allocation (higher-priority bundles get their full allocation first)
    2. Compression order (lowest-priority bundles are compressed/dropped first)
    3. Injection order (higher-priority bundles are injected first)
    """

    BOOT = "boot"
    OPEN_LOOPS = "open_loops"
    TASK = "task"
    PROCEDURES = "procedures"
    RECENT_HISTORY = "recent_history"
    DEFAULTS = "defaults"
    EVIDENCE = "evidence"

    @property
    def priority(self) -> int:
        """Lower number = higher priority."""
        _priority = {
            BundleType.BOOT: 0,
            BundleType.OPEN_LOOPS: 1,
            BundleType.TASK: 2,
            BundleType.PROCEDURES: 3,
            BundleType.RECENT_HISTORY: 4,
            BundleType.DEFAULTS: 5,
            BundleType.EVIDENCE: 6,
        }
        return _priority[self]

    @classmethod
    def by_priority(cls) -> List[BundleType]:
        """Return bundle types ordered from highest to lowest priority."""
        return sorted(cls, key=lambda b: b.priority)

    @classmethod
    def reverse_priority(cls) -> List[BundleType]:
        """Return bundle types ordered from lowest to highest priority (compress first)."""
        return sorted(cls, key=lambda b: -b.priority)


# ---------------------------------------------------------------------------
# Budget allocation weights
# ---------------------------------------------------------------------------

# Default budget allocation weights per bundle type.
# These define what fraction of total budget each bundle type targets.
DEFAULT_BUDGET_WEIGHTS: Dict[BundleType, float] = {
    BundleType.BOOT: 0.15,
    BundleType.OPEN_LOOPS: 0.15,
    BundleType.TASK: 0.25,
    BundleType.PROCEDURES: 0.10,
    BundleType.RECENT_HISTORY: 0.15,
    BundleType.DEFAULTS: 0.10,
    BundleType.EVIDENCE: 0.10,
}


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

@dataclass
class Bundle:
    """A grouped collection of memory items of the same type.

    Attributes:
        bundle_type: The classification of this bundle.
        items: List of memory item dicts in this bundle.
        token_count: Total tokens consumed by items at current compression.
        compression_level: The compression level applied to items in this bundle.
    """

    bundle_type: BundleType
    items: List[Dict[str, Any]] = field(default_factory=list)
    token_count: int = 0
    compression_level: CompressionLevel = CompressionLevel.EVIDENCE


@dataclass
class CompiledContext:
    """Result of context compilation — bundles ready for injection.

    Attributes:
        bundles: Ordered list of bundles (by injection priority, highest first).
        total_tokens: Total tokens consumed across all bundles.
        budget_remaining: Tokens remaining from the original budget.
        metadata: Additional compilation metadata (e.g., items dropped, etc.).
    """

    bundles: List[Bundle] = field(default_factory=list)
    total_tokens: int = 0
    budget_remaining: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Config duck-type
# ---------------------------------------------------------------------------

class _ConfigProtocol:
    """Structural type expected from CortexConfig for context compiler."""
    context_compiler_enabled: bool
    context_compiler_max_bundles: int
    context_compiler_budget_overflow_strategy: str


# ---------------------------------------------------------------------------
# Context Compiler
# ---------------------------------------------------------------------------

class ContextCompiler:
    """Compile retrieval results into injection-ready bundles with token budgeting.

    The compiler performs four steps:
    1. Classify each item into a BundleType
    2. Allocate token budget across bundles (priority-weighted)
    3. Compress items within each bundle to fit allocation
    4. Enforce strict total token cap (compress then drop lowest priority)

    Args:
        token_budget: TokenBudgetManager for token counting and compression.
        config: CortexConfig (uses context_compiler.* fields).
    """

    # Relevance score threshold for task-mode classification
    TASK_RELEVANCE_THRESHOLD: float = 0.7

    # Stability threshold for defaults classification
    STABILITY_THRESHOLD: str = "stable"

    # Recent history window (24 hours)
    RECENT_HISTORY_HOURS: int = 24

    def __init__(
        self,
        token_budget: TokenBudgetManager,
        config: Optional[_ConfigProtocol] = None,
    ) -> None:
        self._token_budget = token_budget
        self._enabled = True
        self._max_bundles = 7
        self._overflow_strategy = "compress_then_drop"

        if config is not None:
            if hasattr(config, "context_compiler_enabled"):
                self._enabled = config.context_compiler_enabled
            if hasattr(config, "context_compiler_max_bundles"):
                self._max_bundles = config.context_compiler_max_bundles
            if hasattr(config, "context_compiler_budget_overflow_strategy"):
                self._overflow_strategy = config.context_compiler_budget_overflow_strategy

    def compile(
        self,
        items: List[Dict[str, Any]],
        budget_tokens: int,
        task_mode: Optional[TaskMode] = None,
    ) -> CompiledContext:
        """Compile items into injection-ready bundles within budget.

        Args:
            items: List of memory item dicts. Expected keys include:
                - content (str): The text content
                - type (str, optional): Memory type (identity, commitment, etc.)
                - source (str, optional): Source identifier (handoff, etc.)
                - relevance_score (float, optional): Retrieval relevance score
                - task_mode (str, optional): Associated task mode
                - stability (str, optional): Stability level
                - created_at (str, optional): ISO 8601 timestamp
                - session_id (str, optional): Session identifier
            budget_tokens: Total token budget for all bundles combined.
            task_mode: Optional current task mode for classification hints.

        Returns:
            CompiledContext with ordered bundles fitting within budget.
        """
        if not items or budget_tokens <= 0:
            return CompiledContext(
                bundles=[],
                total_tokens=0,
                budget_remaining=budget_tokens if budget_tokens > 0 else 0,
                metadata={"items_input": 0, "items_output": 0},
            )

        # Step 1: Classify items into bundle types
        classified: Dict[BundleType, List[Dict[str, Any]]] = defaultdict(list)
        for item in items:
            bundle_type = self._classify_item(item, task_mode)
            classified[bundle_type].append(item)

        # Step 2: Allocate budget across populated bundles
        budget_allocation = self._allocate_budget(classified, budget_tokens)

        # Step 2b: Sort items within each bundle by salience (high > medium > low)
        _salience_order = {"high": 0, "medium": 1, "low": 2}
        for bundle_type in classified:
            classified[bundle_type].sort(
                key=lambda x: _salience_order.get(x.get("salience", ""), 1)
            )

        # Step 3: Build bundles with compression to fit allocation
        bundles: List[Bundle] = []
        total_used = 0

        for bundle_type in BundleType.by_priority():
            if bundle_type not in classified:
                continue

            bundle_budget = budget_allocation.get(bundle_type, 0)
            if bundle_budget <= 0:
                continue

            bundle = self._compress_bundle(
                bundle_type,
                classified[bundle_type],
                bundle_budget,
            )

            if bundle.items and bundle.token_count > 0:
                bundles.append(bundle)
                total_used += bundle.token_count

        # Step 4: Enforce strict total token cap
        if total_used > budget_tokens:
            bundles, total_used = self._enforce_budget(
                bundles, budget_tokens
            )

        # Hard invariant: budget must NEVER be exceeded
        assert total_used <= budget_tokens, (
            f"Budget exceeded: {total_used} > {budget_tokens}"
        )

        return CompiledContext(
            bundles=bundles,
            total_tokens=total_used,
            budget_remaining=budget_tokens - total_used,
            metadata={
                "items_input": len(items),
                "items_output": sum(len(b.items) for b in bundles),
                "bundles_count": len(bundles),
                "overflow_strategy": self._overflow_strategy,
            },
        )

    # ------------------------------------------------------------------
    # Classification
    # ------------------------------------------------------------------

    def _classify_item(
        self,
        item: Dict[str, Any],
        task_mode: Optional[TaskMode] = None,
    ) -> BundleType:
        """Classify a single item into a BundleType using rules-based logic.

        Classification rules (checked in priority order):
        1. source="handoff" or type="identity" → boot
        2. type="commitment" or type="open_loop" → open_loops
        2b. memory_class="open_loop" → open_loops
        3. type="procedure" or type="runbook" → procedures
        3b. memory_class="procedural" → procedures
        4. High relevance + matching task mode → task
        5. Recent session (last 24h) → recent_history
        5b. memory_class="episodic" → recent_history (recency boost, bypasses 24h check)
        6. type="preference" or (type="fact" + high stability) → defaults
        7. Everything else → evidence
        """
        item_type = item.get("type", "")
        item_source = item.get("source", "")
        relevance = item.get("relevance_score", 0.0)
        stability = item.get("stability", "")
        created_at = item.get("created_at", "")

        # Rule 1: Boot — handoff or identity
        if item_source == "handoff" or item_type == "identity":
            return BundleType.BOOT

        # Rule 2: Open loops — commitments
        if item_type in ("commitment", "open_loop"):
            return BundleType.OPEN_LOOPS

        # Rule 2b: memory_class open_loop
        if item.get("memory_class") == "open_loop":
            return BundleType.OPEN_LOOPS

        # Rule 3: Procedures
        if item_type in ("procedure", "runbook"):
            return BundleType.PROCEDURES

        # Rule 3b: memory_class procedural
        if item.get("memory_class") == "procedural":
            return BundleType.PROCEDURES

        # Rule 4: Task — high relevance + task mode match
        if task_mode is not None and relevance >= self.TASK_RELEVANCE_THRESHOLD:
            item_task_mode = item.get("task_mode", "")
            if item_task_mode == task_mode.value or item_task_mode == "":
                return BundleType.TASK

        # Rule 5: Recent history — last 24h
        if created_at and self._is_recent(created_at):
            return BundleType.RECENT_HISTORY

        # Rule 5b: episodic memories get recency boost
        if item.get("memory_class") == "episodic":
            return BundleType.RECENT_HISTORY

        # Rule 6: Defaults — preferences or stable facts
        if item_type == "preference":
            return BundleType.DEFAULTS
        if item_type == "fact" and stability == self.STABILITY_THRESHOLD:
            return BundleType.DEFAULTS

        # Rule 7: Everything else → evidence
        return BundleType.EVIDENCE

    def _is_recent(self, created_at: str) -> bool:
        """Check if a timestamp is within the recent history window (24h)."""
        try:
            ts = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
            cutoff = datetime.now(timezone.utc) - timedelta(hours=self.RECENT_HISTORY_HOURS)
            return ts >= cutoff
        except (ValueError, TypeError):
            return False

    # ------------------------------------------------------------------
    # Budget allocation
    # ------------------------------------------------------------------

    def _allocate_budget(
        self,
        classified: Dict[BundleType, List[Dict[str, Any]]],
        total_budget: int,
    ) -> Dict[BundleType, int]:
        """Allocate token budget across populated bundles using priority weights.

        Only bundles that actually have items receive allocation. Unused
        budget from empty bundles is redistributed proportionally to
        populated bundles.

        Args:
            classified: Mapping of BundleType → items in that bundle.
            total_budget: Total tokens to allocate.

        Returns:
            Dict mapping BundleType → allocated tokens. Sum equals total_budget
            (or less, if rounding causes small shortfalls).
        """
        if not classified or total_budget <= 0:
            return {}

        # Only allocate to bundles that have items
        populated = {bt for bt in classified if classified[bt]}
        if not populated:
            return {}

        # Calculate weights for populated bundles only
        total_weight = sum(
            DEFAULT_BUDGET_WEIGHTS.get(bt, 0.0) for bt in populated
        )

        if total_weight <= 0:
            # Fallback: equal distribution
            per_bundle = total_budget // len(populated)
            return {bt: per_bundle for bt in populated}

        # Allocate proportionally based on weights
        allocation: Dict[BundleType, int] = {}
        allocated_so_far = 0

        sorted_bundles = sorted(populated, key=lambda b: b.priority)

        for i, bt in enumerate(sorted_bundles):
            weight = DEFAULT_BUDGET_WEIGHTS.get(bt, 0.0)
            if i == len(sorted_bundles) - 1:
                # Last bundle gets remainder to ensure sum == total_budget
                allocation[bt] = total_budget - allocated_so_far
            else:
                tokens = int(total_budget * weight / total_weight)
                allocation[bt] = tokens
                allocated_so_far += tokens

        return allocation

    # ------------------------------------------------------------------
    # Bundle compression
    # ------------------------------------------------------------------

    def _compress_bundle(
        self,
        bundle_type: BundleType,
        items: List[Dict[str, Any]],
        budget: int,
    ) -> Bundle:
        """Compress items within a bundle to fit the allocated budget.

        Uses TokenBudgetManager.compress_to_fit() for actual compression.
        Items are processed in order (assumed pre-sorted by relevance).

        Args:
            bundle_type: Type of this bundle.
            items: Memory items for this bundle.
            budget: Token budget allocated to this bundle.

        Returns:
            Bundle with compressed items fitting within budget.
        """
        if not items or budget <= 0:
            return Bundle(bundle_type=bundle_type)

        # Extract content strings for compression
        contents = [item.get("content", "") for item in items]
        contents = [c for c in contents if c]  # Skip empty content

        if not contents:
            return Bundle(bundle_type=bundle_type)

        # Use token budget manager to compress items to fit
        compressed = self._token_budget.compress_to_fit(
            items=contents,
            target_tokens=budget,
        )

        if not compressed:
            return Bundle(bundle_type=bundle_type)

        # Build the bundle with compressed items
        total_tokens = sum(ci.token_count for ci in compressed)

        # Determine overall compression level (worst/highest compression used)
        if compressed:
            worst_compression = min(ci.compression for ci in compressed)
        else:
            worst_compression = CompressionLevel.EVIDENCE

        # Rebuild items with compressed content
        compressed_items: List[Dict[str, Any]] = []
        for ci in compressed:
            compressed_items.append({
                "content": ci.content,
                "compression": ci.compression.name,
                "token_count": ci.token_count,
                "original_tokens": ci.original_tokens,
            })

        return Bundle(
            bundle_type=bundle_type,
            items=compressed_items,
            token_count=total_tokens,
            compression_level=worst_compression,
        )

    # ------------------------------------------------------------------
    # Over-budget enforcement
    # ------------------------------------------------------------------

    def _enforce_budget(
        self,
        bundles: List[Bundle],
        budget: int,
    ) -> tuple[List[Bundle], int]:
        """Enforce strict total token cap by compressing/dropping lowest-priority bundles.

        Strategy (compress_then_drop):
        1. Sort bundles by reverse priority (lowest priority first)
        2. Try to re-compress lowest-priority bundles at POINTER level
        3. If still over budget, drop lowest-priority bundles entirely

        Args:
            bundles: Current list of bundles.
            budget: Hard token budget cap.

        Returns:
            (filtered_bundles, total_tokens) where total_tokens <= budget.
        """
        total = sum(b.token_count for b in bundles)
        if total <= budget:
            return bundles, total

        # Sort by reverse priority for processing (lowest priority first)
        bundles_by_priority = sorted(bundles, key=lambda b: -b.bundle_type.priority)

        # Pass 1: Re-compress lowest-priority bundles at POINTER level
        for bundle in bundles_by_priority:
            if total <= budget:
                break

            # Re-compress this bundle's items at POINTER level
            contents = [item.get("content", "") for item in bundle.items if item.get("content")]
            if not contents:
                continue

            old_tokens = bundle.token_count
            # Target the actual remaining budget instead of arbitrary halving.
            # Calculate how many tokens this bundle needs to shed to fit.
            overshoot = total - budget
            target = max(1, old_tokens - overshoot)
            compressed = self._token_budget.compress_to_fit(
                items=contents,
                target_tokens=target,
                min_compression=CompressionLevel.POINTER,
            )

            new_tokens = sum(ci.token_count for ci in compressed)
            total = total - old_tokens + new_tokens
            bundle.token_count = new_tokens
            bundle.compression_level = CompressionLevel.POINTER
            bundle.items = [
                {
                    "content": ci.content,
                    "compression": ci.compression.name,
                    "token_count": ci.token_count,
                    "original_tokens": ci.original_tokens,
                }
                for ci in compressed
            ]

        # Pass 2: Drop lowest-priority bundles if still over
        if total > budget:
            kept: List[Bundle] = []
            running = 0
            # Re-sort by priority (highest first) to keep important bundles
            for bundle in sorted(bundles, key=lambda b: b.bundle_type.priority):
                if running + bundle.token_count <= budget:
                    kept.append(bundle)
                    running += bundle.token_count
                else:
                    logger.info(
                        "Dropping bundle %s (%d tokens) to fit budget",
                        bundle.bundle_type.value, bundle.token_count,
                    )
            bundles = kept
            total = running

        return bundles, total
