#!/usr/bin/env python3
"""推論の記録に成功/失敗を付ける判定の画面 (依頼 183).

  python3 pi05/verdict_viewer/server.py            # http://localhost:8765/ をブラウザで開く
  python3 pi05/verdict_viewer/server.py --no-open  # ブラウザを開かない

- 元の記録 (runs/pi05_policy/) は読むだけ. 判定は pi05_policy_verdicts_v1.csv に 1 行ずつ追記するだけ
- 書く前に KIOXIA のマウントを確かめ, なければ書かずに画面へエラーを返す
- localhost だけで待ち受ける
"""
import argparse
import csv
import json
import mimetypes
import re
import sys
import threading
import webbrowser
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import unquote, urlparse

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import verdicts as V  # noqa: E402

STEM_RE = re.compile(r"^pi05_[0-9_]+$")
FRAME_RE = re.compile(r"^[0-9A-Za-z_]+\.jpg$")
MAX_STATE_POINTS = 1500


def state_series(run_dir):
    """関節 6 本と爪の時系列. recording/states.csv があればそれ, なければ log.jsonl の state."""
    rd = Path(run_dir)
    p = V.safe_file(rd, "recording", "states.csv")
    t, joints, grip = [], [], []
    stages = []
    if p is not None:
        t0 = None
        with open(p, newline="") as f:
            for r in csv.DictReader(f):
                try:
                    ts = float(r["ts"])
                except (TypeError, ValueError):
                    continue
                t0 = ts if t0 is None else t0
                ev = r.get("event_type")
                if ev == "actual":
                    try:
                        joints.append([float(r[f"j{i}"]) for i in range(6)])
                    except (TypeError, ValueError):
                        continue
                    t.append(round(ts - t0, 3))
                elif ev == "gripper_cmd" and r.get("gripper_state") not in (None, ""):
                    try:
                        grip.append([round(ts - t0, 3), float(r["gripper_state"])])
                    except ValueError:
                        pass
                elif ev == "stage":
                    stages.append([round(ts - t0, 3), r.get("stage", "")])
        if len(t) > MAX_STATE_POINTS:
            k = len(t) / MAX_STATE_POINTS
            idx = sorted(set(int(i * k) for i in range(MAX_STATE_POINTS)) | {len(t) - 1})
            t = [t[i] for i in idx]
            joints = [joints[i] for i in idx]
        return {"source": "recording/states.csv", "t": t, "joints": joints, "grip": grip,
                "grip_label": "爪の命令 (0 開 / 1 閉)", "stages": stages, "t0": t0}
    t0 = None
    for x in V.read_log(rd):
        if "chunk" not in x or "state" not in x:
            continue
        try:
            st = json.loads(x["state"]) if isinstance(x["state"], str) else x["state"]
            ts = float(x["t"])
        except Exception:
            continue
        if not isinstance(st, list) or len(st) < 6:
            continue
        t0 = ts if t0 is None else t0
        t.append(round(ts - t0, 3))
        joints.append([float(v) for v in st[:6]])
        if len(st) >= 7:
            grip.append([round(ts - t0, 3), float(st[6])])
    return {"source": "log.jsonl の state (推論の前に読んだ値)", "t": t, "joints": joints, "grip": grip,
            "grip_label": "state の爪 (7 番目の値)", "stages": [], "t0": t0}


def video_times(run_dir, t0):
    """top.mp4 の各コマの時刻 (states の始まりからの秒). 再生位置のしるしに使う."""
    p = V.safe_file(run_dir, "recording", "top_frame_log.csv")
    if t0 is None or p is None:
        return None
    out = []
    try:
        with open(p, newline="") as f:
            for r in csv.DictReader(f):
                out.append(round(float(r["capture_ts"]) - t0, 3))
    except Exception:
        return None
    return out


