import os
from pathlib import Path

import mujoco
import numpy as np
import pytest

from rm75_ik import DualArmAssembly, pose_errors
from rm75_ik.mujoco_backend import CONTROLLED_HOME_Q_RAD, DualArmMuJoCo
from rm75_ik.mujoco_model import build_normalized_dual_mjcf
from rm75_ik.mujoco_trajectories import (
    cartesian_demo_targets,
    joint_demo_trajectory,
    se3_target_trajectory,
    solve_pose_trajectory,
)
from rm75_ik import RM75IkSolver, RM75Kinematics, teleop_joint_limits


def test_normalized_model_has_standard_dual_arm_structure():
    xml, assets = build_normalized_dual_mjcf()
    model = mujoco.MjModel.from_xml_string(xml, assets)

    assert model.nq == model.nv == model.njnt == 14
    assert model.nu == 0
    assert model.nmesh == 27
    assert len([key for key in assets if key.endswith(".obj")]) == 19
    assert [model.joint(index).name for index in range(14)] == [
        f"{arm}_joint_{joint}"
        for arm in ("left", "right")
        for joint in range(1, 8)
    ]


@pytest.mark.parametrize("arm", ["left", "right"])
def test_mujoco_flange_matches_stage1_dual_assembly(arm):
    scene = DualArmMuJoCo(controlled_arm=arm)
    assembly = DualArmAssembly.from_source_urdf()
    q = np.deg2rad([20, -10, 30, 40, -20, 15, 80])

    scene.set_arm_configuration(arm, q)
    errors = pose_errors(scene.get_flange_pose(arm), assembly.forward(arm, q))

    assert errors[0] < 1e-9
    assert errors[1] < 1e-9


def test_invalid_configuration_does_not_change_state():
    scene = DualArmMuJoCo()
    before = scene.get_arm_configuration("left")

    with pytest.raises(ValueError):
        scene.set_arm_configuration("left", np.full(7, np.nan))

    np.testing.assert_array_equal(scene.get_arm_configuration("left"), before)


def test_playback_moves_only_controlled_arm():
    scene = DualArmMuJoCo(controlled_arm="left")
    inactive_before = scene.get_arm_configuration("right")
    trajectory = joint_demo_trajectory(CONTROLLED_HOME_Q_RAD, 20)

    result = scene.play_trajectory(trajectory)

    assert result.samples == 20
    np.testing.assert_array_equal(
        scene.get_arm_configuration("right"), inactive_before
    )
    np.testing.assert_allclose(
        scene.get_arm_configuration("left"), trajectory[-1], atol=0.0
    )


def test_se3_target_trajectory_preserves_endpoints():
    kinematics = RM75Kinematics(limits=teleop_joint_limits())
    start = kinematics.forward(CONTROLLED_HOME_Q_RAD)
    target = start.copy()
    target.translation = target.translation + np.array([0.04, 0.01, 0.0])

    trajectory = se3_target_trajectory(start, target, 25)

    assert pose_errors(trajectory[0], start) == pytest.approx((0.0, 0.0), abs=1e-12)
    assert pose_errors(trajectory[-1], target) == pytest.approx((0.0, 0.0), abs=1e-12)


def test_manual_drag_dynamics_moves_only_controlled_arm():
    scene = DualArmMuJoCo(controlled_arm="left")
    controlled_before = scene.get_arm_configuration("left")
    inactive_before = scene.get_arm_configuration("right")
    scene.configure_manual_drag()
    body_id = mujoco.mj_name2id(
        scene.model, mujoco.mjtObj.mjOBJ_BODY, "left_link_7"
    )
    scene.data.xfrc_applied[body_id, :3] = [0.0, 10.0, 0.0]

    for _ in range(50):
        scene.step_manual_drag()

    assert np.max(np.abs(scene.get_arm_configuration("left") - controlled_before)) > 1e-5
    np.testing.assert_array_equal(
        scene.get_arm_configuration("right"), inactive_before
    )


@pytest.mark.parametrize("kind", ["line", "arc", "orientation", "combined"])
def test_demo_cartesian_trajectories_are_solvable(kind):
    kinematics = RM75Kinematics(limits=teleop_joint_limits())
    solver = RM75IkSolver(kinematics)
    targets = cartesian_demo_targets(
        kind, kinematics.forward(CONTROLLED_HOME_Q_RAD), 30
    )

    solutions = solve_pose_trajectory(solver, targets, CONTROLLED_HOME_Q_RAD)

    assert solutions.shape == (30, 7)
    assert np.all(np.isfinite(solutions))


@pytest.mark.skipif(
    os.environ.get("MUJOCO_GL") != "egl",
    reason="EGL rendering test requires MUJOCO_GL=egl",
)
def test_offscreen_render_contains_both_arms():
    scene = DualArmMuJoCo()
    rgb = scene.render(640, 360)
    segmentation = scene.render(640, 360, segmentation=True)
    pairs = np.unique(segmentation.reshape(-1, 2), axis=0)
    names = {
        mujoco.mj_id2name(scene.model, mujoco.mjtObj.mjOBJ_GEOM, int(object_id))
        for object_id, object_type in pairs
        if object_type == mujoco.mjtObj.mjOBJ_GEOM and object_id >= 0
    }

    assert rgb.std() > 5.0
    assert any(name.startswith("left_") for name in names if name)
    assert any(name.startswith("right_") for name in names if name)


@pytest.fixture(scope="module")
def realman_sdk_root():
    value = os.environ.get("REALMAN_SDK_ROOT")
    if not value:
        pytest.skip("REALMAN_SDK_ROOT is not set")
    return Path(value)


def test_quick_stage2_validation_passes(realman_sdk_root):
    from rm75_ik.realman_reference import RealManFkReference
    from rm75_ik.stage2_validation import Stage2Settings, Stage2Validator

    validator = Stage2Validator(
        RealManFkReference(realman_sdk_root), settings=Stage2Settings.quick()
    )
    summary = validator.run()

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