#!/usr/bin/env python3
"""test_pad_correction.py — run_policy_client.py の p (ゲームパッドで補正. 2026-11-05 LOG450) の確かめ.

本物の pygame のループ (SDL_VIDEODRIVER=dummy) に偽のコントローラーの押しを流し, 補正の結果の記録・
爪の記録・巻き戻し用の目標・ドライランの仮の姿勢・実機の送り口 (daemon を偽物にする) を見る.
  SDL_VIDEODRIVER=dummy SDL_AUDIODRIVER=dummy python3 pi05/test_pad_correction.py
"""
import functools
import io
import json
import os
import sys
import unittest
from pathlib import Path
from unittest import mock

os.environ.setdefault("SDL_VIDEODRIVER", "dummy")
os.environ.setdefault("SDL_AUDIODRIVER", "dummy")
_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE))
sys.path.insert(0, str(_HERE.parent / "tools"))

import numpy as np  # noqa: E402

import run_policy_client as rpc  # noqa: E402
import gamepad_teleop as gt  # noqa: E402


class FakeJoystick:
    """test_gamepad_teleop.FakeJoystick と同じ (あちらは読み込むと刻みを 0.5° に書き換えるので写した)."""

    def __init__(self, idx):
        self.axes = [0.0, 0.0, 0.0, 0.0, -1.0, -1.0]

    def get_instance_id(self):
        return 7

    def get_name(self):
        return "PS4 Controller"

    def get_guid(self):
        return "fake-guid"

    def get_numaxes(self):
        return 6

    def get_numbuttons(self):
        return 16

    def get_numhats(self):
        return 0

    def get_axis(self, a):
        return self.axes[a]

    def get_button(self, b):
        return 0

HOME = [0.0, 0.0, 0.0, 0.0, 0.0, -30.0]


class FakeRecorder:
    def __init__(self):
        self.events = []

    def set_mode(self, m):
        self.events.append(("mode", m))

    def gripper(self, s):
        self.events.append(("gripper", s))


def script_frames(script):
    import pygame

    def btn(b, down=True):
        t = pygame.JOYBUTTONDOWN if down else pygame.JOYBUTTONUP
        pygame.event.post(pygame.event.Event(t, instance_id=7, button=b, joy=0))

    def on_frame(n, tp):
        if n == 1:
            pygame.event.post(pygame.event.Event(pygame.JOYDEVICEADDED, device_index=0))
        for b in script.get(n, ()):
            btn(b); btn(b, False)
        if n == 600:
            pygame.event.post(pygame.event.Event(pygame.QUIT))
    return on_frame


LO = np.array([-60.0, -10, -10, -60, -60, -60])
HI = np.array([60.0, 60, 60, 60, 60, 0])


class TestPadDry(unittest.TestCase):
    def test_dry_flow(self):
        arm = rpc.ArmIO(False, None, sim_q=HOME)
        rec = FakeRecorder()
        log_f = io.StringIO()
        on_frame = script_frames({3: [gt.BTN_RIGHT], 40: [gt.BTN_RIGHT], 80: [gt.BTN_CROSS], 140: [gt.BTN_SQUARE]})
        orig = gt.run_embedded
        with mock.patch.object(gt, "run_embedded", functools.partial(orig, on_frame=on_frame)), \
                mock.patch.object(gt.pygame.joystick, "Joystick", FakeJoystick), \
                mock.patch.object(gt, "GRIPPER_SLEEP", 0.0):
            res, grip = rpc.run_pad_correction(arm, 0, 1, 4, log_f, rec, LO, HI)
        self.assertEqual(res["action"], "accept")
        self.assertEqual(res["device"], "gamepad")
        self.assertEqual(res["pad_end_reason"], "finish_button")
        self.assertEqual(res["steps_sent"], 2)
        self.assertEqual(res["sent_targets"][-1][0], 2.0)
        self.assertTrue(res["grip_used"])
        self.assertEqual((res["grip_before"], res["grip_after"], grip), (0, 1, 1))
        self.assertAlmostEqual(arm.sim_q[0], 2.0)                    # 推論に戻る姿勢 = 直した姿勢
        self.assertEqual(rec.events[0], ("mode", rpc.MODE_CORRECT))
        self.assertEqual(rec.events[-1], ("mode", rpc.MODE_MODEL))
        rows = [json.loads(x) for x in log_f.getvalue().splitlines()]
        self.assertTrue(rows and all(r["type"] == "pad" and r["correction_id"] == 1 and r["after_chunk"] == 4
                                     for r in rows))
        self.assertGreaterEqual(res["duration_s"], 0.0)
        self.assertNotIn("outside_envelope", res)

    def test_outside_envelope_warns(self):
        arm = rpc.ArmIO(False, None, sim_q=[59.5, 0, 0, 0, 0, -30])
        on_frame = script_frames({3: [gt.BTN_RIGHT], 60: [gt.BTN_RIGHT], 120: [gt.BTN_SQUARE]})
        orig = gt.run_embedded
        out = []
        with mock.patch.object(gt, "run_embedded", functools.partial(orig, on_frame=on_frame)), \
                mock.patch.object(gt.pygame.joystick, "Joystick", FakeJoystick), \
                mock.patch("builtins.print", lambda *a, **k: out.append(" ".join(map(str, a)))):
            res, _ = rpc.run_pad_correction(arm, 0, 2, 0, io.StringIO(), None, LO, HI)
        self.assertEqual(res["outside_envelope"], [0])
        self.assertTrue(any("学習データの範囲の外" in x for x in out))