class App:
    def __init__(self, runs_dir, csv_path):
        self.runs_dir = Path(runs_dir)
        self.csv_path = Path(csv_path)
        self.lock = threading.Lock()
        self.done_ops = {}  # 画面の 1 回の操作 (op_id) → 書いた行. 応答が届かず押し直しても 2 回書かない (codex R7)

    def runs_summary(self):
        rows = V.read_rows(self.csv_path)
        latest = V.current_by_stem(rows)  # 確定の行があればそれ (依頼 185)
        ai = V.latest_by_stem(rows, labelers=("ai",))
        human = V.latest_by_stem(rows, labelers=V.HUMAN_LABELERS)
        out = []
        for stem in V.list_run_stems(self.runs_dir):
            rd = self.runs_dir / stem
            lt, a, h = latest.get(stem), ai.get(stem), human.get(stem)
            out.append({
                "run_stem": stem,
                "has_mp4": V.has_top_mp4(rd),
                "n_frames": len(V.frame_files(rd)),
                "outcome": lt["outcome"] if lt else "",
                "labeler": lt["labeler"] if lt else "",
                "ai_outcome": a["outcome"] if a else "",
                "human_outcome": h["outcome"] if h else "",
                "human_at": h["labeled_at"] if h else "",  # 人が確定した時刻 (自分が確定した回の一覧. 依頼 187)
                "confirmed": bool(lt and lt["labeler"] in V.CONFIRMED_LABELERS),
            })
        return out

    def stats(self):
        """確定した回だけの成功率 (依頼 185). 記録のフォルダにある回だけ数える."""
        stems = set(V.list_run_stems(self.runs_dir))
        return V.success_rate([r for r in V.read_rows(self.csv_path) if r["run_stem"] in stems])

    def run_detail(self, stem):
        rd = self.runs_dir / stem
        rows = [r for r in V.read_rows(self.csv_path) if r["run_stem"] == stem]
        ai = [r for r in rows if r["labeler"] == "ai"]
        j = V.read_run_json(rd)
        log = V.read_log(rd)
        st = state_series(rd)
        return {
            "run_stem": stem,
            "has_mp4": V.has_top_mp4(rd),
            "frames": V.frame_files(rd),
            "states": st,
            "video_t": video_times(rd, st.get("t0")) if V.has_top_mp4(rd) else None,
            "ai": ai[-1] if ai else None,
            "latest": V.current_by_stem(rows).get(stem),
            "history": rows,
            "executed_chunks": V.executed_chunks(log),
            "end_reason": (j or {}).get("end_reason", ""),
            "model": next((x.get("model") for x in log if x.get("type") == "model"), "")
                     or ((j or {}).get("args") or {}).get("model", ""),
        }

    def add_verdict(self, body):
        stem = body.get("run_stem", "")
        if (not STEM_RE.match(stem) or not (self.runs_dir / stem).is_dir()
                or not V.within(self.runs_dir, self.runs_dir / stem)):
            raise ValueError(f"run_stem がありません: {stem!r}")
        outcome = body.get("outcome")
        if outcome not in V.OUTCOMES:
            raise ValueError(f"outcome が正しくありません: {outcome!r}")
        rd = self.runs_dir / stem
        op_id = str(body.get("op_id") or "")
        with self.lock:
            if op_id and op_id in self.done_ops:
                print(f"[SAVE ] {stem} 同じ操作の送り直し. 書いていない (op_id {op_id})")
                return self.done_ops[op_id]
            ai = V.latest_by_stem(V.read_rows(self.csv_path), labelers=("ai",)).get(stem)
            default = "yes" if outcome == "success" else "unknown"
            pick = body.get("pick") or default
            place = body.get("place") or default
            row = {"labeled_at": V.now_iso(), "run_stem": stem, "pick": pick, "place": place,
                   "outcome": outcome, "labeler": V.human_labeler_for(outcome, ai),
                   "ai_outcome": ai["outcome"] if ai else "", "ai_reason": ai["ai_reason"] if ai else "",
                   "evidence": "top.mp4" if V.has_top_mp4(rd) else ("frames" if V.frame_files(rd) else ""),
                   "note": str(body.get("note", "")).replace("\n", " ").strip()}
            written = V.append_rows([row], self.csv_path, forbid_dirs=(self.runs_dir,))[0]
            if op_id:
                self.done_ops[op_id] = written
            print(f"[SAVE ] {written['run_stem']} {written['outcome']} pick={written['pick']} "
                  f"place={written['place']} labeler={written['labeler']}")
            return written


