"""Quality filter — keeps only high-value trajectories for LoRA training.

Criteria for a high-quality training example:
1. Tool calls succeeded (tool calls present, no error entries)
2. Multi-step tasks completed (≥2 messages + ≥1 tool call)
3. No low-confidence signals (confidence < 0.5 on any Timmy message)
4. Minimum meaningful exchange (≥1 user message + ≥1 Timmy message)

Refs: #1105
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from enum import StrEnum

from timmy_automations.retrain.trajectory_exporter import Trajectory

logger = logging.getLogger(__name__)

_MIN_CONFIDENCE = 0.5


class TrajectoryQuality(StrEnum):
    """Quality classification for a trajectory."""

    HIGH = "high"      # Multi-step + tool success — ideal training data
    MEDIUM = "medium"  # Single exchange, no errors — acceptable
    LOW = "low"        # Error-prone or trivial — skip


@dataclass
class QualityResult:
    """Result of quality assessment for a single trajectory."""

    trajectory: Trajectory
    quality: TrajectoryQuality
    score: float
    reasons: list[str]

    @property
    def is_trainable(self) -> bool:
        return self.quality in (TrajectoryQuality.HIGH, TrajectoryQuality.MEDIUM)


class QualityFilter:
    """Filters trajectories to keep only those worth training on.

    Scoring:
    - +1 pt:  base score for any valid clean exchange (no errors)
    - +3 pts: multi-step task (≥2 messages + ≥1 tool call)
    - +2 pts: tool calls present and no errors
    - +1 pt:  decision recorded (deliberate choice made)
    - -2 pts: any error entry
    - -1 pt:  any low-confidence response (confidence < 0.5)

    HIGH ≥ 4, MEDIUM 1–3, LOW ≤ 0
    """

    def __init__(self, min_confidence: float = _MIN_CONFIDENCE):
        self._min_confidence = min_confidence

    def assess(self, trajectory: Trajectory) -> QualityResult:
        """Score and classify a single trajectory."""
        score = 0.0
        reasons: list[str] = []

        # Minimum viable exchange check
        user_msgs = [m for m in trajectory.messages if m.get("role") == "user"]
        timmy_msgs = [m for m in trajectory.messages if m.get("role") == "timmy"]

        if not user_msgs or not timmy_msgs:
            return QualityResult(
                trajectory=trajectory,
                quality=TrajectoryQuality.LOW,
                score=0.0,
                reasons=["Missing user or assistant messages — not a valid exchange"],
            )

        # Multi-step bonus
        if trajectory.is_multi_step:
            score += 3.0
            reasons.append(
                f"Multi-step task: {trajectory.message_count} messages, "
                f"{trajectory.tool_call_count} tool calls"
            )

        # Base score for any clean exchange (user + timmy, no tool call required)
        if trajectory.error_count == 0:
            score += 1.0
            reasons.append("Clean exchange (no errors)")

        # Tool call quality
        if trajectory.tool_call_count > 0:
            if trajectory.error_count == 0:
                score += 2.0
                reasons.append(
                    f"All {trajectory.tool_call_count} tool call(s) succeeded"
                )
            else:
                score -= 2.0
                reasons.append(
                    f"{trajectory.error_count} error(s) during {trajectory.tool_call_count} tool call(s)"
                )
        elif trajectory.error_count > 0:
            score -= 2.0
            reasons.append(f"{trajectory.error_count} error(s) with no tool calls")

        # Decision bonus
        if trajectory.decisions:
            score += 1.0
            reasons.append(f"Decisions recorded: {len(trajectory.decisions)}")

        # Confidence penalty
        low_conf = [
            m
            for m in timmy_msgs
            if m.get("confidence") is not None
            and m["confidence"] < self._min_confidence
        ]
        if low_conf:
            score -= len(low_conf)
            reasons.append(
                f"{len(low_conf)} low-confidence response(s) (threshold={self._min_confidence})"
            )

        # Classify
        if score >= 4.0:
            quality = TrajectoryQuality.HIGH
        elif score >= 1.0:
            quality = TrajectoryQuality.MEDIUM
        else:
            quality = TrajectoryQuality.LOW

        return QualityResult(
            trajectory=trajectory,
            quality=quality,
            score=score,
            reasons=reasons,
        )

    def filter(
        self, trajectories: list[Trajectory]
    ) -> tuple[list[QualityResult], dict[str, int]]:
        """Assess all trajectories and return trainable ones with stats.

        Returns:
            (trainable_results, stats_dict) where stats_dict has keys
            'total', 'high', 'medium', 'low', 'accepted'.
        """
        results = [self.assess(t) for t in trajectories]
        trainable = [r for r in results if r.is_trainable]

        stats = {
            "total": len(results),
            "high": sum(1 for r in results if r.quality == TrajectoryQuality.HIGH),
            "medium": sum(1 for r in results if r.quality == TrajectoryQuality.MEDIUM),
            "low": sum(1 for r in results if r.quality == TrajectoryQuality.LOW),
            "accepted": len(trainable),
        }

        logger.info(
            "Quality filter: %d/%d accepted (high=%d medium=%d low=%d)",
            stats["accepted"],
            stats["total"],
            stats["high"],
            stats["medium"],
            stats["low"],
        )

        return trainable, stats
