#!/usr/bin/env python3
"""推論の記録の判定ファイルに, 今ある判定の写しと AI の先回り判定を足す (依頼 183).

  python3 pi05/verdict_viewer/seed_verdicts.py --run-input      # {run_stem}.json の outcome を写す (labeler=run_input)
  python3 pi05/verdict_viewer/seed_verdicts.py --ai             # 印のない回を AI が判定 (labeler=ai)
  python3 pi05/verdict_viewer/seed_verdicts.py --ai --dry-run   # API を呼ばずに素材だけ数える
  python3 pi05/verdict_viewer/seed_verdicts.py --rule           # エイイチの決め (依頼 185) の行を足す (labeler=rule)

- 元の記録は読むだけ. 判定ファイルには追記だけ (verdicts.append_rows)
- すでに同じ種類の行がある回は足さない (何度流しても同じ行が増えない)
- AI の素材: recording/top.mp4 があれば, 推論が動いていた区間 (top_frame_log.csv の mode が wait / home
  でないコマ) の終わりまでから 6 コマ (後ろ寄り). なければ frames/ (各チャンクの入力画像) から最後を含む 6 枚.
  2 枚に満たなければ API を呼ばずに not_judgeable
  (終わりの wait / home では人がブロックを箱から出していることがあるので, そこは使わない)
- API は今の vlm_judge (dynamic_tracking_limits/place/place_modules/vlm_judge.py) と同じモデル・キーの読み方.
  使えない (キーがない・通信やクレジットのエラー) 回は行を書かずに飛ばし, 最後に数を出す
- --rule (依頼 185): 送ったチャンク 0 の回 → not_applicable (決め 2), AI が frames/ だけ・素材なしで
  not_judgeable とした回 → not_judgeable (決め 3). 両方なら 2. 確定の行 (run_input・人・rule) がある回には足さない
"""
import argparse
import base64
import csv
import json
import os
import sys
import time
import urllib.error
import urllib.request
from collections import Counter
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
sys.path.insert(0, str(HERE.parent.parent / "dynamic_tracking_limits" / "place" / "place_modules"))
import verdicts as V  # noqa: E402
import vlm_judge  # noqa: E402

N_IMAGES = 6
MIN_IMAGES = 2
MP4_FRACTIONS = (0.3, 0.5, 0.65, 0.8, 0.9, 1.0)
MP4_WIDTH = 960
API_TIMEOUT_SEC = 90.0
API_MAX_TOKENS = 600

PROMPT = """画像は mechArm 270 (ロボットアーム) の作業台を真上から見たもので, 時間の順に並んでいる ({n} 枚・{source}).
課題は「赤いブロックをつかみ, 白い箱の中に置く」(put the red block in the white box). AI の方策が動かしている.
{source_note}
この回について, 次を判定してください.
- pick: 赤いブロックを爪でつかんで持ち上げられたか (yes / no / unknown)
- place: 最後に赤いブロックが白い箱の中にあるか (yes / no / unknown)
- outcome: success (最後に赤いブロックが白い箱の中にある) / failure (箱の中にない・つかめていない・落とした) /
  unknown (画像はあるが判別できない) / not_judgeable (画像が終わりまで写っておらず判定できない)
以下の JSON だけで返答してください. 前置きや後置きの説明は一切含めないこと.
{{"pick": "yes|no|unknown", "place": "yes|no|unknown", "outcome": "success|failure|unknown|not_judgeable", "reason": "日本語で 1〜2 文の短い理由"}}"""

NOTE_MP4 = ("上カメラの録画から, 方策が動いていた区間の後ろ寄りのコマを選んだ. 最後の 1 枚は方策が止まる直前. "
            "人の手が写ることがある (補正や置き直し).")
NOTE_FRAMES = ("各画像は方策に渡した入力画像で, チャンク (動きのひとまとまり) を送る前に撮ったもの. "
               "最後の 1 枚は最後のチャンクを送る前なので, その後の動きは写っていない. "
               "赤いブロックが写っていない画像があってもよい (爪の中や箱の中にあることがある).")


def mp4_active_end(run_dir, n_frames):
    """top_frame_log.csv で mode が wait / home でない最後のコマ. 読めなければ最後のコマ."""
    p = V.safe_file(run_dir, "recording", "top_frame_log.csv")
    last = None
    try:
        if p is None:
            raise FileNotFoundError
        with open(p, newline="") as f:
            for r in csv.DictReader(f):
                if r.get("mode") not in ("wait", "home"):
                    last = int(r["frame_idx"])
    except Exception:
        last = None
    if last is None:
        return n_frames - 1
    return max(0, min(last, n_frames - 1))


