import os
from dataclasses import replace
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pinocchio as pin
import pytest

from rm75_ik import pose_errors
from rm75_ik.mujoco_robot import MujocoRobot
from rm75_ik.teleop_config import load_dual_arm_profiles
from rm75_ik.teleop_control import (
    ControllerSample,
    DualArmQpTeleopController,
    RelativePoseMapper,
    SafetyState,
)


ROOT = Path(__file__).resolve().parents[2]


@pytest.fixture(scope="module")
def profiles():
    return load_dual_arm_profiles(
        ROOT / "xr_rm_bringup/config/dual_arm_rm75.yaml",
        ROOT / "xr_rm_bringup/config/peripherals_rm75.yaml",
    )


def sample(arm, grip, position=(0.0, 0.0, 0.0), quaternion=(0, 0, 0, 1)):
    return ControllerSample(
        hand=arm,
        grip=grip,
        trigger=0.0,
        position_m=np.asarray(position, dtype=float),
        quaternion_xyzw=np.asarray(quaternion, dtype=float),
    )


def test_profiles_use_expected_tools_and_mapping(profiles):
    assert profiles["left"].active_tool_name == "minisci"
    assert profiles["right"].active_tool_name == "omnipic"
    np.testing.assert_array_equal(
        profiles["left"].xr_to_robot,
        [[0, -1, 0], [0, 0, 1], [-1, 0, 0]],
    )
    np.testing.assert_array_equal(
        profiles["right"].xr_to_robot,
        [[0, 1, 0], [0, 0, 1], [1, 0, 0]],
    )


@pytest.mark.parametrize(
    ("arm", "expected_delta"),
    [("left", [-0.01, 0.0, 0.0]), ("right", [0.01, 0.0, 0.0])],
)
def test_relative_position_mapping_matches_ros_configuration(
    profiles, arm, expected_delta
):
    profile = replace(
        profiles[arm],
        scale=1.0,
        deadband_m=0.0,
        target_filter_alpha=1.0,
        target_filter_alpha_fast=1.0,
        max_linear_speed_m_s=10.0,
        orientation_filter_alpha=1.0,
        max_orientation_speed_rad_s=10.0,
    )
    mapper = RelativePoseMapper(profile)
    start = profile.initial_tcp
    mapper.map(sample(arm, True), start, 0.1)

    mapped = mapper.map(sample(arm, True, [0.0, 0.01, 0.0]), start, 0.1)

    np.testing.assert_allclose(
        mapped.target_tcp.translation - start.translation,
        expected_delta,
        atol=1e-12,
    )


def test_mujoco_initializes_both_configured_tool_tcp_poses(profiles):
    robot = MujocoRobot(profiles)
    try:
        for arm in ("left", "right"):
            position_error, orientation_error = pose_errors(
                robot.get_tcp_pose(arm), profiles[arm].initial_tcp
            )
            assert position_error < 1e-3
            assert orientation_error < np.deg2rad(0.1)
            diagnostic = robot.initial_pose_diagnostics[arm]
            assert diagnostic.configured_joint_position_error_m > 0.05
    finally:
        robot.close()


def test_mujoco_dual_command_is_atomic(profiles):
    robot = MujocoRobot(profiles)
    try:
        before = robot.read_joint_positions().positions_rad
        with pytest.raises(ValueError):
            robot.command_joint_positions(
                {
                    "left": before["left"] + 0.01,
                    "right": np.full(7, np.nan),
                }
            )
        after = robot.read_joint_positions().positions_rad
        np.testing.assert_array_equal(after["left"], before["left"])
        np.testing.assert_array_equal(after["right"], before["right"])
    finally:
        robot.close()