class TestPadExecute(unittest.TestCase):
    """実機の送り口: daemon を偽物にして, 1 関節の send_angle と, 爪の指令の瞬間の記録を見る."""

    def test_execute_sends_and_records_gripper(self):
        q = {"now": np.array(HOME, dtype=float)}
        calls = []

        def req(c, **k):
            calls.append(c)
            if c["cmd"] == "send_angle":
                q["now"][c["joint_id"] - 1] = c["angle"]
                return {"ok": True}
            if c["cmd"] == "get_angles":
                return {"ok": True, "angles": list(q["now"])}
            return {"ok": True, "result": None}

        rec = FakeRecorder()
        arm = rpc.ArmIO(True, rec)
        on_frame = script_frames({3: [gt.BTN_LEFT], 60: [gt.BTN_CROSS], 120: [gt.BTN_SQUARE]})
        orig = gt.run_embedded
        gt.load_daemon_modules()
        self.assertIs(gt.rpc, sys.modules["run_policy_client"])
        with mock.patch.object(gt, "run_embedded", functools.partial(orig, on_frame=on_frame)), \
                mock.patch.object(gt.pygame.joystick, "Joystick", FakeJoystick), \
                mock.patch.object(rpc, "read_angles", lambda: q["now"].copy()), \
                mock.patch.object(gt, "daemon_request", req), \
                mock.patch.object(rpc, "daemon_set_gripper_value", lambda v: calls.append({"cmd": "grip", "v": v})), \
                mock.patch.object(rpc, "GRIPPER_SLEEP", 0.0):
            res, grip = rpc.run_pad_correction(arm, 0, 3, 7, io.StringIO(), rec, LO, HI)
        self.assertEqual(res["pad_end_reason"], "finish_button")
        self.assertIn({"cmd": "send_angle", "joint_id": 1, "angle": -1.0, "speed": gt.SPEED, "wait": rpc.SEND_WAIT},
                      calls)
        self.assertIn({"cmd": "grip", "v": rpc.GRIP_VALUE}, calls)
        self.assertIn(("gripper", 1), rec.events)
        self.assertEqual(grip, 1)
        self.assertEqual(res["sent_targets"], [[-1.0, 0.0, 0.0, 0.0, 0.0, -30.0]])


class TestConfirmP(unittest.TestCase):
    """確認で p を押すと補正に入り, 終えたあとの b でその補正を巻き戻せる (アームなし --dry-keys)."""

    def test_p_then_b(self):
        import tempfile
        import time
        from test_rewind_correct import run_client

        def fake_pad(arm, grip, corr_id, chunk, log_f, recorder, lo, hi):
            start = arm.read()[:6].copy()
            t0 = time.time()
            tgt = start.copy(); tgt[0] += 3.0
            arm.sim_q = tgt.copy()
            return ({"type": "correction", "action": "accept", "device": "gamepad", "correction_id": corr_id,
                     "after_chunk": chunk, "t_start": t0, "t_end": time.time(), "start_pose": rpc._fmt6(start),
                     "final_target": rpc._fmt6(tgt), "measured_after": rpc._fmt6(tgt), "steps_sent": 1,
                     "grip_before": grip, "grip_after": grip, "grip_used": False, "regrip_used": False,
                     "sent_targets": [rpc._fmt6(tgt)]}, grip)
        td = Path(tempfile.mkdtemp())
        with mock.patch.object(rpc, "run_pad_correction", fake_pad):
            code, world, run, lines, out = run_client(
                rpc, ["--dry-keys", "--steps-per-chunk", "10", "--max-chunks", "4"], ["", "p", "b", "q"], td)
        self.assertEqual(code, 0)
        req = [x for x in lines if x.get("correction_requested")]
        self.assertEqual([x["correction_device"] for x in req], ["gamepad"])
        rw = [x for x in lines if x.get("type") == "rewind"]
        self.assertEqual([(x["kind"], x["correction_id"]) for x in rw], [("correction", 1)])
        self.assertIn("p=ゲームパッドで補正", rpc.CONFIRM_PROMPT)


if __name__ == "__main__":
    unittest.main()