def make_handler(app):
    class H(BaseHTTPRequestHandler):
        def log_message(self, fmt, *args):
            pass

        def _send(self, code, data, ctype="application/json; charset=utf-8", extra=None):
            if not isinstance(data, (bytes, bytearray)):
                data = json.dumps(data, ensure_ascii=False).encode("utf-8")
            self.send_response(code)
            self.send_header("Content-Type", ctype)
            self.send_header("Content-Length", str(len(data)))
            self.send_header("Cache-Control", "no-store")
            for k, v in (extra or {}).items():
                self.send_header(k, v)
            self.end_headers()
            if self.command != "HEAD":
                self.wfile.write(data)

        def _file(self, stem, *rel):
            # リンクでその回の記録の外へ出ない (codex R6・3 回目 R2)
            path = V.safe_file(app.runs_dir / stem, *rel) if V.within(app.runs_dir, app.runs_dir / stem) else None
            if path is None:
                return self._send(404, {"error": "not found"})
            size = path.stat().st_size
            ctype = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
            rng = self.headers.get("Range")
            start, end = 0, size - 1
            if rng:
                m = re.match(r"bytes=(\d*)-(\d*)", rng)
                if m:
                    if m.group(1):
                        start = int(m.group(1))
                        if m.group(2):
                            end = min(int(m.group(2)), size - 1)
                    elif m.group(2):
                        start = max(0, size - int(m.group(2)))
                if start > end or start >= size:
                    self.send_response(416)
                    self.send_header("Content-Range", f"bytes */{size}")
                    self.end_headers()
                    return
            length = end - start + 1
            self.send_response(206 if rng else 200)
            self.send_header("Content-Type", ctype)
            self.send_header("Accept-Ranges", "bytes")
            self.send_header("Content-Length", str(length))
            if rng:
                self.send_header("Content-Range", f"bytes {start}-{end}/{size}")
            self.end_headers()
            if self.command == "HEAD":
                return
            with open(path, "rb") as f:
                f.seek(start)
                left = length
                while left > 0:
                    chunk = f.read(min(1 << 20, left))
                    if not chunk:
                        break
                    try:
                        self.wfile.write(chunk)
                    except (BrokenPipeError, ConnectionResetError):
                        return
                    left -= len(chunk)

        def do_HEAD(self):
            self.do_GET()

        def do_GET(self):
            u = urlparse(self.path)
            parts = [unquote(p) for p in u.path.split("/") if p]
            try:
                if not parts:
                    return self._send(200, (HERE / "index.html").read_bytes(), "text/html; charset=utf-8")
                if parts == ["api", "runs"]:
                    return self._send(200, {"runs": app.runs_summary(), "stats": app.stats(), "csv": str(app.csv_path)})
                if (len(parts) == 3 and parts[:2] == ["api", "run"] and STEM_RE.match(parts[2])
                        and (app.runs_dir / parts[2]).is_dir() and V.within(app.runs_dir, app.runs_dir / parts[2])):
                    return self._send(200, app.run_detail(parts[2]))
                if len(parts) == 3 and parts[0] == "media" and STEM_RE.match(parts[1]) and parts[2] == "top.mp4":
                    return self._file(parts[1], "recording", "top.mp4")
                if (len(parts) == 4 and parts[0] == "media" and STEM_RE.match(parts[1]) and parts[2] == "frames"
                        and FRAME_RE.match(parts[3])):
                    return self._file(parts[1], "frames", parts[3])
                return self._send(404, {"error": "not found"})
            except Exception as e:  # 画面に理由を返す
                return self._send(500, {"error": f"{type(e).__name__}: {e}"})

        def do_POST(self):
            if urlparse(self.path).path != "/api/verdict":
                return self._send(404, {"error": "not found"})
            try:
                n = int(self.headers.get("Content-Length") or 0)
                body = json.loads(self.rfile.read(n).decode("utf-8") or "{}")
                row = app.add_verdict(body)
                return self._send(200, {"ok": True, "row": row})
            except (V.NotMounted, V.BrokenCsv) as e:
                print(f"[ABORT] {e}")
                return self._send(503, {"error": str(e)})
            except ValueError as e:
                return self._send(400, {"error": str(e)})
            except Exception as e:
                return self._send(500, {"error": f"{type(e).__name__}: {e}"})
    return H


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--port", type=int, default=8765)
    ap.add_argument("--no-open", action="store_true", help="ブラウザを開かない")
    ap.add_argument("--runs-dir", default=str(V.RUNS_DIR))
    ap.add_argument("--test-csv", default=None,
                    help="画面の確かめ用. 判定を保存先テーブルの CSV でなくこのファイルに書く (ふだんは付けない)")
    a = ap.parse_args()
    if not Path(a.runs_dir).is_dir():
        print(f"[ABORT] 記録のフォルダがありません ({a.runs_dir}). KIOXIA をつないでください")
        return 2
    csv_path = Path(a.test_csv) if a.test_csv else V.VERDICTS_CSV
    if a.test_csv and V.within(a.runs_dir, csv_path):
        print(f"[ABORT] --test-csv が記録のフォルダの中です ({csv_path}). 元の記録には書きません")
        return 2
    if a.test_csv:
        print(f"[TEST ] 確かめ用: 判定は {csv_path} に書きます (保存先テーブルの CSV には書きません)")
    app = App(a.runs_dir, csv_path)
    srv = ThreadingHTTPServer(("127.0.0.1", a.port), make_handler(app))
    url = f"http://localhost:{a.port}/"
    print(f"[VIEW ] {url}  (判定は {csv_path} に追記. Ctrl-C で終わる)")
    if not a.no_open:
        threading.Timer(0.5, lambda: webbrowser.open(url)).start()
    try:
        srv.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        srv.server_close()
    return 0


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