#!/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
from pathlib import Path
from zoneinfo import ZoneInfo
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}


IL_TZ = ZoneInfo("Asia/Jerusalem")


def _ts() -> str:
    """Filename-safe report stamp in Israel time (INV-UI9: human-facing display).

    This script runs on the host, which is UTC, so a process-clock stamp would
    mislead the chair. We pin Asia/Jerusalem and suffix ``IL`` (instead of a bare
    ``Z``) so the displayed/filename time is unambiguously Israel-local.
    """
    return datetime.now(IL_TZ).strftime("%Y%m%dT%H%M%S-IL")


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"]
    il_now = datetime.now(IL_TZ).strftime("%Y-%m-%d %H:%M")
    lines = [
        f"# #208 — כיול model×effort מול הסופיים — {ts}\n",
        f"> נוצר: {il_now} (שעון ישראל · Asia/Jerusalem)\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()))