"""bandito judge — calibrate, run, overview, runs."""

from __future__ import annotations

import logging
import os
from typing import Any

from . import get_client
from .observe import format_ago
from ..storage import get_store

logger = logging.getLogger(__name__)


def run_overview(project: str, *, json_output: bool = False) -> None:
    import json as json_mod

    client = get_client()
    if client is None:
        return

    # Fetch calibration status
    cal_status = client.get_calibration(project)
    cal = cal_status["calibration"] if cal_status and cal_status.get("calibration") else None

    # Fetch calibration traces (disagreements)
    cal_traces = []
    try:
        cal_traces = client.list_calibration_traces(project)
    except RuntimeError:
        pass

    # Fetch judge runs for usage summary
    runs = []
    try:
        runs = client.list_judge_runs(project)
    except RuntimeError:
        pass

    if json_output:
        disagreements = [t for t in cal_traces if not t.get("agreed", True)]
        output = {
            "calibration": None,
            "usage": {
                "total_runs": len(runs),
                "total_scored": sum(r.get("total_scored", 0) for r in runs),
                "avg_pass_rate": round(
                    sum(r.get("passed", 0) for r in runs) / max(sum(r.get("total_scored", 0) for r in runs), 1), 4
                ) if runs else None,
                "latest_run": runs[0] if runs else None,
            },
        }
        if cal:
            status = "not calibrated"
            if cal_status.get("stale"):
                status = "stale"
            elif cal.get("tpr", 0) >= 0.9 and cal.get("tnr", 0) >= 0.9:
                status = "well-calibrated"
            else:
                status = "needs-work"
            output["calibration"] = {
                "status": status,
                "stale": cal_status.get("stale", False),
                "tpr": cal.get("tpr"),
                "tnr": cal.get("tnr"),
                "agreement": cal.get("agreement"),
                "kappa": cal.get("kappa"),
                "total": cal.get("total"),
                "tp": cal.get("tp"),
                "tn": cal.get("tn"),
                "fp": cal.get("fp"),
                "fn": cal.get("fn"),
                "calibrated_at": cal.get("calibrated_at"),
                "grades_since_calibration": cal_status.get("grades_since_calibration", 0),
                "disagreements": disagreements,
            }
        print(json_mod.dumps(output, indent=2, default=str))
        return

    # Human-readable output
    print(f"\n  Judge Overview: {project}\n")

    # -- Calibration section --
    print(f"  \u2500\u2500 Calibration {'\u2500' * 39}")
    print()
    if cal is None:
        print("  Not yet calibrated.")
        print(f"  Run:  bandito judge calibrate --project {project}")
    else:
        tpr = cal.get("tpr", 0)
        tnr = cal.get("tnr", 0)

        if cal_status.get("stale"):
            status_str = "stale \u2014 rubric edited after last calibration"
        elif tpr >= 0.9 and tnr >= 0.9:
            status_str = "well-calibrated"
        else:
            weak = "TPR" if tpr < tnr else "TNR"
            status_str = f"{weak} below 90% \u2014 refine rubric and re-calibrate"

        print(f"  Status:        {status_str}")
        tpr_mark = "\u2713" if tpr >= 0.9 else "\u2717"
        tnr_mark = "\u2713" if tnr >= 0.9 else "\u2717"
        print(f"  {tpr_mark} TPR (recall):  {tpr:.0%}   (target: \u226590%)")
        print(f"  {tnr_mark} TNR (spec):    {tnr:.0%}   (target: \u226590%)")
        print(f"    Agreement:   {cal.get('agreement', 0):.0%}")
        print(f"    Kappa:       {cal.get('kappa', 0):.2f}")
        print(f"    Traces:      {cal.get('total', 0)}")

        cal_at = cal.get("calibrated_at", "")
        if cal_at:
            print(f"    Calibrated:  {format_ago(cal_at)}")

        grades_since = cal_status.get("grades_since_calibration", 0)
        if grades_since:
            print(f"    New grades:  {grades_since} since last calibration")

        # Disagreements
        disagreements = [t for t in cal_traces if not t.get("agreed", True)]
        if disagreements:
            print()
            print(f"  Disagreements ({len(disagreements)}):")
            for d in disagreements:
                kind = d.get("kind", "?")
                critique = d.get("judge_critique", "")
                if len(critique) > 80:
                    critique = critique[:80] + "\u2026"
                print(f"    {d.get('trace_id', '?')}  {kind}  \"{critique}\"")

    # -- Usage section --
    print()
    print(f"  \u2500\u2500 Usage {'\u2500' * 45}")
    print()
    if not runs:
        print("  No judge runs yet.")
        if cal and cal.get("tpr", 0) >= 0.9 and cal.get("tnr", 0) >= 0.9:
            print(f"  Run:  bandito judge run --project {project}")
    else:
        total_scored = sum(r.get("total_scored", 0) for r in runs)
        total_passed = sum(r.get("passed", 0) for r in runs)
        avg_rate = total_passed / total_scored if total_scored > 0 else 0.0
        latest = runs[0]
        latest_rate = latest.get("pass_rate", 0)
        latest_at = latest.get("created_at", "")

        print(f"  Total runs:       {len(runs)}")
        print(f"  Total scored:     {total_scored}")
        print(f"  Avg pass rate:    {avg_rate:.0%}")

        delta = latest_rate - avg_rate
        if delta > 0.005:
            delta_str = f" (+{delta:.0%} vs avg)"
        elif delta < -0.005:
            delta_str = f" ({delta:.0%} vs avg)"
        else:
            delta_str = ""

        latest_scored = latest.get("total_scored", 0)
        ago = format_ago(latest_at) if latest_at else ""
        print(f"  Latest run:       {latest_scored} traces \u00b7 {latest_rate:.0%}{delta_str}" + (f" \u00b7 {ago}" if ago else ""))

    print()


