"""
Judge Selector — auto-selects applicable judges per span type.

Ported from evals/worker/src/features/evaluation/judgeSelector.ts.
Maps span types to judge categories, enabling automatic evaluation
without requiring explicit template selection.

Enhanced with signal-based selection: FastDetector produces lightweight
signal hints (e.g., "refusal", "hallucination_marker") that help the
selector pick the RIGHT platform judges for each step.
"""

import logging
from dataclasses import dataclass, field
from typing import Any, Dict, List, Set

logger = logging.getLogger("aigie.judge.selector")


@dataclass
class JudgeCategory:
    """A category of judge evaluation."""

    name: str
    description: str
    template_ids: List[str] = field(default_factory=list)  # Platform template IDs
    weight: float = 1.0  # Relative importance


@dataclass
class PlatformJudge:
    """A predefined judge synced from the Kytte platform."""

    judge_id: str
    name: str
    category: str
    template_id: str | None = None
    description: str = ""
    # Which signal hints this judge is designed to handle
    signal_triggers: List[str] = field(default_factory=list)
    # Which span types this judge applies to (empty = all)
    span_types: List[str] = field(default_factory=list)
    # Skills this judge leverages
    skills: List[str] = field(default_factory=list)
    priority: int = 0  # Higher = higher priority


# Default judge categories mapped by span type
_SPAN_TYPE_JUDGES: Dict[str, List[str]] = {
    "llm": ["accuracy", "quality", "relevance", "safety-monitor"],
    "tool": ["tool-usage", "error-handling", "efficiency"],
    "agent": [
        "accuracy",
        "quality",
        "relevance",
        "tool-usage",
        "error-handling",
        "efficiency",
        "safety-monitor",
    ],
    "chain": ["quality", "relevance", "efficiency"],
    "retrieval": ["relevance", "accuracy"],
    "embedding": ["efficiency"],
}

# Judge categories that can run as fast heuristics (no LLM needed)
_HEURISTIC_JUDGES: Set[str] = {
    "error-handling",
    "tool-usage",
    "efficiency",
    "safety-monitor",
}

# Signal hints → which judge categories they suggest
# This is the default mapping. Platform judges can override via signal_triggers.
_SIGNAL_TO_JUDGES: Dict[str, List[str]] = {
    "refusal": ["quality", "safety-monitor"],
    "safety_triggered": ["safety-monitor", "quality"],
    "hallucination_marker": ["accuracy", "quality"],
    "lazy_response": ["quality", "relevance"],
    "suspiciously_short": ["quality", "relevance"],
}