def mp4_images(run_dir):
    import cv2
    cap = cv2.VideoCapture(str(Path(run_dir) / "recording" / "top.mp4"))
    try:
        n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if n <= 0:
            return [], {}
        end = mp4_active_end(run_dir, n)
        idxs = sorted(set(int(round(end * fr)) for fr in MP4_FRACTIONS))
        out = []
        for i in idxs:
            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
            ok, img = cap.read()
            if not ok or img is None:
                continue
            h, w = img.shape[:2]
            if w > MP4_WIDTH:
                img = cv2.resize(img, (MP4_WIDTH, int(h * MP4_WIDTH / w)))
            ok, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 85])
            if ok:
                out.append(buf.tobytes())
        return out, {"mp4_frames": n, "active_end": end, "picked": idxs}
    finally:
        cap.release()


def frames_images(run_dir):
    names = V.frame_files(run_dir)
    if not names:
        return [], {}
    if len(names) <= N_IMAGES:
        pick = names
    else:
        step = (len(names) - 1) / (N_IMAGES - 1)
        pick = [names[int(round(i * step))] for i in range(N_IMAGES)]
    out = [(Path(run_dir) / "frames" / nm).read_bytes() for nm in pick]
    return out, {"frames": len(names), "picked": pick}


def material(run_dir):
    """(evidence, 画像の bytes のリスト, メモ). 素材が足りなければ evidence='' で画像は空."""
    if V.has_top_mp4(run_dir):
        imgs, info = mp4_images(run_dir)
        if len(imgs) >= MIN_IMAGES:
            return "top.mp4", imgs, info
    imgs, info = frames_images(run_dir)
    if len(imgs) >= MIN_IMAGES:
        return "frames", imgs, info
    return "", [], info


def api_key():
    k = os.environ.get("ANTHROPIC_API_KEY")
    if not k:
        p = os.path.expanduser("~/.anthropic_key")
        if os.path.isfile(p):
            k = open(p).read().strip()
            # 「export ANTHROPIC_API_KEY=...」の形で置かれていることがある (2026-11-14 に確かめた)
            if k.startswith("export "):
                k = k[len("export "):].strip()
            if k.startswith("ANTHROPIC_API_KEY="):
                k = k[len("ANTHROPIC_API_KEY="):].strip()
            k = k.strip("'\"")
    return k or None


class ApiUnavailable(RuntimeError):
    pass


def call_api(key, imgs, evidence):
    source = "上カメラの録画 top.mp4 のコマ" if evidence == "top.mp4" else "frames/ の入力画像"
    note = NOTE_MP4 if evidence == "top.mp4" else NOTE_FRAMES
    content = [{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg",
                                            "data": base64.b64encode(b).decode("ascii")}} for b in imgs]
    content.append({"type": "text", "text": PROMPT.format(n=len(imgs), source=source, source_note=note)})
    payload = {"model": vlm_judge.VLM_MODEL, "max_tokens": API_MAX_TOKENS,
               "messages": [{"role": "user", "content": content}]}
    req = urllib.request.Request(vlm_judge.VLM_API_URL, data=json.dumps(payload).encode("utf-8"),
                                 headers={"x-api-key": key, "anthropic-version": vlm_judge.VLM_ANTHROPIC_VERSION,
                                          "content-type": "application/json"}, method="POST")
    last_err = ""
    for attempt in range(3):
        try:
            with urllib.request.urlopen(req, timeout=API_TIMEOUT_SEC) as resp:
                body = json.loads(resp.read().decode("utf-8"))
            text = "".join(c.get("text", "") for c in body.get("content", []) if c.get("type") == "text")
            return vlm_judge._parse_json_from_text(text), text
        except urllib.error.HTTPError as e:
            last_err = f"http_{e.code}"
            if e.code in (429, 500, 502, 503, 504, 529):
                time.sleep(5 * (attempt + 1))
                continue
            raise ApiUnavailable(last_err)
        except Exception as e:
            last_err = type(e).__name__
            time.sleep(3)
    raise ApiUnavailable(last_err)


