"""Dataset comparison engine — side-by-side quality diff.

Runs orbit analyze on two datasets and produces a structured diff
showing what improved, what degraded, and what's new.
"""

from __future__ import annotations

import json
import subprocess
from dataclasses import dataclass, field


@dataclass
class MetricDelta:
    """A single metric comparison."""

    name: str
    value_a: float | str | None
    value_b: float | str | None
    delta: float | None  # numeric diff, or None for non-numeric
    direction: str  # "improved", "degraded", "unchanged", "new", "removed"
    display: str  # human-readable diff string


@dataclass
class CompareResult:
    """Full comparison between two datasets."""

    dataset_a: str
    dataset_b: str
    grade_a: str
    grade_b: str
    score_a: int
    score_b: int
    score_delta: int
    verdict: str  # one-line summary
    metrics: list[MetricDelta]
    new_strengths: list[str]
    resolved_issues: list[str]
    remaining_issues: list[str]
    new_issues: list[str]


def _analyze_dataset(repo_id: str, policy: str = "auto") -> dict | None:
    """Run orbit analyze --json on a dataset and return parsed results."""
    try:
        result = subprocess.run(
            [
                "orbit", "analyze", repo_id,
                "--skip-embeddings",
                "--policy", policy,
                "--json",
            ],
            capture_output=True,
            text=True,
            timeout=180,
        )
        if result.returncode != 0:
            return None
        return json.loads(result.stdout)
    except (subprocess.TimeoutExpired, json.JSONDecodeError):
        return None


def _direction(a: float | None, b: float | None, higher_is_better: bool = True) -> str:
    if a is None or b is None:
        return "unchanged"
    diff = b - a
    if abs(diff) < 0.005:
        return "unchanged"
    if higher_is_better:
        return "improved" if diff > 0 else "degraded"
    else:
        return "improved" if diff < 0 else "degraded"


def _arrow(direction: str) -> str:
    return {
        "improved": "↑",
        "degraded": "↓",
        "unchanged": "→",
        "new": "+",
        "removed": "−",
    }.get(direction, "?")


def compare_datasets(
    dataset_a: str,
    dataset_b: str,
    policy: str = "auto",
) -> CompareResult:
    """Compare two datasets by running orbit analyze on both.

    Args:
        dataset_a: First dataset (baseline/before).
        dataset_b: Second dataset (comparison/after).
        policy: Policy type for analysis.

    Returns:
        CompareResult with structured diff.
    """
    report_a = _analyze_dataset(dataset_a, policy)
    report_b = _analyze_dataset(dataset_b, policy)

    if report_a is None or report_b is None:
        failed = dataset_a if report_a is None else dataset_b
        return CompareResult(
            dataset_a=dataset_a,
            dataset_b=dataset_b,
            grade_a="?",
            grade_b="?",
            score_a=0,
            score_b=0,
            score_delta=0,
            verdict=f"Could not analyze {failed}",
            metrics=[],
            new_strengths=[],
            resolved_issues=[],
            remaining_issues=[],
            new_issues=[],
        )

    ra = report_a.get("readiness", {})
    rb = report_b.get("readiness", {})
    grade_a = ra.get("grade", "?")
    grade_b = rb.get("grade", "?")
    score_a = ra.get("score", 0)
    score_b = rb.get("score", 0)
    score_delta = score_b - score_a

    # Build metric comparisons
    metrics = []

    def add_metric(name: str, key_path: list[str], higher_is_better: bool = True, fmt: str = ".2f"):
        va = report_a
        vb = report_b
        for k in key_path:
            va = va.get(k) if isinstance(va, dict) else None
            vb = vb.get(k) if isinstance(vb, dict) else None
        if va is None and vb is None:
            return
        d = _direction(va, vb, higher_is_better)
        delta = None
        if isinstance(va, (int, float)) and isinstance(vb, (int, float)):
            delta = vb - va
            display = f"{va:{fmt}} → {vb:{fmt}} ({_arrow(d)} {delta:+{fmt}})"
        else:
            display = f"{va} → {vb}"
        metrics.append(MetricDelta(name=name, value_a=va, value_b=vb, delta=delta, direction=d, display=display))

    add_metric("Score", ["readiness", "score"], fmt="d")
    add_metric("Episodes", ["num_episodes"], fmt="d")
    add_metric("Quality Score", ["quality_score"])
    add_metric("Action Divergence", ["action_divergence", "score"], higher_is_better=False)
    add_metric("Policy Fit", ["policy_fit", "fit_score"])
    add_metric("Consistency", ["consistency_diagnostics", "overall_consistency"])
    add_metric("Episode Consistency", ["coverage", "episode_consistency"])

    # Strengths/problems diff
    strengths_a = set(ra.get("strengths", []))
    strengths_b = set(rb.get("strengths", []))
    problems_a = set(ra.get("problems", []))
    problems_b = set(rb.get("problems", []))

    new_strengths = list(strengths_b - strengths_a)
    resolved_issues = list(problems_a - problems_b)
    remaining_issues = list(problems_a & problems_b)
    new_issues = list(problems_b - problems_a)

    # Verdict
    if score_delta > 10:
        verdict = f"{dataset_b} is significantly better (+{score_delta} points). "
    elif score_delta > 0:
        verdict = f"{dataset_b} is slightly better (+{score_delta} points). "
    elif score_delta == 0:
        verdict = "Both datasets are equivalent. "
    elif score_delta > -10:
        verdict = f"{dataset_b} is slightly worse ({score_delta} points). "
    else:
        verdict = f"{dataset_b} is significantly worse ({score_delta} points). "

    grade_order = {"A": 4, "B": 3, "C": 2, "D": 1, "F": 0}
    if grade_order.get(grade_b, 0) >= 3:
        verdict += "Ready for training."
    elif grade_order.get(grade_b, 0) >= 2:
        verdict += "Consider cleaning before training."
    else:
        verdict += "Needs more work before training."

    return CompareResult(
        dataset_a=dataset_a,
        dataset_b=dataset_b,
        grade_a=grade_a,
        grade_b=grade_b,
        score_a=score_a,
        score_b=score_b,
        score_delta=score_delta,
        verdict=verdict,
        metrics=metrics,
        new_strengths=new_strengths,
        resolved_issues=resolved_issues,
        remaining_issues=remaining_issues,
        new_issues=new_issues,
    )


