#!/usr/bin/env python3
"""policy_recorder.py — 実機推論 (run_policy_client.py --execute) の間, 上カメラと実測角を 1 秒 10 コマで残す (2026-10-29).

お手本 (record_teach_demo.py + recorder.py --top) と同じ形で書くので,
build_demo_episodes.build_one() がそのまま読める:
  {run_dir}/recording/states.csv          recorder.py と同じ列. actual / gripper_cmd / stage 行.
                                          stage 列 = その時の状態 (model / rewind / correct / home / wait)
  {run_dir}/recording/top.mp4             上カメラ (元の大きさ. 10 コマ/秒)
  {run_dir}/recording/top_frame_log.csv   frame_idx, capture_ts (+ mode 列. 1 コマごとの状態)
  {run_dir}/recording/top_meta.json       コマ数・実効 fps など

アームへは何も送らない (関節角を読むだけ). daemon のソケットは 1 本を本体と共用するので,
install_daemon_lock() で 1 回のやりとりごとに鍵をかける (読み書きが混ざらないように).
"""
import csv
import json
import threading
import time
from pathlib import Path

import cv2

import visual_servo_pick as vsp

RECORD_FPS = 10.0
FRAME_INTERVAL = 1.0 / RECORD_FPS
ANGLE_INTERVAL = 1.0 / RECORD_FPS
FOURCC = "avc1"                 # recorder.py と同じ
STATES_COLS = ["ts", "event_type", "stage",
               "j0", "j1", "j2", "j3", "j4", "j5",
               "speed", "gripper_state",
               "ts_send", "ts_arrival", "attempt", "arrived",
               "note"]          # recorder.py の STATES_COLS と同じ
FRAME_LOG_COLS = ["frame_idx", "capture_ts", "mode"]
FLUSH_INTERVAL = 1.0
STOP_JOIN_SEC = 5.0             # stop() でスレッドの終わりを待つ上限
FRAME_WAIT_SEC = 2.0            # 本体が画像を求めたとき, 新しいコマを待つ上限


def install_daemon_lock():
    """visual_servo_pick._daemon_request を「1 回のやりとりごとに鍵をかける」版に差し替える.

    daemon_get_angles / daemon_send_angles / daemon_set_gripper* はどれも
    モジュールの _daemon_request を呼ぶので, ここを差し替えれば全部に効く.
    記録スレッドの get_angles と本体の send_angles が同じソケットで混ざらないようにする.
    中身 (送る内容・タイムアウト) は変えない. 2 回呼んでも 1 回だけ差し替える.
    """
    if getattr(vsp._daemon_request, "_policy_locked", False):
        return
    orig = vsp._daemon_request
    lock = threading.RLock()

    def _locked(*a, **k):
        with lock:
            return orig(*a, **k)
    _locked._policy_locked = True
    _locked._orig = orig
    vsp._daemon_request = _locked


