import sys
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
from lerobot.datasets.lerobot_dataset import LeRobotDataset

from lerobot_robot_ufactory.robots.uf_robot.uf_robot_config import UFRobotConfig
from lerobot_robot_ufactory.scripts import uf_lerobot_record as record_module
from lerobot_robot_ufactory.scripts.uf_lerobot_record import (
    _manual_action_from_observation,
    _update_manual_gripper_key_state,
    _update_manual_gripper_target,
    _prepare_dataset_root,
    _prepare_recording_episode,
    get_cfg,
)


class FakeXArm:
    def __init__(self, robot_ip):
        self.robot_ip = robot_ip
        self.connected = True
        self.axis = 6
        self.error_code = 0
        self.mode = 0
        self.initial_point = [0.0, -30.0, 0.0, 0.0, 0.0, 30.0]
        self._arm = type("FakeArmTransport", (), {"_baud_checkset": False})()
        self.gripper_position = 800
        self.calls = []

    def motion_enable(self, **kwargs):
        self.calls.append(("motion_enable", kwargs))

    def clean_error(self):
        self.calls.append(("clean_error",))

    def set_teach_sensitivity(self, sensitivity):
        self.calls.append(("set_teach_sensitivity", sensitivity))
        return 0

    def set_mode(self, mode):
        self.calls.append(("set_mode", mode))
        self.mode = mode
        return 0

    def set_state(self, state):
        self.calls.append(("set_state", state))
        return 0

    def get_initial_point(self):
        self.calls.append(("get_initial_point",))
        return 0, self.initial_point

    def set_servo_angle(self, **kwargs):
        self.calls.append(("set_servo_angle", kwargs))
        return 0

    def get_err_warn_code(self):
        return 0, [0, 0]

    def set_linear_spd_limit_factor(self, factor):
        self.calls.append(("set_linear_spd_limit_factor", factor))
        return 0

    def set_gripper_enable(self, enable):
        self.calls.append(("set_gripper_enable", enable))
        return 0

    def set_gripper_mode(self, mode):
        self.calls.append(("set_gripper_mode", mode))
        return 0

    def set_gripper_speed(self, speed):
        self.calls.append(("set_gripper_speed", speed))
        return 0

    def set_gripper_position(self, position, **kwargs):
        self.calls.append(("set_gripper_position", position, kwargs))
        self.gripper_position = position
        return 0

    def get_gripper_position(self):
        self.calls.append(("get_gripper_position",))
        return 0, self.gripper_position

    def getset_tgpio_modbus_data(self, data):
        self.calls.append(("getset_tgpio_modbus_data", data))
        return 0, []

    def get_joint_states(self, is_radian=True, num=3):
        positions = np.arange(6, dtype=np.float64)
        velocities = np.zeros(6, dtype=np.float64)
        return 0, [positions, velocities, velocities]

    def disconnect(self):
        self.calls.append(("disconnect",))
        self.connected = False


def test_manual_mode_robot_enters_teaching_mode_without_sending_actions(monkeypatch, tmp_path):
    from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module

    arm = FakeXArm("192.168.1.245")
    monkeypatch.setattr(uf_robot_module, "XArmAPI", lambda robot_ip: arm)
    monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None)

    config = UFRobotConfig(
        id="test_manual_robot",
        calibration_dir=tmp_path,
        robot_ip=arm.robot_ip,
        robot_dof=6,
        control_space="joint",
        gripper_type=0,
        manual_mode=True,
        teach_sensitivity=4,
    )
    robot = uf_robot_module.UFRobot(config)

    robot.connect()
    assert robot.is_connected
    assert arm.mode == 2
    assert ("set_teach_sensitivity", 4) in arm.calls
    assert robot._initial_point == arm.initial_point
    assert not any(call[0] == "set_servo_angle" for call in arm.calls)

    robot.reset_to_initial()
    reset_calls = [call for call in arm.calls if call[0] == "set_servo_angle"]
    assert reset_calls == [
        (
            "set_servo_angle",
            {
                "angle": arm.initial_point,
                "speed": 20,
                "is_radian": False,
                "wait": True,
            },
        )
    ]
    assert arm.mode == 2

    action = {"J1.pos": 1.0}
    assert robot.send_action(action) is action

    observation = robot.get_observation()
    assert observation["J1.pos"] == 0.0
    assert observation["J6.pos"] == 5.0

    robot.disconnect()
    assert not robot.is_connected
    assert arm.mode == 0
    assert ("disconnect",) in arm.calls

    call_count = len(arm.calls)
    robot.disconnect()
    assert len(arm.calls) == call_count


