"""Phase 3 Step 1 — Critical Engagement (5 types) base-8B measurement, STEP-A style (sharp GT-aligned
definitions, NO few-shot). Per-type κ + confusion + overall. Reuses the base-8B cache (:8001).

Run:  python -m eval.step_critical
"""
from __future__ import annotations

import json
import os
import sys

from eval import kappa as K
from eval.prompts_v2 import _ask
from prompt_card.scoring import observable_axes as OA

CE = ["skepticism", "rebuttal", "source_request", "independent_verification", "re_questioning"]
_CE_T = ('{"skepticism": <bool>, "rebuttal": <bool>, "source_request": <bool>, '
         '"independent_verification": <bool>, "re_questioning": <bool>}')


def build_critical_prompt(prev_assistant: str, this_user: str) -> str:
    instr = (
        "Detect critical-engagement TYPES in THE USER's turn as it reacts to what the AI just said. Multiple "
        "types can be true at once; mark true only when genuinely present.\n"
        "- skepticism: doubts/questions the AI's claim WITHOUT giving a reason. ▸ 'really?', 'are you sure?', "
        "'that doesn't sound right'. ✗ 'ok, thanks' (acceptance).\n"
        "- rebuttal: pushes back with the user's OWN counter-argument or correction. ▸ 'that's wrong — if X "
        "were true, Y wouldn't happen', 'no, it returns 0 not 1'. ✗ a bare 'are you sure?' (that is skepticism).\n"
        "- source_request: asks for a citation / evidence / source. ▸ 'what's your source?', 'cite that'. "
        "Note: asking for REASONING ('why did you pick X?') is NOT source_request.\n"
        "- independent_verification: the user states an EXPLICIT external check they performed. ▸ 'I ran it and "
        "it throws on empty input', 'I checked the docs, it says X'. ✗ a bare confident assertion with no stated "
        "check (that is rebuttal).\n"
        "- re_questioning: RE-ASKS the same question because the AI's answer was unsatisfactory. ▸ 'that's not "
        "what I asked — I meant the async case'. ✗ a NEW/different question."
    )
    payload = f"AI just said:\n{prev_assistant}\n\nUser's turn:\n{this_user}"
    return _ask(instr, _CE_T, "Turns", payload)


def main():
    base = os.environ.get("OPENBMB_BASE_URL"); token = os.environ.get("OPENBMB_TOKEN")
    if not base or not token:
        print("ERROR: creds", file=sys.stderr); sys.exit(2)
    from prompt_card.llm.minicpm import MiniCPMClient
    gt = K.load_gt(); convs = K.load_convs()
    client = K.CachedClient(MiniCPMClient(base, token), workers=8)

    prompts, rows = [], []
    for r in gt:
        conv = convs[r["id"]]
        for j, row in enumerate(r["critical"]):
            prev = OA._prev_assistant(conv, int(row["turn"][1:]) - 1)
            users = OA._user_turns(conv)
            utext = users[int(row["turn"][1:]) - 1]
            prompts.append(build_critical_prompt(prev, utext))
            rows.append(set(row["types"]))
    print(f"[critical] {len(prompts)} turn-prompts", flush=True)
    resp = client.run_all(prompts)
    print(f"[critical] 8B done (new {client.misses})", flush=True)

    preds = []
    fail = 0
    for p in prompts:
        d = OA.parse(resp[p], CE)
        if d is None:
            fail += 1; d = {}
        preds.append({t for t in CE if d.get(t)})

    print("\n=== Critical Engagement — base-8B per-type κ (vs Kim ground truth) ===")
    print(f"parse failures: {fail}/{len(prompts)}")
    per_k = {}
    for t in CE:
        yt = [int(t in s) for s in rows]; yp = [int(t in s) for s in preds]
        k = K.cohen_kappa(yt, yp); c = K.binary_counts(yt, yp); p, rc, f1 = K.prf(c)
        per_k[t] = k
        ks = f"{k:+.3f}" if k is not None else "N/A"
        print(f"  {t:24} κ={ks}  pos={sum(yt)}  [TN {c['tn']} FP {c['fp']} FN {c['fn']} TP {c['tp']}]  "
              f"P{'-' if p is None else f'{p:.2f}'} R{'-' if rc is None else f'{rc:.2f}'}")
    # any-CE per turn
    anyt = [int(bool(s)) for s in rows]; anyp = [int(bool(s)) for s in preds]
    anyk = K.cohen_kappa(anyt, anyp)
    # distinct-type count per conv (the product signal): agreement
    gi = 0; gt_counts = []; pr_counts = []
    for r in gt:
        n = len(r["critical"])
        gtypes = set(); ptypes = set()
        for _ in range(n):
            gtypes |= rows[gi]; ptypes |= preds[gi]; gi += 1
        gt_counts.append(len(gtypes)); pr_counts.append(len(ptypes))
    # mean abs error on the 0-5 distinct-type signal
    mae = sum(abs(a - b) for a, b in zip(gt_counts, pr_counts)) / len(gt_counts)
    valid = [per_k[t] for t in CE if per_k[t] is not None]
    headline = sum(valid) / len(valid) if valid else None
    print(f"\n  any-CE-present (per turn) κ = {anyk:+.3f}")
    print(f"  per-type mean κ (headline) = {headline:+.3f}")
    print(f"  distinct-type count (0-5) MAE per conv = {mae:.2f}")

    if headline is None:
        decision = "N/A"
    elif headline >= 0.6:
        decision = "NO LoRA — base 8B solid"
    elif headline >= 0.4:
        decision = "NO LoRA — borderline, document caveats (save time)"
    else:
        decision = "LoRA NEEDED (headline < 0.4)"
    print(f"\n  >>> CRITICAL LoRA DECISION: {decision}")
    json.dump({"per_type": per_k, "any": anyk, "headline": headline, "mae": mae,
               "decision": decision, "parse_fail": fail},
              open(os.path.join(os.path.dirname(__file__), "_cache", "critical.json"), "w"), indent=1, default=str)


if __name__ == "__main__":
    main()