def run_runs(project: str, *, last: int | None = None, json_output: bool = False) -> None:
    import json as json_mod

    client = get_client()
    if client is None:
        return

    runs = []
    try:
        runs = client.list_judge_runs(project)
    except RuntimeError as e:
        print(f"\n  Error: {e}\n")
        return

    if last is not None:
        runs = runs[:last]

    if json_output:
        total_scored = sum(r.get("total_scored", 0) for r in runs)
        total_passed = sum(r.get("passed", 0) for r in runs)
        avg_rate = total_passed / total_scored if total_scored > 0 else 0.0
        print(json_mod.dumps({
            "runs": runs,
            "total_runs": len(runs),
            "total_scored": total_scored,
            "avg_pass_rate": round(avg_rate, 4),
        }, indent=2, default=str))
        return

    if not runs:
        print(f"\n  No judge runs for '{project}'.")
        print(f"  Run:  bandito judge run --project {project}\n")
        return

    print(f"\n  Judge Runs: {project}\n")
    _print_runs_table(runs)


def run_calibrate(
    project: str,
    *,
    max_traces: int | None = None,
    dry_run: bool = False,
    regenerate: bool = False,
) -> None:
    from .judge.calibrate import format_calibration_report, run_calibration
    from .judge.rubric import draft_rubric, parse_rubric_content

    # --- Auth ---
    client = get_client()
    if client is None:
        return

    # --- Load store for traces ---
    try:
        store = get_store(project=project)
    except Exception as e:
        print(f"\n  Could not load traces: {e}\n")
        return

    # --- Rubric: fetch or create inline ---
    rubric_data = client.get_rubric(project)

    if regenerate and rubric_data is not None:
        print(f"\n  Regenerating rubric for '{project}'.\n")
        rubric_data = None

    if rubric_data is None:
        # Inline rubric creation
        print(f"\n  No rubric found for '{project}'. Creating one.\n")

        traces_sample = store.list_traces_filtered(project=project, limit=10)
        if not traces_sample:
            print("  No traces found — cannot draft rubric.")
            print(f"  Instrument your app first, then re-run.\n")
            return

        print(f"  Drafting rubric from {len(traces_sample)} traces...")
        try:
            content = draft_rubric(project, traces_sample, client)
        except RuntimeError as e:
            print(f"\n  Rubric generation failed: {e}\n")
            return

        print(f"\n  Generated rubric ({len(content.splitlines())} lines).\n")
        answer = input("  Edit in $EDITOR before saving? [y/N] ").strip().lower()
        if answer in ("y", "yes"):
            content = _edit_in_editor(content)

        client.save_rubric(project, content)
        print(f"\n  Rubric saved for '{project}'.")
        rubric_content = content
    else:
        rubric_content = rubric_data["content"]

    # --- Show rubric and confirm before calibrating ---
    rubric = parse_rubric_content(rubric_content)

    print(f"\n  \u2500\u2500 Rubric: {project} {'\u2500' * max(0, 44 - len(project))}\n")
    for line in rubric_content.splitlines():
        print(f"  {line}")
    print(f"\n  {'\u2500' * 56}")

    warnings = rubric.validate()
    if warnings:
        for w in warnings:
            print(f"  Warning: {w}")
        print()

    answer = input("  Edit rubric in $EDITOR? [y/N] ").strip().lower()
    if answer in ("y", "yes"):
        rubric_content = _edit_in_editor(rubric_content)
        client.save_rubric(project, rubric_content)
        rubric = parse_rubric_content(rubric_content)
        print("  Rubric updated.\n")

    answer = input("\n  Run calibration? [Y/n] ").strip().lower()
    if answer not in ("", "y", "yes"):
        print(f"\n  Rubric saved. Run calibration later with:")
        print(f"    bandito judge calibrate --project {project}\n")
        return

    # --- Build traces_with_grades from backend evals + local traces ---
    from .retrieval import get_retrieval

    r = get_retrieval(project)
    evals = r.evals()
    traces_with_grades: list[tuple[dict, int]] = []
    for trace_id, ev in evals.items():
        if ev.get("human_score") is None:
            continue
        trace = r.get_trace(trace_id)
        if not trace or trace.get("project") != project:
            continue
        traces_with_grades.append((trace, ev["human_score"]))

    if not traces_with_grades:
        print(f"  No human-graded traces found for '{project}'.")
        print(f"  Grade some first:  bandito tui\n")
        return

    if max_traces and max_traces < len(traces_with_grades):
        traces_with_grades = traces_with_grades[:max_traces]

    # --- Run calibration ---
    print(f"  Calibrating judge on {len(traces_with_grades)} human-graded traces...\n")

    try:
        result = run_calibration(
            project=project,
            traces_with_grades=traces_with_grades,
            client=client,
        )
    except RuntimeError as e:
        print(f"\n  Calibration failed: {e}")
        print(f"  Rubric has been saved but is uncalibrated.")
        print(f"  Run 'bandito judge calibrate --project {project}' to retry.\n")
        return

    print(format_calibration_report(result, traces_with_grades))

    print(f"  Calibration saved.")
    if result["tpr"] >= 0.9 and result["tnr"] >= 0.9:
        print(f"  Ready to score:  bandito judge run --project {project}")
    else:
        print("  Refine your rubric and re-calibrate.")
    print()