def test_robot_reset_uses_sdk_initial_point_in_normal_mode(monkeypatch, tmp_path):
    from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module

    arm = FakeXArm("192.168.1.245")
    monkeypatch.setattr(uf_robot_module, "XArmAPI", lambda robot_ip: arm)
    monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None)

    config = UFRobotConfig(
        id="test_normal_robot",
        calibration_dir=tmp_path,
        robot_ip=arm.robot_ip,
        robot_dof=6,
        control_space="joint",
        gripper_type=0,
    )
    robot = uf_robot_module.UFRobot(config)
    assert not hasattr(config, "start_joints")
    assert not hasattr(config, "start_tcp_pose")
    robot.connect()

    reset_calls = [call for call in arm.calls if call[0] == "set_servo_angle"]
    assert reset_calls == [
        (
            "set_servo_angle",
            {
                "angle": arm.initial_point,
                "speed": 20,
                "is_radian": False,
                "wait": True,
            },
        )
    ]

    robot.disconnect()


def test_normal_mode_waits_for_gripper_to_open_before_control(monkeypatch, tmp_path):
    from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module

    arm = FakeXArm("192.168.1.245")
    arm.gripper_position = 400
    monkeypatch.setattr(uf_robot_module, "XArmAPI", lambda robot_ip: arm)
    monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None)

    config = UFRobotConfig(
        id="test_wait_for_gripper",
        calibration_dir=tmp_path,
        robot_ip=arm.robot_ip,
        robot_dof=6,
        control_space="joint",
        gripper_type=1,
    )
    robot = uf_robot_module.UFRobot(config)
    robot.connect()

    open_calls = [call for call in arm.calls if call[0] == "set_gripper_position"]
    assert open_calls == [("set_gripper_position", 800, {"wait": True})]
    assert robot._last_gripper_command == 0.0

    before_writes = len(
        [call for call in arm.calls if call[0] == "getset_tgpio_modbus_data"]
    )
    robot._send_gripper_action(0.0)
    after_writes = len(
        [call for call in arm.calls if call[0] == "getset_tgpio_modbus_data"]
    )
    assert after_writes == before_writes

    robot.disconnect()


def test_manual_mode_config_rejects_cartesian_control(tmp_path):
    with pytest.raises(ValueError, match="control_space='joint'"):
        UFRobotConfig(
            id="test_manual_robot",
            calibration_dir=tmp_path,
            robot_dof=6,
            control_space="cartesian",
            manual_mode=True,
        )


def test_manual_gripper_speed_is_configurable_and_non_negative(tmp_path):
    config = UFRobotConfig(
        id="test_manual_robot",
        calibration_dir=tmp_path,
        robot_dof=6,
        manual_mode=True,
        manual_gripper_speed=0.25,
    )
    assert config.manual_gripper_speed == 0.25

    with pytest.raises(ValueError, match="manual_gripper_speed"):
        UFRobotConfig(
            id="test_manual_robot",
            calibration_dir=tmp_path,
            robot_dof=6,
            manual_mode=True,
            manual_gripper_speed=-0.1,
        )


def test_manual_action_filters_non_action_observation_fields():
    observation = {
        "J1.pos": 1.0,
        "J1.vel": 2.0,
        "gripper.pos": 0.5,
        "camera": np.zeros((2, 2, 3), dtype=np.uint8),
    }
    action_features = {"J1.pos": float, "gripper.pos": float}

    assert _manual_action_from_observation(observation, action_features) == {
        "J1.pos": 1.0,
        "gripper.pos": 0.5,
    }


