"""
TrajectoryTracker — per-trace step tracking with predictive pattern matching.

Builds a running signature as spans complete. Checks against cached
trajectory patterns for early intervention triggers. Target: <1ms.
"""

import logging
from dataclasses import dataclass, field
from typing import Any

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


@dataclass
class TrajectoryPrediction:
    """A predicted failure from trajectory pattern matching."""

    pattern_id: str
    predicted_domain: str
    predicted_root_cause: str
    confidence: float
    recommended_strategy: str
    recommended_fix_content: str


@dataclass
class _TraceState:
    """Internal state for a trace's trajectory."""

    steps: list[dict[str, Any]] = field(default_factory=list)


class TrajectoryTracker:
    """Tracks per-trace trajectories and matches against known failure patterns.

    All operations are local dict lookups — no network calls. Target: <1ms.
    """

    def __init__(self, patterns: list[dict[str, Any]] | None = None):
        self._patterns: list[dict[str, Any]] = patterns or []
        self._traces: dict[str, _TraceState] = {}
        self._by_length: dict[int, list[dict[str, Any]]] = {}
        self._index_patterns()

    def record_step(
        self,
        trace_id: str,
        step_index: int,
        signals: list[str],
        span_type: str,
    ) -> TrajectoryPrediction | None:
        """Record a step and check for trajectory pattern match.

        Returns a prediction if the current trajectory matches a known
        failure pattern, None otherwise.
        """
        state = self._traces.setdefault(trace_id, _TraceState())
        state.steps.append(
            {
                "step": step_index,
                "signals": signals,
                "span_type": span_type,
            }
        )

        current_length = len(state.steps)
        candidates = self._by_length.get(current_length, [])

        for pattern in candidates:
            if self._matches(state.steps, pattern):
                failure = pattern.get("predicted_failure", {})
                return TrajectoryPrediction(
                    pattern_id=pattern["id"],
                    predicted_domain=failure.get("issue_domain", ""),
                    predicted_root_cause=failure.get("root_cause", ""),
                    confidence=pattern.get("confidence", 0.0),
                    recommended_strategy=pattern.get("recommended_strategy", ""),
                    recommended_fix_content=pattern.get("recommended_fix_content", ""),
                )
        return None

    def clear_trace(self, trace_id: str) -> None:
        """Clean up state for a completed trace."""
        self._traces.pop(trace_id, None)

    def update_patterns(self, patterns: list[dict[str, Any]]) -> None:
        """Replace patterns (called by PatternCache on sync)."""
        self._patterns = patterns
        self._index_patterns()

    def pattern_count(self) -> int:
        return len(self._patterns)

    def _index_patterns(self) -> None:
        """Index patterns by prefix_length for O(1) lookup."""
        self._by_length.clear()
        for p in self._patterns:
            length = p.get("prefix_length", 0)
            self._by_length.setdefault(length, []).append(p)

    def _matches(self, steps: list[dict], pattern: dict) -> bool:
        """Check if trace steps match a pattern's prefix_signature."""
        prefix = pattern.get("prefix_signature", [])
        if len(steps) < len(prefix):
            return False

        for i, expected in enumerate(prefix):
            actual = steps[i]
            if expected.get("span_type") and actual.get("span_type") != expected["span_type"]:
                return False
            expected_signals = set(expected.get("signals", []))
            actual_signals = set(actual.get("signals", []))
            if expected_signals and not expected_signals.intersection(actual_signals):
                return False
        return True
