"""Confidence scoring utilities for retrieval.

Provides confidence scoring for retrieval results and low-confidence
callback hooks for triggering "ask follow-up" behavior.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Any, Callable, Dict, List, Optional, Protocol

logger = logging.getLogger(__name__)


class LowConfidenceCallback(Protocol):
    """Protocol for low-confidence callback handlers.

    Callbacks are invoked when retrieval confidence falls below threshold.
    Use this to trigger "ask follow-up" behavior or clarification requests.
    """

    def __call__(
        self,
        query: str,
        results: List[Any],
        confidence: float,
        threshold: float,
        context: Dict[str, Any],
    ) -> Optional[Dict[str, Any]]:
        """Handle low-confidence retrieval result.

        Args:
            query: The original query string
            results: Retrieved results (may be empty or low quality)
            confidence: The computed confidence score
            threshold: The threshold that was not met
            context: Additional context (namespace, weights, etc.)

        Returns:
            Optional dict with follow-up action, e.g.:
            {"action": "ask_clarification", "prompt": "Did you mean...?"}
        """
        ...


@dataclass
class LowConfidenceAction:
    """Result from a low-confidence callback."""

    action: str  # "ask_clarification", "expand_search", "return_anyway", etc.
    prompt: Optional[str] = None
    suggested_queries: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ConfidenceConfig:
    """Configuration for confidence scoring and callbacks.

    Attributes:
        threshold: Minimum confidence score (0.0-1.0) before triggering callback
        weights: Weight distribution for confidence components
        callbacks: List of callbacks to invoke on low confidence
        namespace_thresholds: Per-namespace threshold overrides
    """

    threshold: float = 0.5
    weights: Optional["ConfidenceWeights"] = None
    callbacks: List[LowConfidenceCallback] = field(default_factory=list)
    namespace_thresholds: Dict[str, float] = field(default_factory=dict)

    def get_threshold(self, namespace: Optional[str] = None) -> float:
        """Get threshold for a specific namespace."""
        if namespace and namespace in self.namespace_thresholds:
            return self.namespace_thresholds[namespace]
        return self.threshold


@dataclass(frozen=True)
class ConfidenceWeights:
    alpha: float = 0.4
    beta: float = 0.3
    gamma: float = 0.2
    delta: float = 0.1

    def normalised(self) -> "ConfidenceWeights":
        total = self.alpha + self.beta + self.gamma + self.delta
        if total <= 0:
            raise ValueError("Confidence weights must sum to a positive value")
        return ConfidenceWeights(
            alpha=self.alpha / total,
            beta=self.beta / total,
            gamma=self.gamma / total,
            delta=self.delta / total,
        )


def freshness_score(
    *,
    timestamp: datetime,
    now: datetime | None = None,
    horizon: timedelta | None = None,
) -> float:
    reference = now or datetime.now()
    if horizon is None:
        horizon = timedelta(days=7)
    age = reference - timestamp
    if age.total_seconds() <= 0:
        return 1.0
    if age >= horizon:
        return 0.0
    return max(0.0, 1.0 - (age.total_seconds() / horizon.total_seconds()))


def compute_confidence(
    *,
    vector_similarity: float,
    graph_path_score: float,
    trust_score: float,
    freshness: float,
    weights: ConfidenceWeights | None = None,
) -> float:
    norm_weights = (weights or ConfidenceWeights()).normalised()
    return (
        norm_weights.alpha * vector_similarity
        + norm_weights.beta * graph_path_score
        + norm_weights.gamma * trust_score
        + norm_weights.delta * freshness
    )


class ConfidenceEvaluator:
    """Evaluates retrieval confidence and triggers callbacks on low confidence.

    Example usage:
        def ask_followup(query, results, confidence, threshold, context):
            return {"action": "ask_clarification", "prompt": f"Did you mean {query}?"}

        config = ConfidenceConfig(threshold=0.6, callbacks=[ask_followup])
        evaluator = ConfidenceEvaluator(config)

        result = evaluator.evaluate(
            query="user preferences",
            results=retrieved_episodes,
            vector_similarity=0.45,
            graph_path_score=0.3,
            trust_score=0.8,
            freshness=0.9,
        )

        if result.is_low_confidence:
            # Handle low confidence - e.g., ask user for clarification
            print(result.callback_results)
    """

    def __init__(self, config: Optional[ConfidenceConfig] = None):
        """Initialize the evaluator with configuration.

        Args:
            config: Confidence configuration. Uses defaults if not provided.
        """
        self.config = config or ConfidenceConfig()

    def evaluate(
        self,
        *,
        query: str,
        results: List[Any],
        vector_similarity: float,
        graph_path_score: float,
        trust_score: float,
        freshness: float,
        namespace: Optional[str] = None,
        extra_context: Optional[Dict[str, Any]] = None,
    ) -> "ConfidenceResult":
        """Evaluate confidence and trigger callbacks if below threshold.

        Args:
            query: The original query string
            results: Retrieved results
            vector_similarity: Semantic similarity score (0-1)
            graph_path_score: Graph path relevance score (0-1)
            trust_score: Trust/reliability score (0-1)
            freshness: Recency score (0-1)
            namespace: Optional namespace for threshold lookup
            extra_context: Additional context to pass to callbacks

        Returns:
            ConfidenceResult with score, threshold status, and callback results
        """
        weights = self.config.weights
        confidence = compute_confidence(
            vector_similarity=vector_similarity,
            graph_path_score=graph_path_score,
            trust_score=trust_score,
            freshness=freshness,
            weights=weights,
        )

        threshold = self.config.get_threshold(namespace)
        is_low = confidence < threshold

        callback_results: List[LowConfidenceAction] = []

        if is_low and self.config.callbacks:
            context: Dict[str, Any] = {
                "namespace": namespace,
                "weights": weights,
                "vector_similarity": vector_similarity,
                "graph_path_score": graph_path_score,
                "trust_score": trust_score,
                "freshness": freshness,
                **(extra_context or {}),
            }

            for callback in self.config.callbacks:
                try:
                    result = callback(
                        query=query,
                        results=results,
                        confidence=confidence,
                        threshold=threshold,
                        context=context,
                    )
                    if result:
                        action = LowConfidenceAction(
                            action=result.get("action", "unknown"),
                            prompt=result.get("prompt"),
                            suggested_queries=result.get("suggested_queries", []),
                            metadata=result.get("metadata", {}),
                        )
                        callback_results.append(action)
                except Exception as e:
                    logger.warning(
                        "Low-confidence callback failed",
                        extra={"error": str(e), "query": query},
                    )

        return ConfidenceResult(
            confidence=confidence,
            threshold=threshold,
            is_low_confidence=is_low,
            callback_results=callback_results,
            components=ConfidenceComponents(
                vector_similarity=vector_similarity,
                graph_path_score=graph_path_score,
                trust_score=trust_score,
                freshness=freshness,
            ),
        )

    def register_callback(self, callback: LowConfidenceCallback) -> None:
        """Register a new low-confidence callback.

        Args:
            callback: Callback function to register
        """
        self.config.callbacks.append(callback)

    def set_threshold(
        self,
        threshold: float,
        namespace: Optional[str] = None,
    ) -> None:
        """Set confidence threshold.

        Args:
            threshold: New threshold value (0.0-1.0)
            namespace: If provided, sets namespace-specific threshold
        """
        if not 0.0 <= threshold <= 1.0:
            raise ValueError("Threshold must be between 0.0 and 1.0")

        if namespace:
            self.config.namespace_thresholds[namespace] = threshold
        else:
            self.config.threshold = threshold


@dataclass
class ConfidenceComponents:
    """Individual components that make up a confidence score."""

    vector_similarity: float
    graph_path_score: float
    trust_score: float
    freshness: float


@dataclass
class ConfidenceResult:
    """Result of confidence evaluation.

    Attributes:
        confidence: The computed confidence score (0.0-1.0)
        threshold: The threshold that was applied
        is_low_confidence: Whether confidence is below threshold
        callback_results: Results from any triggered callbacks
        components: Individual score components
    """

    confidence: float
    threshold: float
    is_low_confidence: bool
    callback_results: List[LowConfidenceAction]
    components: ConfidenceComponents

    @property
    def should_ask_followup(self) -> bool:
        """Check if any callback suggests asking for clarification."""
        return any(
            r.action in ("ask_clarification", "ask_followup")
            for r in self.callback_results
        )

    @property
    def suggested_prompt(self) -> Optional[str]:
        """Get first suggested prompt from callbacks."""
        for result in self.callback_results:
            if result.prompt:
                return result.prompt
        return None


def create_ask_clarification_callback(
    prompt_template: str = "Could you clarify what you mean by '{query}'?",
) -> Callable[..., Dict[str, Any]]:
    """Create a simple clarification callback.

    Args:
        prompt_template: Template for clarification prompt. Use {query} placeholder.

    Returns:
        Callback function that returns clarification action
    """

    def callback(
        query: str,
        results: List[Any],
        confidence: float,
        threshold: float,
        context: Dict[str, Any],
    ) -> Dict[str, Any]:
        return {
            "action": "ask_clarification",
            "prompt": prompt_template.format(query=query),
            "metadata": {
                "confidence": confidence,
                "threshold": threshold,
                "result_count": len(results),
            },
        }

    return callback


def create_expand_search_callback(
    expansion_factor: float = 1.5,
) -> Callable[..., Dict[str, Any]]:
    """Create a callback that suggests expanding the search.

    Args:
        expansion_factor: Factor to multiply k by for expanded search

    Returns:
        Callback function that returns expand_search action
    """

    def callback(
        query: str,
        results: List[Any],
        confidence: float,
        threshold: float,
        context: Dict[str, Any],
    ) -> Dict[str, Any]:
        return {
            "action": "expand_search",
            "metadata": {
                "expansion_factor": expansion_factor,
                "original_result_count": len(results),
            },
        }

    return callback