def test_manual_gripper_keys_update_target_in_expected_direction_and_bounds():
    key_state = {"close": False, "open": False}

    _update_manual_gripper_key_state(type("Key", (), {"char": "C"})(), True, key_state)
    assert key_state == {"close": True, "open": False}
    assert _update_manual_gripper_target(0.5, key_state, speed=1.0, fps=10) == pytest.approx(0.6)

    _update_manual_gripper_key_state(type("Key", (), {"char": "C"})(), False, key_state)
    _update_manual_gripper_key_state(type("Key", (), {"char": "o"})(), True, key_state)
    assert _update_manual_gripper_target(0.05, key_state, speed=1.0, fps=10) == 0.0

    _update_manual_gripper_key_state(type("Key", (), {"char": "c"})(), True, key_state)
    assert _update_manual_gripper_target(0.99, key_state, speed=1.0, fps=10) == 0.99


def test_manual_mode_initializes_gripper_without_opening_and_sends_only_gripper(monkeypatch, tmp_path):
    from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module

    arm = FakeXArm("192.168.1.245")
    monkeypatch.setattr(uf_robot_module, "XArmAPI", lambda robot_ip: arm)
    monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None)

    config = UFRobotConfig(
        id="test_manual_gripper_robot",
        calibration_dir=tmp_path,
        robot_ip=arm.robot_ip,
        robot_dof=6,
        control_space="joint",
        gripper_type=1,
        manual_mode=True,
    )
    robot = uf_robot_module.UFRobot(config)
    robot.connect()

    assert ("set_gripper_enable", True) in arm.calls
    assert ("set_gripper_mode", 0) in arm.calls
    assert ("set_gripper_speed", 5000) in arm.calls
    assert not any(call[0] == "set_gripper_position" for call in arm.calls)

    robot.send_action({"J1.pos": 1.0, "gripper.pos": 0.5})

    assert any(call[0] == "getset_tgpio_modbus_data" for call in arm.calls)
    assert not any(call[0] == "set_servo_angle" for call in arm.calls)
    robot.disconnect()


def test_gripper_command_is_only_sent_after_target_changes(monkeypatch, tmp_path):
    from lerobot_robot_ufactory.robots.uf_robot import uf_robot as uf_robot_module

    arm = FakeXArm("192.168.1.245")
    monkeypatch.setattr(uf_robot_module, "XArmAPI", lambda robot_ip: arm)
    monkeypatch.setattr(uf_robot_module.time, "sleep", lambda _: None)

    config = UFRobotConfig(
        id="test_gripper_command_threshold",
        calibration_dir=tmp_path,
        robot_ip=arm.robot_ip,
        robot_dof=6,
        control_space="joint",
        gripper_type=1,
        manual_mode=True,
        gripper_command_threshold=0.01,
    )
    robot = uf_robot_module.UFRobot(config)
    robot.connect()

    robot.send_action({"gripper.pos": 0.5})
    robot.send_action({"gripper.pos": 0.505})
    robot.send_action({"gripper.pos": 0.52})

    writes = [call for call in arm.calls if call[0] == "getset_tgpio_modbus_data"]
    assert len(writes) == 2
    robot.disconnect()


def test_manual_record_config_has_no_teleop(monkeypatch):
    config_path = Path("config/manual_mode/xarm7_manual_record_config.yaml").resolve()
    monkeypatch.setattr(
        sys,
        "argv",
        ["uf-lerobot-record", "--config_path", str(config_path)],
    )

    config = get_cfg()

    assert config.robot.manual_mode is True
    assert config.robot.robot_dof == 7
    assert config.robot.manual_gripper_speed == 0.5
    assert config.teleop is None
    assert config.dataset.fps == 30


def test_prepare_dataset_root_leaves_new_root_for_lerobot_create(tmp_path):
    root = tmp_path / "nested" / "dataset"
    cfg = SimpleNamespace(dataset=SimpleNamespace(root=root), resume=False)

    _prepare_dataset_root(cfg)

    assert not root.exists()


def test_prepare_dataset_root_rejects_incomplete_resume(tmp_path):
    root = tmp_path / "dataset"
    (root / "meta").mkdir(parents=True)
    (root / "meta" / "info.json").write_text("{}")
    cfg = SimpleNamespace(dataset=SimpleNamespace(root=root), resume=True)

    with pytest.raises(RuntimeError, match="meta/tasks.parquet"):
        _prepare_dataset_root(cfg)


def test_prepare_dataset_root_rejects_resume_when_root_is_missing(tmp_path):
    cfg = SimpleNamespace(dataset=SimpleNamespace(root=tmp_path / "missing"), resume=True)

    with pytest.raises(RuntimeError, match="does not exist"):
        _prepare_dataset_root(cfg)