def run_judge_cmd(
    project: str,
    *,
    sample: int | None = None,
    rejudge: bool = False,
    dry_run: bool = False,
    dataset: str | None = None,
) -> None:
    import random

    from .judge.run import format_run_summary, run_judge
    from ..storage import get_store

    # --- Auth ---
    client = get_client()
    if client is None:
        return

    # --- Resolve dataset if specified ---
    dataset_trace_ids = None
    if dataset:
        datasets = client.list_datasets(project=project)
        ds = None
        for d in datasets:
            if d["name"] == dataset or str(d["id"]) == dataset:
                ds = d
                break
        if ds is None:
            print(f"\n  Dataset '{dataset}' not found for project '{project}'\n")
            return
        full = client.get_dataset(ds["id"])
        dataset_trace_ids = {item["trace_id"] for item in full.get("items", [])}
        print(f"\n  Running judge on dataset '{ds['name']}' ({len(dataset_trace_ids)} traces)")

    # --- Check rubric exists ---
    rubric_data = client.get_rubric(project)
    if rubric_data is None:
        print(f"\n  No rubric found for '{project}'.")
        print(f"  Run:  bandito judge calibrate --project {project}\n")
        return

    # --- Hard calibration gate ---
    cal_status = client.get_calibration(project)
    if cal_status is None or not cal_status.get("calibrated"):
        print(f"\n  Judge has not been calibrated for '{project}'.")
        answer = input("  Run calibrate first? [Y/n] ").strip().lower()
        if answer in ("", "y", "yes"):
            run_calibrate(project)
            cal_status = client.get_calibration(project)
            if cal_status is None or not cal_status.get("calibrated"):
                print("  Calibration did not pass. Cannot run judge.\n")
                return
        else:
            return
    elif cal_status.get("stale"):
        print(f"\n  Calibration is stale \u2014 rubric was edited after last calibration.")
        answer = input("  Re-calibrate first? [Y/n] ").strip().lower()
        if answer in ("", "y", "yes"):
            run_calibrate(project)
            cal_status = client.get_calibration(project)
            if cal_status is None or not cal_status.get("calibrated"):
                print("  Calibration did not pass. Cannot run judge.\n")
                return
        else:
            return

    # --- Grade staleness nudge ---
    grades_since = cal_status.get("grades_since_calibration", 0)
    if grades_since >= 100:
        print(f"\n  {grades_since} new human grades since last calibration.")
        answer = input("  Re-calibrate to incorporate them? [Y/n] ").strip().lower()
        if answer in ("", "y", "yes"):
            run_calibrate(project)
            cal_status = client.get_calibration(project)
            if cal_status is None or not cal_status.get("calibrated"):
                print("  Calibration did not pass. Cannot run judge.\n")
                return

    # --- Load traces ---
    from .retrieval import get_retrieval

    try:
        r = get_retrieval(project)
    except Exception as e:
        print(f"\n  Could not load traces: {e}\n")
        return

    # Filter to dataset traces if specified
    if dataset_trace_ids:
        all_loaded = []
        for trace in r.traces(tier="trace", limit=10000):
            if trace.get("trace_id") in dataset_trace_ids:
                all_loaded.append(trace)
        print(f"  Loaded {len(all_loaded)} traces from dataset")

    # Build trace pool (cap at 5000 for safety)
    _TRACE_POOL_LIMIT = 5000
    evals = r.evals()

    invalid_scores = 0
    for tid, ev in evals.items():
        score = ev.get("judge_score")
        if score is not None and score not in (0, 1):
            invalid_scores += 1
            logger.warning(f"Invalid judge_score {score} for {tid}, skipping")
    if invalid_scores:
        logger.warning(f"Found {invalid_scores} traces with invalid judge scores")

    all_loaded = r.traces(tier="trace", limit=_TRACE_POOL_LIMIT)
    if len(all_loaded) >= _TRACE_POOL_LIMIT:
        print(f"\n  Note: loaded {_TRACE_POOL_LIMIT} most recent traces (project may have more).")
    if rejudge:
        traces = all_loaded
    else:
        judged_ids = {
            tid for tid, ev in evals.items()
            if ev.get("judge_score") is not None
        }
        traces = [
            t for t in all_loaded
            if t["trace_id"] not in judged_ids
        ]

    if not traces:
        if rejudge:
            print(f"\n  No traces found for '{project}'.\n")
        else:
            print(f"\n  No unjudged traces for '{project}'.")
            print("  Use --rejudge to re-score all traces.\n")
        return

    if sample and sample < len(traces):
        traces = random.sample(traces, sample)

    _run_judge_parallel(traces, client, project, rejudge)

    # Backend stores evals + judge run summary during execution

    # Post-run: grading coverage callout
    all_summaries = r.traces(tier="summary")
    total_traces = len(all_summaries)
    human_graded = sum(
        1 for ev in evals.values()
        if ev.get("human_score") is not None
    )

    if total_traces > 0:
        pct = human_graded / total_traces * 100
        print(f"  {total_traces} total traces \u00b7 {human_graded} human-graded ({pct:.0f}%)")
        if pct < 15:
            print()
            print("  Low human grading coverage. For better judge accuracy:")
            print(f"    1. Grade in TUI:   bandito tui --project {project}")
            print(f"    2. Re-calibrate:   bandito judge calibrate --project {project}")
        print()