def print_compare(result: CompareResult) -> None:
    """Print a formatted comparison report."""
    from rich.panel import Panel

    from orbit.utils.display import console

    # Header
    console.print()
    console.print(Panel(
        f"[bold]ORBIT Compare[/bold]\n"
        f"{result.dataset_a}  →  {result.dataset_b}",
        border_style="cyan",
    ))

    # Grade change
    ga_color = {"A": "green", "B": "green", "C": "yellow", "D": "yellow", "F": "red"}.get(result.grade_a, "white")
    gb_color = {"A": "green", "B": "green", "C": "yellow", "D": "yellow", "F": "red"}.get(result.grade_b, "white")
    delta_color = "green" if result.score_delta > 0 else ("red" if result.score_delta < 0 else "dim")
    delta_str = f"{result.score_delta:+d}" if result.score_delta != 0 else "no change"

    console.print(
        f"\n  Grade: [{ga_color}]{result.grade_a}[/{ga_color}] ({result.score_a})"
        f"  →  [{gb_color}]{result.grade_b}[/{gb_color}] ({result.score_b})"
        f"  [{delta_color}]({delta_str})[/{delta_color}]"
    )

    # Metrics table
    if result.metrics:
        console.print()
        console.print(f"  {'Metric':<24} {'Before':<12} {'After':<12} {'Change'}")
        console.print(f"  {'─' * 60}")
        for m in result.metrics:
            color = {"improved": "green", "degraded": "red", "unchanged": "dim"}.get(m.direction, "white")
            console.print(f"  {m.name:<24} [{color}]{m.display}[/{color}]")

    # Resolved issues
    if result.resolved_issues:
        console.print(f"\n  [bold green]RESOLVED ISSUES[/bold green]")
        for issue in result.resolved_issues:
            console.print(f"    [green]✓[/green] {issue}")

    # New strengths
    if result.new_strengths:
        console.print(f"\n  [bold green]NEW STRENGTHS[/bold green]")
        for s in result.new_strengths:
            console.print(f"    [green]✓[/green] {s}")

    # Remaining issues
    if result.remaining_issues:
        console.print(f"\n  [yellow]REMAINING ISSUES[/yellow]")
        for issue in result.remaining_issues:
            console.print(f"    [yellow]⚠[/yellow] {issue}")

    # New issues
    if result.new_issues:
        console.print(f"\n  [red]NEW ISSUES[/red]")
        for issue in result.new_issues:
            console.print(f"    [red]✗[/red] {issue}")

    # Verdict
    console.print(f"\n  [bold]Verdict:[/bold] {result.verdict}")
    console.print()