def test_prepare_dataset_root_resumes_complete_dataset_without_prompt(tmp_path, monkeypatch):
    root = tmp_path / "dataset"
    (root / "meta" / "episodes" / "chunk-000").mkdir(parents=True)
    (root / "data" / "chunk-000").mkdir(parents=True)
    (root / "meta" / "info.json").write_text("{}")
    (root / "meta" / "tasks.parquet").write_bytes(b"tasks")
    (root / "meta" / "episodes" / "chunk-000" / "file-000.parquet").write_bytes(b"episodes")
    (root / "data" / "chunk-000" / "file-000.parquet").write_bytes(b"data")
    cfg = SimpleNamespace(dataset=SimpleNamespace(root=root), resume=True)
    monkeypatch.setattr(record_module.sys.stdin, "isatty", lambda: True)

    _prepare_dataset_root(cfg)


def test_manual_record_loop_writes_actual_state_as_action(tmp_path):
    class FakeRobot:
        name = "fake_manual_robot"
        robot_type = name
        action_features = {"J1.pos": float, "J2.pos": float}
        observation_features = action_features

        def __init__(self):
            self.observation_count = 0
            self.sent_actions = []

        def get_observation(self):
            self.observation_count += 1
            value = float(self.observation_count)
            return {"J1.pos": value, "J2.pos": value + 1}

        def send_action(self, action):
            self.sent_actions.append(action.copy())
            return action

    robot = FakeRobot()
    action_pipeline, robot_pipeline, observation_pipeline = record_module.make_default_processors()
    features = record_module.combine_feature_dicts(
        record_module.aggregate_pipeline_dataset_features(
            pipeline=action_pipeline,
            initial_features=record_module.create_initial_features(action=robot.action_features),
            use_videos=False,
        ),
        record_module.aggregate_pipeline_dataset_features(
            pipeline=observation_pipeline,
            initial_features=record_module.create_initial_features(
                observation=robot.observation_features
            ),
            use_videos=False,
        ),
    )
    dataset = LeRobotDataset.create(
        "test/manual-record",
        fps=30,
        features=features,
        root=tmp_path / "dataset",
        robot_type=robot.robot_type,
        use_videos=False,
    )

    record_module.record_loop(
        robot=robot,
        events={"exit_early": False},
        fps=30,
        teleop_action_processor=action_pipeline,
        robot_action_processor=robot_pipeline,
        robot_observation_processor=observation_pipeline,
        dataset=dataset,
        control_time_s=0.001,
        single_task="test task",
        manual_mode=True,
    )

    assert dataset.episode_buffer["size"] == 1
    assert robot.sent_actions == [{"J1.pos": 2.0, "J2.pos": 3.0}]
    assert dataset.episode_buffer["action"][0].tolist() == [2.0, 3.0]

    dataset.save_episode()
    dataset.finalize()


def test_manual_recording_episode_resets_before_recording():
    class FakeRobot:
        def __init__(self):
            self.calls = []

        def reset_to_initial(self):
            self.calls.append("reset_to_initial")

    robot = FakeRobot()

    _prepare_recording_episode(robot, teleop=None, is_uf_teleop=False, manual_mode=True)

    assert robot.calls == ["reset_to_initial"]


def test_manual_record_loop_applies_keyboard_gripper_target():
    class FakeRobot:
        name = "fake_manual_robot"
        robot_type = name
        action_features = {"J1.pos": float, "gripper.pos": float}

        def __init__(self):
            self.sent_actions = []

        def get_observation(self):
            return {"J1.pos": 1.0, "gripper.pos": 0.5}

        def send_action(self, action):
            self.sent_actions.append(action.copy())
            return action

    robot = FakeRobot()
    action_pipeline, robot_pipeline, observation_pipeline = record_module.make_default_processors()

    record_module.record_loop(
        robot=robot,
        events={"exit_early": False},
        fps=10,
        teleop_action_processor=action_pipeline,
        robot_action_processor=robot_pipeline,
        robot_observation_processor=observation_pipeline,
        control_time_s=0.001,
        manual_mode=True,
        manual_gripper_keys={"close": True, "open": False},
        manual_gripper_speed=1.0,
    )

    assert robot.sent_actions == [{"J1.pos": 1.0, "gripper.pos": 0.6}]