class PolicyRecorder:
    """カメラを読むスレッドと関節角を読むスレッドを持つ. 本体は grab() で最新のコマを受け取る."""

    def __init__(self, run_dir, cap, read_angles_fn, to_model_fn):
        self.dir = Path(run_dir) / "recording"
        self.dir.mkdir(parents=True, exist_ok=True)
        self.cap = cap
        self.read_angles_fn = read_angles_fn
        self.to_model_fn = to_model_fn
        self.mode = "model"
        self.stop_flag = False
        self.stopped_cleanly = None     # stop() のあと: スレッドが止まって書き切れたか
        self.csv_lock = threading.Lock()
        self.frame_cv = threading.Condition()
        self.latest = None              # (frame, capture_ts)
        self.writer = None
        self.size = None
        self.n_frames = 0
        self.n_actual = 0
        self.n_actual_fail = 0
        self.write_errors = 0           # mp4 を開けない / 書けなかった回数 (codex R2 A5)
        self.started = time.time()
        self.last_flush = self.started
        self.states_f = (self.dir / "states.csv").open("w", newline="")
        self.states_w = csv.writer(self.states_f)
        self.states_w.writerow(STATES_COLS)
        self.flog_f = (self.dir / "top_frame_log.csv").open("w", newline="")
        self.flog_w = csv.writer(self.flog_f)
        self.flog_w.writerow(FRAME_LOG_COLS)
        self._state_row("stage", time.time())
        self.t_cam = threading.Thread(target=self._cam_loop, daemon=True, name="record_camera")
        self.t_ang = threading.Thread(target=self._angle_loop, daemon=True, name="record_angles")

    def start(self):
        self.t_cam.start()
        self.t_ang.start()

    # ---- 本体から呼ぶ -------------------------------------------------------
    def set_mode(self, mode):
        if mode == self.mode:
            return
        self.mode = mode
        self._state_row("stage", time.time())

    def gripper(self, state):
        """グリッパーに指令した瞬間に呼ぶ (build_demo_episodes は指令の時刻で 0/1 を切り替える)."""
        self._state_row("gripper_cmd", time.time(), gripper_state=int(state))

    def grab(self):
        """呼んだあとに撮れたコマを返す (古いコマを渡さない). 戻り値は grab_rgb と同じ (rgb, bgr_small)."""
        t_req = time.time()
        deadline = t_req + FRAME_WAIT_SEC
        with self.frame_cv:
            while not self.stop_flag:
                if self.latest is not None and self.latest[1] >= t_req:
                    return self.to_model_fn(self.latest[0])
                left = deadline - time.time()
                if left <= 0:
                    break
                self.frame_cv.wait(left)
        return None, None

    def stop(self):
        """スレッドを止めて書き切る. 何度呼んでもよい."""
        if self.stop_flag:
            return
        self.stop_flag = True
        with self.frame_cv:
            self.frame_cv.notify_all()
        for t in (self.t_cam, self.t_ang):
            if t.is_alive():
                t.join(timeout=STOP_JOIN_SEC)
        ended = time.time()
        alive = [t.name for t in (self.t_cam, self.t_ang) if t.is_alive()]
        # codex R1-5: スレッドが止まっていないのに mp4 を閉じると, あとから write() が走って競合する.
        # 止まらなかったときは mp4 を閉じず (カメラも本体で閉じない: alive() を見る),
        # 記録を「異常」として残す (build_correction_episodes はこの run を使わない).
        self.stopped_cleanly = not alive and not self.write_errors
        with self.csv_lock:
            for f in (self.states_f, self.flog_f):
                try:
                    f.flush(); f.close()
                except Exception:
                    pass
        if self.writer is not None and not alive:
            try:
                self.writer.release()
            except Exception:
                pass
        el = ended - self.started
        meta = {"frames": self.n_frames, "elapsed_sec": round(el, 3),
                "effective_fps": round(self.n_frames / el, 3) if el > 0 else 0.0,
                "width": self.size[0] if self.size else None,
                "height": self.size[1] if self.size else None,
                "actual_samples": self.n_actual, "actual_read_failures": self.n_actual_fail,
                "actual_hz": round(self.n_actual / el, 2) if el > 0 else 0.0,
                "target_fps": RECORD_FPS, "status": ("thread_not_stopped" if alive else "write_error" if self.write_errors else "ok"),
                "write_errors": self.write_errors,
                "threads_alive": alive}
        (self.dir / "top_meta.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False))
        if alive:
            print(f"[RECORD] [WARN] 記録スレッドが止まりません ({', '.join(alive)}). "
                  "mp4 を閉じずに終えます (この記録は学習データに使いません)")
        print(f"[RECORD] 記録終了 コマ {self.n_frames} ({meta['effective_fps']} fps) "
              f"実測角 {self.n_actual} 回 ({meta['actual_hz']} Hz)  → {self.dir}")

    # ---- 中身 ---------------------------------------------------------------
    def alive(self):
        """記録スレッドがまだ動いているか (stop() のあとに見る. 動いていればカメラを閉じない)."""
        return self.t_cam.is_alive() or self.t_ang.is_alive()

    def _state_row(self, ev, ts, angles=None, gripper_state=None):
        row = [f"{ts:.6f}", ev, self.mode]
        for k in range(6):
            row.append("" if angles is None else f"{float(angles[k]):.3f}")
        row += ["", "" if gripper_state is None else str(gripper_state), "", "", "", "", ""]
        with self.csv_lock:
            if self.states_f.closed:
                return
            self.states_w.writerow(row)
            self._maybe_flush(ts)

    def _maybe_flush(self, now):
        if now - self.last_flush >= FLUSH_INTERVAL:
            try:
                self.states_f.flush(); self.flog_f.flush()
            except Exception:
                pass
            self.last_flush = now

    def _cam_loop(self):
        last_written = 0.0
        while not self.stop_flag:
            ok, frame = self.cap.read()
            ts = time.time()
            if self.stop_flag:
                break
            if not ok or frame is None:
                time.sleep(0.02)
                continue
            with self.frame_cv:
                self.latest = (frame, ts)
                self.frame_cv.notify_all()
            if ts - last_written < FRAME_INTERVAL - 1e-3:
                continue
            last_written = ts
            try:
                if self.writer is None:
                    h, w = frame.shape[:2]
                    self.size = (w, h)
                    self.writer = cv2.VideoWriter(str(self.dir / "top.mp4"),
                                                  cv2.VideoWriter_fourcc(*FOURCC), RECORD_FPS, (w, h))
                    if not self.writer.isOpened():
                        raise RuntimeError(f"VideoWriter を開けない ({FOURCC})")
                if not self.writer.isOpened():
                    raise RuntimeError("VideoWriter が閉じている")
                self.writer.write(frame)
                with self.csv_lock:
                    if not self.flog_f.closed:
                        self.flog_w.writerow([self.n_frames, f"{ts:.6f}", self.mode])
                        self._maybe_flush(ts)
                self.n_frames += 1
            except Exception as e:
                self.write_errors += 1
                if self.write_errors <= 3:
                    print(f"[RECORD] [WARN] コマの書き込みに失敗: {e}")

    def _angle_loop(self):
        nxt = time.time()
        while not self.stop_flag:
            try:
                a = self.read_angles_fn()
            except Exception:
                a = None
            ts = time.time()
            if a is None:
                self.n_actual_fail += 1
            else:
                self._state_row("actual", ts, angles=a)
                self.n_actual += 1
            nxt += ANGLE_INTERVAL
            wait = nxt - time.time()
            if wait > 0:
                time.sleep(wait)
            else:
                nxt = time.time()       # 遅れた分は取り戻さない (詰めて読まない)
