#!/usr/bin/env python3
from __future__ import annotations

import importlib.util
import json
from pathlib import Path
import struct
import sys
import time
from types import ModuleType
import unittest


class Dummy:
    pass


def install_module(name: str, **attributes: object) -> ModuleType:
    module = ModuleType(name)
    for key, value in attributes.items():
        setattr(module, key, value)
    sys.modules[name] = module
    return module


install_module("rclpy")
install_module("rclpy.node", Node=object)
for package, message_name in (
    ("brainco_hand_msgs.msg", "MotorStatus"),
    ("brainco_hand_msgs.msg", "SetMotorMulti"),
    ("diagnostic_msgs.msg", "DiagnosticStatus"),
    ("geometry_msgs.msg", "TwistStamped"),
    ("ros2_bridge_msgs.msg", "ArmStatus"),
    ("sensor_msgs.msg", "JointState"),
    ("std_srvs.srv", "Trigger"),
):
    module = sys.modules.get(package) or install_module(package)
    setattr(module, message_name, Dummy)


fake_omnisocket = install_module(
    "omnisocket", CONTROL_DEFAULTS={}, MSG_TYPE_BINARY=2
)

module_path = Path(__file__).with_name("tg3_local_teleop.py")
spec = importlib.util.spec_from_file_location("tg3_local_teleop", module_path)
assert spec is not None and spec.loader is not None
teleop = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = teleop
spec.loader.exec_module(teleop)


class FakeSession:
    mode = "idle"
    connect_count = 0
    sequence = 0
    sent_one = False
    next_instance_id = 0
    events: list[str] = []

    def __init__(self) -> None:
        type(self).next_instance_id += 1
        self.instance_id = type(self).next_instance_id

    def connect(self, **_kwargs: object) -> None:
        type(self).connect_count += 1
        type(self).events.append(f"connect:{self.instance_id}")

    def stats(self) -> dict[str, int]:
        return {"connected": 1, "registered": 1}

    def recv(self, timeout_ms: int) -> tuple[str, int, bytes] | None:
        if timeout_ms == 0:
            return None
        time.sleep(0.003)
        if self.mode == "idle":
            return None
        if self.mode == "one_then_idle" and type(self).sent_one:
            return None
        type(self).sent_one = True
        type(self).sequence += 1
        data = {
            "arm": {
                "position": {"left": [0.0] * 7, "right": [0.0] * 7}
            }
        }
        payload = json.dumps(data).encode()
        packet = struct.pack(
            "!4sQQI",
            b"TG3A",
            type(self).sequence,
            time.time_ns(),
            len(payload),
        ) + payload
        return "expected-sender", 2, packet

    def close(self) -> None:
        type(self).events.append(f"close:{self.instance_id}")
        return None


def config(refresh_s: float) -> dict[str, object]:
    return {
        "transport": "omnisocket",
        "omnisocket_server": "127.0.0.1:14049",
        "omnisocket_peer_id": "robot",
        "omnisocket_expected_sender": "expected-sender",
        "omnisocket_max_packet_age_ms": 300.0,
        "omnisocket_idle_session_refresh_s": refresh_s,
    }


class IdleSessionRefreshTest(unittest.TestCase):
    def setUp(self) -> None:
        FakeSession.connect_count = 0
        FakeSession.sequence = 0
        FakeSession.sent_one = False
        FakeSession.next_instance_id = 0
        FakeSession.events = []
        fake_omnisocket.Session = FakeSession

    def run_source(self, mode: str, run_s: float = 0.13) -> dict[str, object]:
        FakeSession.mode = mode
        source = teleop.LatestArmData(config(0.02))
        source.start()
        time.sleep(run_s)
        source.close()
        return source.metrics()

    def test_idle_session_is_periodically_reconnected(self) -> None:
        metrics = self.run_source("idle")
        self.assertGreaterEqual(FakeSession.connect_count, 2)
        self.assertGreaterEqual(metrics["idle_session_refreshes"], 1)
        self.assertEqual(
            FakeSession.events[:3],
            ["connect:1", "connect:2", "close:1"],
        )

    def test_continuous_valid_business_frames_prevent_refresh(self) -> None:
        metrics = self.run_source("active", 0.08)
        self.assertEqual(FakeSession.connect_count, 1)
        self.assertEqual(metrics["idle_session_refreshes"], 0)
        self.assertGreater(metrics["frames_accepted"], 1)

    def test_session_refreshes_after_last_valid_frame(self) -> None:
        metrics = self.run_source("one_then_idle")
        self.assertGreaterEqual(FakeSession.connect_count, 2)
        self.assertGreaterEqual(metrics["idle_session_refreshes"], 1)
        self.assertEqual(metrics["frames_accepted"], 1)

    def test_refresh_boundary_and_disable_switch(self) -> None:
        due = teleop.LatestArmData._idle_session_refresh_due
        self.assertFalse(due(10.0, 0.0, 0.0))
        self.assertFalse(due(1.999, 0.0, 2.0))
        self.assertTrue(due(2.0, 0.0, 2.0))


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