#!/usr/bin/env python3
"""#208 (WS5 / Q1, INV-G8 eval-harness) — model×effort calibration vs the finals.

Empirically picks the per-block reasoning `effort` whose regenerated block lands
CLOSEST to עו"ד דפנה תמיר's signed final, over the EXISTING `draft_final_pairs`
ledger. This is the calibration INV-G8 / 07-learning §0.7 ask for: stop choosing
the per-block effort defaults (#204: ה=medium, ו=medium, ז=high, ח=medium, ט=high)
"by feel".

WHAT IT MEASURES (per (case, block, effort) cell — the A/B grid):
  Regenerate `block` for `case` via block_writer.write_block(effort_override=…)
  (the PRODUCTION generation path → claude_session.query → `claude -p`, pinned
  Opus 4.8, local-only), then score the regenerated block against the matching
  SECTION of the final via services.style_distance.block_distance_to_final:
    • change_percent           — word-diff regen↔final-section (compute_diff_stats)
    • anti_pattern_total       — lessons.ANTI_PATTERNS hits in the regen (cleanest
                                 style signal — see §0.7 warning below)
    • golden_ratio_deviation_pp — structural-weight gap vs the final
    • distance                 — normalized composite (lower = closer to Dafna)
  No parallel metric path: it reuses style_distance + learning_loop (G2).

RECOMMENDATION: for each block, the effort with the lowest MEAN composite distance
across cases (ties → fewer anti-patterns → lower change_percent). Reported next to
the #204 current default so a regression/improvement is visible.

⚠️ SAMPLE-SIZE CAVEAT (honored, not hidden): very few cases have an uploaded final
(draft_final_pairs.final_text non-empty). The report prints n_finals PROMINENTLY and
labels the output DIRECTIONAL EVIDENCE, not a regression. With n<3 per block the
recommendation is advisory only; the chair/operator decides whether to adopt.

⚠️ change_percent mixes style with content completeness (07-learning §0.7): the chair
sometimes doubles length for missing substance. anti_pattern_total is the cleaner
style signal — the report surfaces both, and the composite down-weights neither
silently.

GENERATION PATH (do not violate — reference_claude_generation_path / claude_session
docstring): write_block → claude_session.query → `claude -p` uses the local claude.ai
session. It runs ONLY on the host where the `claude` CLI exists (NOT the legal-ai
container, NOT a symlinked worktree without CLI access). Hence the live A/B is
host-only; --self-test proves the measurement logic offline with zero model calls.

Usage (mcp-server venv; live needs POSTGRES + the `claude` CLI on the host):
    PY=/home/chaim/legal-ai/mcp-server/.venv/bin/python
    $PY scripts/calibrate_effort.py --self-test            # offline proof, no DB/CLI
    POSTGRES_PASSWORD=… POSTGRES_HOST=127.0.0.1 POSTGRES_PORT=5433 \
        $PY scripts/calibrate_effort.py                    # live A/B over finals
    … --efforts low,medium,high,xhigh                      # override the effort grid
    … --blocks block-he,block-vav                          # restrict to some blocks
    … --case 8137-11-24                                    # a single case
    … --repeats 2                                          # avg N gens/cell (noise)
    … --dry-run                                            # plan the grid, no model calls
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from statistics import mean

logger = logging.getLogger(__name__)

REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "mcp-server" / "src"))

if "POSTGRES_URL" not in os.environ:
    os.environ["POSTGRES_URL"] = (
        f"postgres://{os.environ.get('POSTGRES_USER','legal_ai')}:"
        f"{os.environ.get('POSTGRES_PASSWORD','')}@"
        f"{os.environ.get('POSTGRES_HOST','127.0.0.1')}:"
        f"{os.environ.get('POSTGRES_PORT','5433')}/"
        f"{os.environ.get('POSTGRES_DB','legal_ai')}"
    )

OUT_DIR = REPO_ROOT / "data" / "eval"

# Only the AI blocks that map to a golden-ratio section can be scored against the
# final's matching section (split_final_by_section). Template blocks (א-ד, יב) are
# deterministic template-fill — no effort knob. block-yod (discussion) is EXCLUDED
# from the defaults: it requires an approved direction (brainstorm → approve_direction)
# that calibration cases lack, so block_writer.write_block raises and it is NOT
# calibratable standalone — and it's out of WS5's interim-draft scope. It stays in
# VALID_BLOCKS so a user can still force it (--blocks block-yod) at their own risk;
# the per-cell guard in _run() keeps such a failure non-fatal.
CALIBRATABLE_BLOCKS = ["block-he", "block-vav", "block-zayin", "block-chet", "block-tet", "block-yod-alef"]
VALID_BLOCKS = CALIBRATABLE_BLOCKS + ["block-yod"]
DEFAULT_EFFORTS = ["low", "medium", "high", "xhigh"]
VALID_EFFORTS = {"low", "medium", "high", "xhigh", "max"}


# ── pure helpers (offline-testable) ──────────────────────────────────────────
def recommend_effort(cells: list[dict]) -> dict | None:
    """Pick the best effort for ONE block from its scored cells.

    cells: [{"effort","distance","anti_pattern_total","change_percent","n"}].
    Lowest mean composite distance wins; ties broken by fewer anti-patterns,
    then lower change_percent. Pure → unit-tested in --self-test.
    """
    if not cells:
        return None
    ranked = sorted(
        cells,
        key=lambda c: (c["distance"], c["anti_pattern_total"], c["change_percent"]),
    )
    return ranked[0]


def aggregate_cell(per_run: list[dict]) -> dict:
    """Mean each metric across repeated generations of the same (case, block, effort)."""
    if not per_run:
        return {"distance": 1.0, "anti_pattern_total": 0.0, "change_percent": 100.0,
                "golden_ratio_deviation_pp": None, "n": 0}
    ratios = [r["golden_ratio_deviation_pp"] for r in per_run if r.get("golden_ratio_deviation_pp") is not None]
    return {
        "distance": round(mean(r["distance"] for r in per_run), 4),
        "anti_pattern_total": round(mean(r["anti_pattern_total"] for r in per_run), 2),
        "change_percent": round(mean(r["change_percent"] for r in per_run), 2),
        "golden_ratio_deviation_pp": round(mean(ratios), 2) if ratios else None,
        "n": len(per_run),
    }


def _current_default(block_id: str) -> str | None:
    from legal_mcp.services.block_writer import BLOCK_CONFIG, DEFAULT_EFFORT
    cfg = BLOCK_CONFIG.get(block_id, {})
    if cfg.get("model") != "ai":
        return None
    return cfg.get("effort", DEFAULT_EFFORT)


# ── self-test (no DB, no model) ──────────────────────────────────────────────
def _self_test() -> int:
    ok = True

    def chk(name, cond):
        nonlocal ok
        ok = ok and cond
        print(f"  {name:42} {'ok' if cond else 'FAIL'}")

    # block_distance_to_final: regen identical to final-section ⇒ change ~0,
    # distance dominated by anti-patterns (0 here) ⇒ ~0.
    from legal_mcp.services.style_distance import (
        block_distance_to_final, split_final_by_section,
    )
    final_section = "לפנינו ערר על החלטת הוועדה המקומית. " * 40
    d_same = block_distance_to_final("block-he", final_section, final_section, "rejection",
                                     section_target_total_words=len(final_section.split()))
    chk("identical regen ⇒ change_percent==0", d_same["change_percent"] == 0.0)
    chk("identical regen ⇒ anti==0", d_same["anti_pattern_total"] == 0)
    chk("identical regen ⇒ distance small", d_same["distance"] < 0.05)

    # A regen full of anti-patterns (markdown headers / bullet lists) scores worse
    # than clean continuous narrative, holding the final fixed.
    clean = "אנו סבורים כי דין הערר להידחות. כידוע, הלכה פסוקה היא. " * 20
    dirty = "## כותרת\n- נקודה ראשונה\n- נקודה שנייה\n* עוד נקודה\n### תת\n" * 10
    d_clean = block_distance_to_final("block-yod", clean, final_section, "rejection",
                                      section_target_total_words=len(final_section.split()))
    d_dirty = block_distance_to_final("block-yod", dirty, final_section, "rejection",
                                      section_target_total_words=len(final_section.split()))
    chk("dirty regen has more anti-patterns", d_dirty["anti_pattern_total"] > d_clean["anti_pattern_total"])
    chk("dirty regen ⇒ larger distance", d_dirty["distance"] > d_clean["distance"])

    # ratio deviation is None when no total provided (never fabricated)
    d_noratio = block_distance_to_final("block-he", clean, final_section, "rejection")
    chk("no total ⇒ ratio deviation None", d_noratio["golden_ratio_deviation_pp"] is None)

    # split_final_by_section returns mapped golden sections only.
    sample_final = (
        "רקע עובדתי\nהמקרקעין נשוא הערר. " * 10 + "\n\n"
        "תמצית טענות הצדדים\nהעוררים טוענים כי. " * 10 + "\n\n"
        "דיון והכרעה\nאנו סבורים כי. " * 10
    )
    secs = split_final_by_section(sample_final)
    chk("split maps to golden sections", set(secs).issubset(
        {"background", "claims", "discussion", "summary"}))
    chk("split found ≥1 section", len(secs) >= 1)

    # recommend_effort: lowest distance wins; tie → fewer anti-patterns.
    rec = recommend_effort([
        {"effort": "low", "distance": 0.40, "anti_pattern_total": 5, "change_percent": 40, "n": 1},
        {"effort": "high", "distance": 0.20, "anti_pattern_total": 2, "change_percent": 25, "n": 1},
        {"effort": "xhigh", "distance": 0.20, "anti_pattern_total": 1, "change_percent": 22, "n": 1},
    ])
    chk("recommend picks lowest distance", rec["distance"] == 0.20)
    chk("recommend tie → fewer anti-patterns", rec["effort"] == "xhigh")
    chk("recommend empty ⇒ None", recommend_effort([]) is None)

    # aggregate_cell: means + ratio drops Nones, keeps n.
    agg = aggregate_cell([
        {"distance": 0.2, "anti_pattern_total": 2, "change_percent": 20, "golden_ratio_deviation_pp": 3.0},
        {"distance": 0.4, "anti_pattern_total": 4, "change_percent": 30, "golden_ratio_deviation_pp": None},
    ])
    chk("aggregate distance mean", agg["distance"] == 0.3)
    chk("aggregate n counted", agg["n"] == 2)
    chk("aggregate ratio skips None", agg["golden_ratio_deviation_pp"] == 3.0)

    print("ALL PASS" if ok else "*** FAILURES ***")
    return 0 if ok else 1


# ── live A/B (host-only — needs DB + `claude` CLI) ───────────────────────────
async def _finals_for_calibration(case_filter: str | None) -> list[dict]:
    """draft_final_pairs whose final_text is populated (the held-out comparison set)."""
    from legal_mcp.services import db
    pairs = await db.list_draft_final_pairs(limit=500)
    out: list[dict] = []
    for p in pairs:
        if case_filter and p.get("case_number") != case_filter:
            continue
        full = await db.get_draft_final_pair(p["id"])
        if full and (full.get("final_text") or "").strip():
            out.append(full)
    return out


async def _score_cell(case_id, block_id: str, effort: str, final_section: str,
                      final_total_words: int, outcome: str, repeats: int) -> dict:
    """Generate `block_id` at `effort` `repeats` times; score each vs the final section."""
    from legal_mcp.services import block_writer
    from legal_mcp.services.style_distance import block_distance_to_final
    runs: list[dict] = []
    for _ in range(repeats):
        res = await block_writer.write_block(case_id, block_id, effort_override=effort)
        scored = block_distance_to_final(
            block_id, res.get("content", ""), final_section, outcome,
            section_target_total_words=final_total_words,
        )
        runs.append(scored)
    agg = aggregate_cell(runs)
    agg["effort"] = effort
    agg["runs"] = runs
    return agg


async def _run(args, ts: str) -> dict:
    from uuid import UUID
    from legal_mcp.services import db
    from legal_mcp.services.lessons import canonical_outcome
    from legal_mcp.services.style_distance import split_final_by_section, _BLOCK_TO_SECTION

    efforts = args.efforts
    blocks = args.blocks
    finals = await _finals_for_calibration(args.case)

    cases_meta = []
    for f in finals:
        case = await db.get_case_by_number(f["case_number"]) if f.get("case_number") else None
        if not case:
            continue
        decision = await db.get_decision_by_case(UUID(case["id"]))
        outcome = canonical_outcome((decision or {}).get("outcome", "rejection"))
        sections = split_final_by_section(f.get("final_text", ""))
        final_total_words = len((f.get("final_text", "") or "").split())
        cases_meta.append({
            "case_number": f["case_number"], "case_id": case["id"],
            "outcome": outcome, "sections": sections, "final_total_words": final_total_words,
        })

    # grid plan: (block → cases that have its section)
    plan: dict[str, list[dict]] = {}
    for block_id in blocks:
        section = _BLOCK_TO_SECTION.get(block_id)
        plan[block_id] = [c for c in cases_meta if section and c["sections"].get(section)]

    total_cells = sum(len(plan[b]) for b in blocks) * len(efforts) * args.repeats
    grid_summary = {
        "n_finals": len(cases_meta),
        "finals": [c["case_number"] for c in cases_meta],
        "blocks": blocks, "efforts": efforts, "repeats": args.repeats,
        "total_generations": total_cells,
        "per_block_n": {b: len(plan[b]) for b in blocks},
    }

    if args.dry_run:
        return {"dry_run": True, "grid": grid_summary, "by_block": {}}

    by_block: dict[str, dict] = {}
    for block_id in blocks:
        section = _BLOCK_TO_SECTION.get(block_id)
        per_effort_runs: dict[str, list[dict]] = {e: [] for e in efforts}
        per_case: list[dict] = []
        for c in plan[block_id]:
            final_section = c["sections"][section]
            case_cells = []
            for effort in efforts:
                # Per-cell guard: ANY failure of a single (case, block, effort) cell
                # — e.g. block-yod raising "ללא כיוון מאושר", or a transient rate-limit
                # mid-run — is logged and skipped, NOT allowed to kill the whole run
                # (INV-G8 eval-harness robustness). Partial results still persist below.
                try:
                    cell = await _score_cell(
                        UUID(c["case_id"]), block_id, effort, final_section,
                        c["final_total_words"], c["outcome"], args.repeats,
                    )
                except Exception as exc:  # noqa: BLE001 — harness must survive any cell failure
                    logger.warning(
                        "calibration cell skipped: case=%s block=%s effort=%s — %s",
                        c["case_number"], block_id, effort, exc,
                    )
                    continue
                per_effort_runs[effort].append(cell)
                case_cells.append({k: cell[k] for k in
                                   ("effort", "distance", "anti_pattern_total",
                                    "change_percent", "golden_ratio_deviation_pp", "n")})
            per_case.append({"case_number": c["case_number"], "cells": case_cells})

        # mean across cases for each effort → one comparable row per effort
        effort_rows = []
        for effort in efforts:
            rows = per_effort_runs[effort]
            if not rows:
                continue
            ratios = [r["golden_ratio_deviation_pp"] for r in rows if r.get("golden_ratio_deviation_pp") is not None]
            effort_rows.append({
                "effort": effort,
                "distance": round(mean(r["distance"] for r in rows), 4),
                "anti_pattern_total": round(mean(r["anti_pattern_total"] for r in rows), 2),
                "change_percent": round(mean(r["change_percent"] for r in rows), 2),
                "golden_ratio_deviation_pp": round(mean(ratios), 2) if ratios else None,
                "n": len(rows),
            })
        rec = recommend_effort(effort_rows)
        by_block[block_id] = {
            "section": section,
            "current_default": _current_default(block_id),
            "recommended": rec["effort"] if rec else None,
            "efforts": effort_rows,
            "per_case": per_case,
        }

        # Incremental persistence (INV-G8): flush the report after every completed
        # block so a crash later in the grid never loses blocks already generated.
        # Blocks not yet done are simply absent from by_block; _write_report tolerates
        # partial results. main() does the final flush once the loop finishes.
        try:
            _write_report({"dry_run": False, "grid": grid_summary, "by_block": by_block}, ts)
        except Exception as exc:  # noqa: BLE001 — a write hiccup must not abort the run
            logger.warning("incremental report write failed after block=%s — %s", block_id, exc)

    return {"dry_run": False, "grid": grid_summary, "by_block": by_block}


def _ts() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def _write_report(result: dict, ts: str) -> tuple[Path, Path]:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    jp = OUT_DIR / f"effort-calibration-{ts}.json"
    mp = OUT_DIR / f"effort-calibration-{ts}.md"
    jp.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")

    g = result["grid"]
    n = g["n_finals"]
    lines = [
        f"# #208 — כיול model×effort מול הסופיים — {ts}\n",
        f"> ⚠️ **גודל-מדגם: {n} סופיים** ({', '.join(g['finals']) or '—'}). "
        "זוהי **עדות-כיוון, לא רגרסיה** — מעט תיקים בעלי סופי-עלוי. "
        "ההמלצה אדוויזורית; ההכרעה בידי היו\"ר/המפעיל.\n",
        f"- בלוקים: {', '.join(g['blocks'])}",
        f"- efforts: {', '.join(g['efforts'])} · repeats/cell: {g['repeats']}",
        f"- סך ייצורי-מודל: {g['total_generations']}",
        "",
    ]
    if result.get("dry_run"):
        lines += ["## DRY-RUN — תכנון הגריד בלבד (ללא ייצור)\n",
                  "| block | #cases | current default |", "|---|---|---|"]
        for b in g["blocks"]:
            lines.append(f"| {b} | {g['per_block_n'].get(b,0)} | {_current_default(b) or '—'} |")
        mp.write_text("\n".join(lines) + "\n", encoding="utf-8")
        return jp, mp

    lines += ["## המלצה per-בלוק (distance נמוך = קרוב יותר לדפנה)\n",
              "| block | section | current | **recommended** | n |", "|---|---|---|---|---|"]
    for b, bd in result["by_block"].items():
        rec = bd.get("recommended") or "—"
        mark = "" if rec == bd.get("current_default") else " ⬅︎"
        n_b = bd["efforts"][0]["n"] if bd.get("efforts") else 0
        lines.append(f"| {b} | {bd.get('section','')} | {bd.get('current_default') or '—'} | **{rec}**{mark} | {n_b} |")
    lines.append("")
    for b, bd in result["by_block"].items():
        lines += [f"### {b} ({bd.get('section','')})\n",
                  "| effort | distance | anti_total | change% | ratioΔpp | n |",
                  "|---|---|---|---|---|---|"]
        for r in bd.get("efforts", []):
            star = " ⭐" if r["effort"] == bd.get("recommended") else ""
            ratio = r["golden_ratio_deviation_pp"]
            lines.append(
                f"| {r['effort']}{star} | {r['distance']:.4f} | {r['anti_pattern_total']} | "
                f"{r['change_percent']} | {ratio if ratio is not None else '—'} | {r['n']} |")
        lines.append("")
    lines.append("> change% מערבב סגנון עם שלמות-תוכן (07-learning §0.7); "
                 "anti_total הוא הסיגנל הנקי-יותר לסגנון.\n")
    mp.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return jp, mp


async def main() -> int:
    ap = argparse.ArgumentParser(description="#208 model/effort calibration harness")
    ap.add_argument("--self-test", action="store_true", help="offline measurement-logic proof (no DB/CLI)")
    ap.add_argument("--dry-run", action="store_true", help="plan the A/B grid over existing finals, no model calls")
    ap.add_argument("--efforts", default=",".join(DEFAULT_EFFORTS),
                    help=f"comma effort grid (default {','.join(DEFAULT_EFFORTS)})")
    ap.add_argument("--blocks", default=",".join(CALIBRATABLE_BLOCKS),
                    help="comma block ids to calibrate")
    ap.add_argument("--case", default=None, help="restrict to a single case_number")
    ap.add_argument("--repeats", type=int, default=1, help="generations per cell (avg out gen noise)")
    args = ap.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    if args.self_test:
        return _self_test()

    args.efforts = [e.strip() for e in args.efforts.split(",") if e.strip()]
    bad = [e for e in args.efforts if e not in VALID_EFFORTS]
    if bad:
        print(f"invalid effort(s): {bad}. valid: {sorted(VALID_EFFORTS)}", file=sys.stderr)
        return 2
    args.blocks = [b.strip() for b in args.blocks.split(",") if b.strip()]
    bad_b = [b for b in args.blocks if b not in VALID_BLOCKS]
    if bad_b:
        print(f"non-calibratable block(s): {bad_b}. valid: {VALID_BLOCKS}", file=sys.stderr)
        return 2

    ts = _ts()
    result = await _run(args, ts)
    jp, mp = _write_report(result, ts)

    g = result["grid"]
    print(f"CALIBRATION: {g['n_finals']} finals — DIRECTIONAL EVIDENCE, not a regression")
    if g["n_finals"] == 0:
        print("  no finals with final_text in draft_final_pairs — upload a signed final first.")
    elif result.get("dry_run"):
        print(f"  dry-run: {g['total_generations']} generations planned across {len(g['blocks'])} blocks")
    else:
        for b, bd in result["by_block"].items():
            print(f"  {b:16} current={bd.get('current_default') or '—':6} "
                  f"→ recommended={bd.get('recommended') or '—'}")
    print(f"  report: {mp}")
    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))