#!/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 回のやりとりごとに鍵をかける (読み書きが混ざらないように).

2026-10-31 依頼 46: 送っている間の角度の抜けについて
  daemon は send_angles を受けると, 送ったあとに角度を読み直す (動いている間は読めずにやり直す) まで
  返事をしない (約 1.6 秒. 爪は約 3.6 秒). daemon は 1 本のつながりを 1 件ずつ処理するので, その間は
  記録スレッドも角度を読めない. 腕はこの間に動き終わる (Log17 の記録で 1.5〜3.6 秒の抜け).
  記録の側で直せるのは次の 2 つだけ (送る指令・到着の判定・待ち時間は変えない):
    - 本体が読んだ角度 (到着の確かめ・送る直前の読み) も actual 行に残す (note=main)
    - daemon とのやりとり 1 回ごとの始まりと長さを残す (send_angles 行 / rpc 行). 送った目標と時刻が
      分かるので, あとから build_correction_episodes.py --gap-fill speed20 で抜けを埋められる
  抜けそのものをなくすには daemon の送ったあとの読み直しを変える必要がある (決めるのはエイイチ).
"""
import csv
import json
import math
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            # 本体が画像を求めたとき, 新しいコマを待つ上限


_rpc_observers = []             # daemon とのやりとり 1 回ごとに呼ぶ関数 (記録用. 依頼 46)


def add_rpc_observer(fn):
    if fn not in _rpc_observers:
        _rpc_observers.append(fn)


def remove_rpc_observer(fn):
    if fn in _rpc_observers:
        _rpc_observers.remove(fn)


def _notify_rpc(cmd_obj, resp, t0, t1, err):
    for fn in list(_rpc_observers):
        try:
            fn(cmd_obj, resp, t0, t1, err)
        except Exception:
            pass                # 記録の失敗で本体を止めない


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 回だけ差し替える.
    依頼 46: やりとりが終わったら (鍵を外してから) add_rpc_observer で足した関数に
    (送った dict, 返事 or None, 始めた時刻, 終わった時刻, 例外 or None) を渡す. 返事と例外はそのまま返す.
    再接続のときの入れ子の ping (外のやりとりの鍵の中) は渡さない (codex 推奨).
    観察は同じスレッドで呼ぶ (states.csv に 1 行書くだけ. 本体へ返るのがその分だけ遅れる).
    """
    if getattr(vsp._daemon_request, "_policy_locked", False):
        return
    orig = vsp._daemon_request
    lock = threading.RLock()
    depth = threading.local()       # 再接続の中の ping (鍵の中で呼ばれる入れ子) は観察に渡さない

    def _locked(*a, **k):
        cmd_obj = a[0] if a else k.get("cmd_obj")
        resp, err, t0, t1 = None, None, None, None
        nested = getattr(depth, "n", 0) > 0
        try:
            depth.n = getattr(depth, "n", 0) + 1
            with lock:
                t0 = time.time()
                try:
                    resp = orig(*a, **k)
                finally:
                    t1 = time.time()
        except BaseException as e:
            err = e
            raise
        finally:
            depth.n -= 1
            if _rpc_observers and t0 is not None and not nested:
                _notify_rpc(cmd_obj, resp, t0, t1, err)
        return resp
    _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, angles_ok=None):
        self.dir = Path(run_dir) / "recording"
        self.angles_ok = angles_ok      # 本体が読んだ角度を残すときの確かめ (None なら 6 個の有限な数)
        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.n_actual_main = 0          # 本体が読んだ角度を残した数 (依頼 46)
        self.n_sends = 0                # send_angles のやりとりを残した数
        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()
        add_rpc_observer(self._on_rpc)

    # ---- 本体から呼ぶ -------------------------------------------------------
    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
        remove_rpc_observer(self._on_rpc)
        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,
                "actual_main_samples": self.n_actual_main, "send_rpcs": self.n_sends,
                "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, speed=None, ts_send=None, note=""):
        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 speed is None else str(speed), "" if gripper_state is None else str(gripper_state),
                "" if ts_send is None else f"{ts_send:.6f}", "", "", "", note]
        with self.csv_lock:
            if self.states_f.closed:
                return False
            self.states_w.writerow(row)
            self._maybe_flush(ts)
        return True

    def _valid(self, a):
        if self.angles_ok is not None:
            return bool(self.angles_ok(a))
        try:
            return len(a) == 6 and all(math.isfinite(float(v)) for v in a)
        except Exception:
            return False

    def _on_rpc(self, cmd_obj, resp, t0, t1, err):
        """daemon とのやりとり 1 回ごと (install_daemon_lock から). 何も送らない・読み直さない.

        get_angles: 本体 (記録スレッド以外) が読めた角度を actual 行に残す (note=main. 記録スレッドの分は
                    _angle_loop が残すので二重にしない)
        send_angles: 始めた時刻 (ts / ts_send)・目標・速度・やりとりの長さ (note) を send_angles 行に残す
        そのほか (ping 以外): rpc 行に cmd とやりとりの長さを残す
        """
        if self.stop_flag or not isinstance(cmd_obj, dict):
            return
        cmd = cmd_obj.get("cmd")
        dur = f"rpc_sec={t1 - t0:.3f}" + ("" if err is None else f" error={type(err).__name__}")
        if cmd == "get_angles":
            if threading.current_thread() is self.t_ang or err is not None or not isinstance(resp, dict):
                return
            a = resp.get("angles")
            if a is None or not self._valid(a):
                return
            if self._state_row("actual", t1, angles=a, note="main"):   # 書けた行だけ数える (codex 推奨)
                self.n_actual_main += 1
        elif cmd == "send_angles":
            ang = cmd_obj.get("angles")
            if self._state_row("send_angles", t0, angles=ang if ang is not None and len(ang) == 6 else None,
                               speed=cmd_obj.get("speed"), ts_send=t0, note=dur):
                self.n_sends += 1
        elif cmd != "ping":
            self._state_row("rpc", t0, ts_send=t0, note=f"cmd={cmd} {dur}")

    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()       # 遅れた分は取り戻さない (詰めて読まない)
