"""
cortex/core/retrieval_policy.py — Retrieval Policy Engine (Issue #362).

Maps task modes to retrieval channel policies. Each policy defines which
retrieval channels to query, budget allocation weights, confidence thresholds,
and max results per channel.

Channels:
    semantic       — vector similarity search
    keyword        — BM25 full-text search
    working_memory — session-scoped in-memory cache
    entity_graph   — entity relationship graph traversal
    temporal       — time-range filtered retrieval

Architecture:
    PolicyEngine takes a TaskModeResult and returns a RetrievalPolicy.
    If the mode confidence is below the configured threshold, FALLBACK_POLICY
    is used instead.

Dependencies: cortex.core.task_mode, cortex.config
Dependents: cortex.core.retrieval
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Dict, Optional

from cortex.config import CortexConfig
from cortex.core.task_mode import TaskMode, TaskModeResult

logger = logging.getLogger(__name__)


@dataclass
class ChannelPolicy:
    """Policy for a single retrieval channel.

    Attributes:
        enabled: Whether this channel should be queried.
        weight: Budget allocation weight (0.0–1.0). Weights across all
                enabled channels are normalized to sum to 1.0.
        min_confidence: Minimum result confidence to include from this channel.
        max_results: Maximum results to return from this channel.
    """

    enabled: bool = True
    weight: float = 0.2
    min_confidence: float = 0.0
    max_results: int = 10


@dataclass
class RetrievalPolicy:
    """Complete retrieval policy for a task mode.

    Defines channel configurations, overall constraints, and the source
    mode that generated this policy.

    Attributes:
        mode: The task mode this policy was generated for (or None for fallback).
        channels: Per-channel policy configuration.
        total_max_results: Hard cap on total results across all channels.
        description: Human-readable description of the policy.
    """

    mode: Optional[TaskMode] = None
    channels: Dict[str, ChannelPolicy] = field(default_factory=dict)
    total_max_results: int = 20
    description: str = ""

    def enabled_channels(self) -> list[str]:
        """Return names of enabled channels, sorted by weight (descending)."""
        return sorted(
            [name for name, ch in self.channels.items() if ch.enabled],
            key=lambda n: self.channels[n].weight,
            reverse=True,
        )

    def normalized_weights(self) -> Dict[str, float]:
        """Return normalized weights for enabled channels (sum to 1.0)."""
        enabled = {n: ch.weight for n, ch in self.channels.items() if ch.enabled}
        total = sum(enabled.values())
        if total == 0:
            return {}
        return {n: w / total for n, w in enabled.items()}


# ---------------------------------------------------------------------------
# Channel names (constants)
# ---------------------------------------------------------------------------

CH_SEMANTIC = "semantic"
CH_KEYWORD = "keyword"
CH_WORKING_MEMORY = "working_memory"
CH_ENTITY_GRAPH = "entity_graph"
CH_TEMPORAL = "temporal"

ALL_CHANNELS = [CH_SEMANTIC, CH_KEYWORD, CH_WORKING_MEMORY, CH_ENTITY_GRAPH, CH_TEMPORAL]


def _make_channels(
    semantic: tuple = (True, 0.3, 0.0, 10),
    keyword: tuple = (True, 0.2, 0.0, 10),
    working_memory: tuple = (True, 0.2, 0.0, 5),
    entity_graph: tuple = (True, 0.15, 0.0, 8),
    temporal: tuple = (True, 0.15, 0.0, 8),
) -> Dict[str, ChannelPolicy]:
    """Helper to build channel dict from tuples (enabled, weight, min_conf, max_results)."""
    return {
        CH_SEMANTIC: ChannelPolicy(*semantic),
        CH_KEYWORD: ChannelPolicy(*keyword),
        CH_WORKING_MEMORY: ChannelPolicy(*working_memory),
        CH_ENTITY_GRAPH: ChannelPolicy(*entity_graph),
        CH_TEMPORAL: ChannelPolicy(*temporal),
    }


# ---------------------------------------------------------------------------
# Default policies per mode
# ---------------------------------------------------------------------------

DEFAULT_POLICIES: Dict[TaskMode, RetrievalPolicy] = {
    TaskMode.CONTINUE_PROJECT: RetrievalPolicy(
        mode=TaskMode.CONTINUE_PROJECT,
        channels=_make_channels(
            semantic=(True, 0.25, 0.0, 10),
            keyword=(True, 0.15, 0.0, 8),
            working_memory=(True, 0.30, 0.0, 10),  # high: resume context
            entity_graph=(True, 0.15, 0.0, 8),
            temporal=(True, 0.15, 0.0, 8),          # recent work
        ),
        total_max_results=25,
        description="Resume ongoing work — prioritize working memory and recent context",
    ),
    TaskMode.DESIGN: RetrievalPolicy(
        mode=TaskMode.DESIGN,
        channels=_make_channels(
            semantic=(True, 0.35, 0.0, 12),         # high: specs, decisions
            keyword=(True, 0.20, 0.0, 10),
            working_memory=(True, 0.10, 0.0, 5),
            entity_graph=(True, 0.25, 0.0, 10),     # high: relationships
            temporal=(True, 0.10, 0.0, 5),
        ),
        total_max_results=25,
        description="Architecture/planning — prioritize specs, decisions, constraints",
    ),
    TaskMode.WRITE: RetrievalPolicy(
        mode=TaskMode.WRITE,
        channels=_make_channels(
            semantic=(True, 0.35, 0.0, 10),         # examples, references
            keyword=(True, 0.25, 0.0, 10),          # style guides
            working_memory=(True, 0.15, 0.0, 5),
            entity_graph=(True, 0.10, 0.0, 5),
            temporal=(True, 0.15, 0.0, 5),
        ),
        total_max_results=20,
        description="Content creation — prioritize examples, style guides, references",
    ),
    TaskMode.PLAN: RetrievalPolicy(
        mode=TaskMode.PLAN,
        channels=_make_channels(
            semantic=(True, 0.30, 0.0, 10),
            keyword=(True, 0.15, 0.0, 8),
            working_memory=(True, 0.15, 0.0, 5),
            entity_graph=(True, 0.15, 0.0, 8),
            temporal=(True, 0.25, 0.0, 10),          # high: timelines, commitments
        ),
        total_max_results=25,
        description="Strategic planning — prioritize goals, commitments, timelines",
    ),
    TaskMode.DECIDE: RetrievalPolicy(
        mode=TaskMode.DECIDE,
        channels=_make_channels(
            semantic=(True, 0.35, 0.0, 12),         # high: past decisions
            keyword=(True, 0.20, 0.0, 10),
            working_memory=(True, 0.15, 0.0, 5),
            entity_graph=(True, 0.20, 0.0, 8),      # constraints, relationships
            temporal=(True, 0.10, 0.0, 5),
        ),
        total_max_results=25,
        description="Decision making — prioritize pros/cons, past decisions, constraints",
    ),
    TaskMode.DEBUG: RetrievalPolicy(
        mode=TaskMode.DEBUG,
        channels=_make_channels(
            semantic=(True, 0.20, 0.0, 8),
            keyword=(True, 0.25, 0.0, 10),          # high: error patterns
            working_memory=(True, 0.30, 0.0, 10),   # high: recent context
            entity_graph=(True, 0.10, 0.0, 5),
            temporal=(True, 0.15, 0.0, 8),           # recent changes
        ),
        total_max_results=25,
        description="Troubleshooting — prioritize error context, logs, recent changes",
    ),
    TaskMode.RECALL_HISTORY: RetrievalPolicy(
        mode=TaskMode.RECALL_HISTORY,
        channels=_make_channels(
            semantic=(True, 0.35, 0.0, 12),         # high: historical facts
            keyword=(True, 0.25, 0.0, 10),
            working_memory=(True, 0.10, 0.0, 5),
            entity_graph=(True, 0.15, 0.0, 8),
            temporal=(True, 0.15, 0.0, 8),
        ),
        total_max_results=25,
        description="Memory lookup — prioritize historical facts, conversations",
    ),
    TaskMode.FUTURE_COMMITMENT: RetrievalPolicy(
        mode=TaskMode.FUTURE_COMMITMENT,
        channels=_make_channels(
            semantic=(True, 0.20, 0.0, 8),
            keyword=(True, 0.20, 0.0, 8),
            working_memory=(True, 0.15, 0.0, 5),
            entity_graph=(True, 0.15, 0.0, 8),
            temporal=(True, 0.30, 0.0, 12),          # high: deadlines, schedule
        ),
        total_max_results=25,
        description="Scheduling/promises — prioritize commitments, deadlines, open loops",
    ),
}

FALLBACK_POLICY = RetrievalPolicy(
    mode=None,
    channels=_make_channels(
        semantic=(True, 0.30, 0.0, 10),
        keyword=(True, 0.25, 0.0, 10),
        working_memory=(True, 0.15, 0.0, 5),
        entity_graph=(True, 0.15, 0.0, 8),
        temporal=(True, 0.15, 0.0, 8),
    ),
    total_max_results=20,
    description="Fallback — conservative balanced policy for unknown/low-confidence modes",
)


class PolicyEngine:
    """Maps TaskModeResults to RetrievalPolicies.

    Uses DEFAULT_POLICIES for known modes with sufficient confidence.
    Falls back to FALLBACK_POLICY when confidence is below threshold or
    the mode is unknown.

    Config overrides (via CortexConfig) can adjust the confidence threshold
    and default mode.
    """

    def __init__(self, config: Optional[CortexConfig] = None) -> None:
        self._confidence_threshold = 0.3
        self._enabled = True

        if config is not None:
            self._confidence_threshold = getattr(
                config, "retrieval_policy_confidence_threshold", 0.3
            )
            self._enabled = getattr(config, "retrieval_policy_enabled", True)

    @property
    def enabled(self) -> bool:
        """Whether the policy engine is active."""
        return self._enabled

    def get_policy(self, mode_result: TaskModeResult) -> RetrievalPolicy:
        """Return the retrieval policy for a classified task mode.

        If the engine is disabled, confidence is below threshold, or the
        mode has no configured policy, returns FALLBACK_POLICY.

        Args:
            mode_result: Classification result from TaskModeClassifier.

        Returns:
            RetrievalPolicy to use for retrieval routing.
        """
        if not self._enabled:
            logger.debug("PolicyEngine disabled — using fallback policy")
            return FALLBACK_POLICY

        if mode_result.confidence < self._confidence_threshold:
            logger.debug(
                "PolicyEngine: confidence %.3f < threshold %.3f — using fallback",
                mode_result.confidence,
                self._confidence_threshold,
            )
            return FALLBACK_POLICY

        policy = DEFAULT_POLICIES.get(mode_result.mode)
        if policy is None:
            logger.warning(
                "PolicyEngine: no policy for mode %s — using fallback",
                mode_result.mode,
            )
            return FALLBACK_POLICY

        logger.debug(
            "PolicyEngine: mode=%s confidence=%.3f → policy=%s",
            mode_result.mode.value,
            mode_result.confidence,
            policy.description,
        )
        return policy