# -- Helpers -------------------------------------------------------------------


def _edit_in_editor(content: str) -> str:
    """Open content in $EDITOR and return the edited result."""
    import subprocess
    import tempfile

    editor = os.environ.get("EDITOR", "vi")
    with tempfile.NamedTemporaryFile(suffix=".md", mode="w", delete=False) as tmp:
        tmp.write(content)
        tmp_path = tmp.name
    try:
        subprocess.run([editor, tmp_path], check=True)
        with open(tmp_path) as f:
            return f.read()
    except (OSError, subprocess.CalledProcessError) as e:
        print(f"  Editor failed: {e}. Using previous content.")
        return content
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass


def _print_runs_table(runs: list[dict]) -> None:
    """Print a formatted table of judge runs with vs-average column."""
    from datetime import datetime

    total_scored = sum(r.get("total_scored", 0) for r in runs)
    total_passed = sum(r.get("passed", 0) for r in runs)
    avg_rate = total_passed / total_scored if total_scored > 0 else 0.0

    print(f"  {'DATE':<22} {'SCORED':>7} {'PASSED':>7} {'FAILED':>7} {'RATE':>7} {'vs AVG':>7}  FLAGS")

    for r in runs:
        created = r.get("created_at", "")
        if created:
            try:
                dt = datetime.fromisoformat(str(created).replace("Z", "+00:00"))
                date_str = dt.strftime("%Y-%m-%d %H:%M")
            except (ValueError, TypeError):
                date_str = str(created)[:16]
        else:
            date_str = "?"

        scored = r.get("total_scored", 0)
        passed = r.get("passed", 0)
        failed = r.get("failed", 0)
        rate = r.get("pass_rate", 0)

        delta = rate - avg_rate
        if delta > 0.005:
            delta_str = f"+{delta:.0%}"
        elif delta < -0.005:
            delta_str = f"{delta:.0%}"
        else:
            delta_str = "="

        flags = []
        if r.get("sample_size"):
            flags.append(f"sample={r['sample_size']}")
        if r.get("rejudge"):
            flags.append("rejudge")
        if r.get("parse_errors", 0) > 0:
            flags.append(f"{r['parse_errors']} errors")
        flag_str = "  " + " ".join(flags) if flags else ""

        print(f"  {date_str:<22} {scored:>7} {passed:>7} {failed:>7} {rate:>6.0%} {delta_str:>7}{flag_str}")

    print(f"\n  {len(runs)} runs \u00b7 {total_scored} total scored \u00b7 avg pass rate: {avg_rate:.0%}")
    print()