class JudgeSelector:
    """Selects applicable judges for a given span.

    Three selection modes:
    1. select(span_type) — basic selection by span type (backward compatible)
    2. select_with_signals(span_type, signals, metadata) — signal-aware selection
       that uses FastDetector hints to pick the right platform judges
    3. select_heuristic_only / select_llm_only — filtered subsets

    Usage:
        selector = JudgeSelector()

        # Basic (backward compatible)
        judges = selector.select("llm")

        # Signal-aware (new)
        judges = selector.select_with_signals(
            "llm",
            signal_hints=["refusal", "hallucination_marker"],
            metadata={"model": "gpt-4"},
        )
    """

    def __init__(
        self,
        custom_mappings: Dict[str, List[str]] | None = None,
        template_id_map: Dict[str, str] | None = None,
    ):
        self._mappings = dict(_SPAN_TYPE_JUDGES)
        if custom_mappings:
            self._mappings.update(custom_mappings)
        # Maps judge category name → platform eval template ID
        self._template_map = template_id_map or {}
        # Platform judges synced from Kytte (judge_id → PlatformJudge)
        self._platform_judges: Dict[str, PlatformJudge] = {}
        # Signal → judge mapping (default + platform overrides)
        self._signal_judge_map: Dict[str, List[str]] = dict(_SIGNAL_TO_JUDGES)

    def select(
        self,
        span_type: str,
        metadata: Dict[str, Any] | None = None,
    ) -> List[str]:
        """Select applicable judge categories for a span type."""
        judges = list(self._mappings.get(span_type, self._mappings.get("llm", [])))

        # Add error-handling judge if span has error indicators
        if metadata:
            has_error = (
                metadata.get("status") == "error"
                or metadata.get("error") is not None
                or metadata.get("status_code", 200) >= 400
            )
            if has_error and "error-handling" not in judges:
                judges.insert(0, "error-handling")

        return judges

    def select_with_signals(
        self,
        span_type: str,
        signal_hints: List[str] | None = None,
        metadata: Dict[str, Any] | None = None,
    ) -> List[str]:
        """Select judges using FastDetector signal hints for smarter targeting.

        Signal hints (e.g., "refusal", "hallucination_marker") help prioritize
        the right platform judges. A "hallucination_marker" signal will boost
        accuracy/quality judges to the top of the list.

        Platform judges with matching skills get higher priority — a judge with
        skills=["hallucination_detection"] will be ranked above a generic
        accuracy judge when "hallucination_marker" signal is present.

        Args:
            span_type: Type of span (llm, tool, agent, etc.)
            signal_hints: Lightweight tags from FastDetector
            metadata: Span metadata (model, status, error, etc.)

        Returns:
            Ordered list of judge categories, with signal-relevant judges first.
        """
        # Start with base span-type selection
        base_judges = self.select(span_type, metadata)

        if not signal_hints:
            return base_judges

        # Collect signal-triggered judges with priority scoring
        # Higher score = more relevant (signal match + skills match + priority)
        category_scores: Dict[str, float] = {}

        for signal in signal_hints:
            # Check platform judges first (they have explicit signal_triggers + skills)
            for pj in self._platform_judges.values():
                if signal in pj.signal_triggers:
                    # Filter by span_type if the judge restricts it
                    if pj.span_types and span_type not in pj.span_types:
                        continue
                    score = category_scores.get(pj.category, 0.0)
                    score += 10.0  # Base signal match
                    score += pj.priority  # Platform-assigned priority
                    # Skill bonus: judges with relevant skills get boosted
                    if pj.skills:
                        skill_match = any(
                            signal.replace("_", "") in s.replace("_", "").replace("-", "")
                            for s in pj.skills
                        )
                        if skill_match:
                            score += 5.0  # Skill relevance bonus
                    category_scores[pj.category] = max(score, category_scores.get(pj.category, 0.0))

            # Fall back to default signal → judge mapping
            for category in self._signal_judge_map.get(signal, []):
                if category not in category_scores:
                    category_scores[category] = 1.0  # Default mapping score

        # Sort signal-triggered categories by score (highest first)
        signal_judges = sorted(
            category_scores.keys(), key=lambda c: category_scores[c], reverse=True
        )

        # Merge: signal-triggered judges first, then remaining base judges
        result: List[str] = list(signal_judges)
        for j in base_judges:
            if j not in result:
                result.append(j)

        return result

    def select_heuristic_only(self, span_type: str) -> List[str]:
        """Select only judges that can run as fast heuristics."""
        all_judges = self.select(span_type)
        return [j for j in all_judges if j in _HEURISTIC_JUDGES]

    def select_llm_only(self, span_type: str) -> List[str]:
        """Select only judges that require LLM evaluation."""
        all_judges = self.select(span_type)
        return [j for j in all_judges if j not in _HEURISTIC_JUDGES]

    def get_template_id(self, judge_category: str) -> str | None:
        """Get platform eval template ID for a judge category."""
        return self._template_map.get(judge_category)

    def update_template_map(self, mapping: Dict[str, str]) -> None:
        """Update the judge category → template ID mapping."""
        self._template_map.update(mapping)

    # ------------------------------------------------------------------
    # Platform judge management
    # ------------------------------------------------------------------

    def update_platform_judges(self, judges: List[Dict[str, Any]]) -> None:
        """Sync predefined judges from the Kytte platform.

        Called by PatternCache during refresh. Each judge dict should contain:
        {judge_id, name, category, template_id, signal_triggers, span_types, skills, priority}
        """
        new_judges: Dict[str, PlatformJudge] = {}
        for j in judges:
            pj = PlatformJudge(
                judge_id=j.get("judge_id", j.get("id", "")),
                name=j.get("name", ""),
                category=j.get("category", "quality"),
                template_id=j.get("template_id"),
                description=j.get("description", ""),
                signal_triggers=j.get("signal_triggers", []),
                span_types=j.get("span_types", []),
                skills=j.get("skills", []),
                priority=j.get("priority", 0),
            )
            new_judges[pj.judge_id] = pj

            # Auto-update template map
            if pj.template_id:
                self._template_map[pj.category] = pj.template_id

        self._platform_judges = new_judges
        logger.debug(f"Updated {len(new_judges)} platform judges")

    def update_signal_map(self, mapping: Dict[str, List[str]]) -> None:
        """Update signal → judge category mapping from platform."""
        self._signal_judge_map.update(mapping)

    def get_platform_judges_for_signal(self, signal: str) -> List[PlatformJudge]:
        """Get platform judges that handle a specific signal."""
        return [pj for pj in self._platform_judges.values() if signal in pj.signal_triggers]

    def get_stats(self) -> Dict[str, Any]:
        """Get selector statistics."""
        return {
            "span_type_mappings": len(self._mappings),
            "platform_judges": len(self._platform_judges),
            "template_mappings": len(self._template_map),
            "signal_mappings": len(self._signal_judge_map),
        }