def test_dual_controller_has_no_grip_jump_and_moves_both_arms(profiles):
    robot = MujocoRobot(profiles)
    controller = DualArmQpTeleopController(robot, profiles)
    try:
        initial = robot.read_joint_positions().positions_rad
        for arm in ("left", "right"):
            controller.update_sample(sample(arm, True), 0.0)
        engaged = controller.step(0.0)
        after_engage = robot.read_joint_positions().positions_rad
        assert engaged.state is SafetyState.ACTIVE
        for arm in ("left", "right"):
            np.testing.assert_allclose(after_engage[arm], initial[arm], atol=1e-10)

        for arm in ("left", "right"):
            controller.update_sample(sample(arm, True, [0, 0.005, 0]), 0.01)
        moved_result = controller.step(0.01)
        moved = robot.read_joint_positions().positions_rad
        assert moved_result.state is SafetyState.ACTIVE
        for arm in ("left", "right"):
            assert np.max(np.abs(moved[arm] - initial[arm])) > 1e-5
    finally:
        controller.close()


def test_active_arm_timeout_faults_both_and_requires_release(profiles):
    robot = MujocoRobot(profiles)
    controller = DualArmQpTeleopController(robot, profiles)
    try:
        for arm in ("left", "right"):
            controller.update_sample(sample(arm, True), 0.0)
        assert controller.step(0.0).state is SafetyState.ACTIVE

        fault = controller.step(1.0)
        assert fault.state is SafetyState.FAULT
        assert "timed out" in fault.reason

        for arm in ("left", "right"):
            controller.update_sample(sample(arm, False), 1.01)
        cleared = controller.step(1.01)
        assert cleared.state is SafetyState.IDLE
    finally:
        controller.close()


def test_reset_restores_initial_state_and_requires_fresh_dual_release(profiles):
    robot = MujocoRobot(profiles)
    controller = DualArmQpTeleopController(robot, profiles)
    try:
        initial = robot.read_joint_positions().positions_rad
        for arm in ("left", "right"):
            controller.update_sample(sample(arm, True), 0.0)
        assert controller.step(0.0).state is SafetyState.ACTIVE

        for arm in ("left", "right"):
            controller.update_sample(sample(arm, True, [0, 0.005, 0]), 0.01)
        assert controller.step(0.01).state is SafetyState.ACTIVE
        moved = robot.read_joint_positions().positions_rad
        assert any(
            np.max(np.abs(moved[arm] - initial[arm])) > 1e-5
            for arm in ("left", "right")
        )

        reset = controller.reset_to_initial()
        assert reset.state is SafetyState.RESETTING
        restored = robot.read_joint_positions().positions_rad
        for arm in ("left", "right"):
            np.testing.assert_array_equal(restored[arm], initial[arm])

        assert controller.step(0.02).state is SafetyState.RESETTING
        controller.update_sample(sample("left", False), 0.03)
        controller.update_sample(sample("right", True), 0.03)
        assert controller.step(0.03).state is SafetyState.RESETTING
        controller.update_sample(sample("right", False), 0.04)
        assert controller.step(0.04).state is SafetyState.IDLE
    finally:
        controller.close()


def test_message_compatible_adapter_rejects_wrong_hand():
    message = SimpleNamespace(
        hand="right",
        grip=True,
        trigger=0.0,
        pose=SimpleNamespace(
            position=SimpleNamespace(x=0.0, y=0.0, z=0.0),
            orientation=SimpleNamespace(x=0.0, y=0.0, z=0.0, w=1.0),
        ),
    )

    with pytest.raises(ValueError, match="does not match"):
        ControllerSample.from_message(message, expected_arm="left")


@pytest.mark.skipif(
    not os.environ.get("REALMAN_SDK_ROOT"),
    reason="REALMAN_SDK_ROOT is not set",
)
def test_quick_stage3_validation_passes():
    from rm75_ik.realman_reference import RealManFkReference
    from rm75_ik.stage3_validation import Stage3Settings, Stage3Validator

    validator = Stage3Validator(
        RealManFkReference(Path(os.environ["REALMAN_SDK_ROOT"])),
        ROOT / "xr_rm_bringup/config/dual_arm_rm75.yaml",
        ROOT / "xr_rm_bringup/config/peripherals_rm75.yaml",
        Stage3Settings.quick(),
    )

    summary = validator.run()

    assert summary["passed"] is True
    assert summary["failure_count"] == 0