def _run_judge_parallel(
    traces: list[dict[str, Any]],
    client: Any,
    project: str,
    rejudge: bool,
) -> None:
    """Run judge with automatic parallelization based on trace count."""
    import concurrent.futures
    from math import ceil

    from .judge.run import format_run_summary, run_judge

    total = len(traces)
    workers = min(4, max(1, ceil(total / 20)))
    chunk_size = ceil(total / workers)
    chunks = [traces[i : i + chunk_size] for i in range(0, total, chunk_size)]

    label = "Re-scoring" if rejudge else "Scoring"
    print(f"\n  {label} {total} traces for '{project}' ({workers} workers)...\n")

    results: list[dict] = []

    def run_chunk(chunk: list[dict[str, Any]]) -> dict:
        try:
            return run_judge(traces=chunk, client=client, project=project, rejudge=rejudge)
        except RuntimeError as e:
            return {"error": str(e), "trace_results": [], "total_scored": 0, "passed": 0, "failed": 0}

    if workers == 1:
        results.append(run_chunk(chunks[0]))
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
            futures = list(executor.map(run_chunk, chunks))
            results = list(futures)

    combined = _combine_results(results)
    print(format_run_summary(combined))

    if combined.get("error_count"):
        print(f"  Warning: {combined['error_count']} chunk(s) failed. Results may be incomplete.")


def _combine_results(results: list[dict]) -> dict:
    """Combine multiple judge run results into one."""
    combined: dict = {
        "trace_results": [],
        "total_scored": 0,
        "passed": 0,
        "failed": 0,
        "parse_errors": 0,
        "error_count": 0,
    }
    for r in results:
        if r.get("error"):
            combined["error_count"] += 1
            logger.warning(f"Judge chunk failed: {r.get('error')}")
            continue
        combined["trace_results"].extend(r.get("trace_results", []))
        combined["total_scored"] += r.get("total_scored", 0)
        combined["passed"] += r.get("passed", 0)
        combined["failed"] += r.get("failed", 0)
        combined["parse_errors"] += r.get("parse_errors", 0)
    return combined