def ai_row(run_stem, parsed, evidence, n_exec):
    def pv(v):
        return v if v in V.PICK_PLACE else "unknown"
    oc = parsed.get("outcome") if isinstance(parsed, dict) else None
    oc = oc if oc in V.AI_OUTCOMES else "unknown"
    reason = str(parsed.get("reason", "")).replace("\n", " ").strip() if isinstance(parsed, dict) else ""
    return {"labeled_at": V.now_iso(), "run_stem": run_stem,
            "pick": pv(parsed.get("pick")), "place": pv(parsed.get("place")), "outcome": oc, "labeler": "ai",
            "ai_outcome": oc, "ai_reason": f"[送ったチャンク {n_exec}] {reason}", "evidence": evidence, "note": ""}


def do_run_input(csv_path, runs_dir, dry_run):
    have = V.latest_by_stem(V.read_rows(csv_path), labelers=("run_input",) + V.HUMAN_LABELERS)
    rows = []
    for stem in V.list_run_stems(runs_dir):
        if stem in have:
            continue
        r = V.run_input_row(stem, V.read_run_json(Path(runs_dir) / stem))
        if r:
            rows.append(r)
    c = Counter(r["outcome"] for r in rows)
    print(f"[COPY ] 写す行 {len(rows)} (success {c['success']} / failure {c['failure']} / unknown {c['unknown']})"
          + (" ※ dry-run なので書いていません" if dry_run else ""))
    if not dry_run:
        written = V.append_rows(rows, csv_path, skip_if_labeled_by=("run_input",) + V.HUMAN_LABELERS,
                                 forbid_dirs=(runs_dir,))
        if len(written) != len(rows):
            print(f"[COPY ] 書くまでの間に人か写しの行が足されていた {len(rows) - len(written)} 回は足さなかった")
    return c


def do_ai(csv_path, runs_dir, dry_run, limit, only):
    rows = V.read_rows(csv_path)
    marked = set(V.latest_by_stem(rows, labelers=("run_input",)))
    have_ai = set(V.latest_by_stem(rows, labelers=("ai",)))
    human = set(V.latest_by_stem(rows, labelers=V.CONFIRMED_LABELERS))  # 確定した回には AI の行を足さない (codex R2・依頼 185)
    targets = []
    for stem in V.list_run_stems(runs_dir):
        j = V.read_run_json(Path(runs_dir) / stem)
        if j and j.get("outcome") not in (None, ""):
            continue  # 印がある回は写しの対象
        if stem in marked or stem in have_ai or stem in human:
            continue
        if only and stem not in only:
            continue
        targets.append(stem)
    if limit:
        targets = targets[:limit]
    key = None if dry_run else api_key()
    counts = Counter()
    print(f"[AI   ] 対象 {len(targets)} 回 (モデル {vlm_judge.VLM_MODEL})")
    for i, stem in enumerate(targets, 1):
        rd = Path(runs_dir) / stem
        n_exec = V.executed_chunks(V.read_log(rd))
        evidence, imgs, info = material(rd)
        if not imgs:
            row = {"labeled_at": V.now_iso(), "run_stem": stem, "pick": "unknown", "place": "unknown",
                   "outcome": "not_judgeable", "labeler": "ai", "ai_outcome": "not_judgeable",
                   "ai_reason": f"[送ったチャンク {n_exec}] 素材が足りない (top.mp4 なし・frames/ {len(V.frame_files(rd))} 枚). API は呼んでいない",
                   "evidence": "", "note": ""}
        elif dry_run:
            print(f"  {i:3d} {stem} {evidence} {len(imgs)} 枚 送ったチャンク {n_exec}")
            counts["dry_run"] += 1
            continue
        elif not key:
            counts["skipped"] += 1
            print(f"  {i:3d} {stem} 飛ばした (API キーがない)")
            continue
        else:
            try:
                parsed, text = call_api(key, imgs, evidence)
            except ApiUnavailable as e:
                counts["skipped"] += 1
                print(f"  {i:3d} {stem} 飛ばした (API: {e})")
                continue
            if not isinstance(parsed, dict):
                counts["skipped"] += 1
                print(f"  {i:3d} {stem} 飛ばした (答えを読めない: {text[:80]!r})")
                continue
            row = ai_row(stem, parsed, evidence, n_exec)
        if dry_run:
            print(f"  {i:3d} {stem} not_judgeable (素材なし) 送ったチャンク {n_exec}")
            counts["not_judgeable"] += 1
            continue
        # 1 回ずつ書く (途中で止まっても済んだ回は残る). API を待つ間に人が確定していたら足さない
        if not V.append_rows([row], csv_path, skip_if_labeled_by=("ai",) + V.CONFIRMED_LABELERS,
                            forbid_dirs=(runs_dir,)):
            counts["skipped_human"] += 1
            print(f"  {i:3d} {stem} 足さなかった (待つ間に人か写しの行が足されていた)")
            continue
        counts[row["outcome"]] += 1
        print(f"  {i:3d} {stem} {row['outcome']:13s} pick={row['pick']} place={row['place']} {evidence or '-'}  {row['ai_reason'][:70]}")
    print("[AI   ] " + " / ".join(f"{k} {counts[k]}" for k in
                                  ("success", "failure", "unknown", "not_judgeable", "skipped", "skipped_human", "dry_run") if counts[k] or k not in ("dry_run", "skipped_human")))
    return counts


def do_rule(csv_path, runs_dir, dry_run):
    """エイイチの決め (2026-11-14 依頼 185) の行を足す. 送ったチャンクは log.jsonl の chunk の行で executed が True の数."""
    rows = V.read_rows(csv_path)
    conf = V.latest_by_stem(rows, labelers=V.CONFIRMED_LABELERS)
    ai = V.latest_by_stem(rows, labelers=("ai",))
    out, counts = [], Counter()
    for stem in V.list_run_stems(runs_dir):
        rd = Path(runs_dir) / stem
        n_exec = V.executed_chunks(V.read_log(rd))
        a = ai.get(stem) or {}
        hit2 = n_exec == 0
        hit3 = a.get("outcome") == "not_judgeable" and a.get("evidence") in ("frames", "")
        if stem in conf:
            if hit2 or hit3:
                counts[f"skip_{conf[stem]['labeler']}"] += 1
                print(f"  {stem} 足さない (確定の行あり: {conf[stem]['labeler']} {conf[stem]['outcome']}・"
                      f"送ったチャンク {n_exec})")
            continue
        r = V.rule_row(stem, n_exec, ai.get(stem))
        if r is None:
            continue
        counts["rule2" if hit2 else "rule3"] += 1
        counts["both"] += hit2 and hit3
        out.append(r)
    print(f"[RULE ] 足す行 {len(out)} (決め2 not_applicable {counts['rule2']} (うち 3 にも当たる {counts['both']}) / "
          f"決め3 not_judgeable {counts['rule3']})" + (" ※ dry-run なので書いていません" if dry_run else ""))
    if not dry_run and out:
        written = V.append_rows(out, csv_path, skip_if_labeled_by=V.CONFIRMED_LABELERS, forbid_dirs=(runs_dir,))
        if len(written) != len(out):
            print(f"[RULE ] 書くまでの間に確定の行が足されていた {len(out) - len(written)} 回は足さなかった")
        counts["written"] = len(written)
    return counts


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--run-input", action="store_true", help="{run_stem}.json の outcome を写す")
    ap.add_argument("--ai", action="store_true", help="印のない回を AI が判定する")
    ap.add_argument("--rule", action="store_true", help="エイイチの決め (依頼 185) の行を足す")
    ap.add_argument("--dry-run", action="store_true", help="書かない・API を呼ばない")
    ap.add_argument("--limit", type=int, default=0, help="AI の対象を先頭から N 回まで")
    ap.add_argument("--stems", nargs="*", default=None, help="AI の対象をこの run_stem だけにする")
    ap.add_argument("--runs-dir", default=str(V.RUNS_DIR))
    a = ap.parse_args()
    if not (a.run_input or a.ai or a.rule):
        ap.error("--run-input か --ai か --rule を付けてください")
    try:
        csv_path = V.VERDICTS_CSV  # 書き先は保存先テーブルの 1 か所だけ (codex R3. テストは関数に渡す)
        V.ensure_ssd_mounted(csv_path)
        if a.run_input:
            do_run_input(csv_path, a.runs_dir, a.dry_run)
        if a.ai:
            do_ai(csv_path, a.runs_dir, a.dry_run, a.limit, set(a.stems) if a.stems else None)
        if a.rule:
            do_rule(csv_path, a.runs_dir, a.dry_run)
    except (V.NotMounted, V.BrokenCsv) as e:
        print(f"[ABORT] {e}")
        return 2
    return 0


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