import mujoco
import pytest

from api.robots import ROBOT_NAMES, ROBOTS, reset_to_robot_keyframe
from worker.sim_runner import RunConfig, _hold_controller, config_for, simulate

# A hinged pole driven by a torque motor, starting 0.3 rad from upright: gravity pulls it over unless held.
POLE = """<mujoco>
  <compiler angle="radian"/>
  <option timestep="0.002"/>
  <worldbody>
    <geom type="plane" size="1 1 0.1"/>
    <body name="base" pos="0 0 0.1">
      <body name="pole">
        <joint name="hinge" type="hinge" axis="0 1 0" range="-3 3"/>
        <geom type="capsule" fromto="0 0 0 0 0 1" size="0.03" mass="1"/>
      </body>
    </body>
  </worldbody>
  <actuator><motor joint="hinge" ctrlrange="-30 30"/></actuator>
</mujoco>"""


def pole_angle_after(tmp_path, controller, home=(0.3,)):
    (tmp_path / "scene.xml").write_text(POLE)
    config = RunConfig(2.0, controller, render=False, root_body="base", fall_height=None, home_qpos=home)
    simulate(tmp_path / "scene.xml", tmp_path / controller, config)
    model = mujoco.MjModel.from_xml_path(str(tmp_path / "scene.xml"))
    data = mujoco.MjData(model)
    reset_to_robot_keyframe(model, data, "base", home)
    hold = _hold_controller(model, data, config)
    if controller == "passive":
        model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_ACTUATION
    for _ in range(1000):
        if controller == "hold":
            hold.apply(data)
        mujoco.mj_step(model, data)
    return float(data.qpos[0])


def test_hold_uses_pd_for_torque_motors(tmp_path):
    assert pole_angle_after(tmp_path, "hold") == pytest.approx(0.3, abs=0.05)
    assert abs(pole_angle_after(tmp_path, "passive")) > 1.0  # falls over without the controller


def test_home_pose_is_used_only_without_a_keyframe(tmp_path):
    model = mujoco.MjModel.from_xml_string(POLE)
    data = mujoco.MjData(model)
    reset_to_robot_keyframe(model, data, "base", (0.3,))
    assert data.qpos[0] == pytest.approx(0.3)
    keyed = mujoco.MjModel.from_xml_string(
        POLE.replace("</mujoco>", '<keyframe><key qpos="-0.2"/></keyframe></mujoco>')
    )
    data = mujoco.MjData(keyed)
    reset_to_robot_keyframe(keyed, data, "base", (0.3,))
    assert data.qpos[0] == pytest.approx(-0.2)


def test_robot_configs():
    assert set(ROBOT_NAMES) == {"unitree_g1", "unitree_h1", "unitree_go2", "franka_emika_panda"}
    arm = config_for("franka_emika_panda", 5, "hold")
    assert arm.root_body == "link0" and arm.fall_height is None and arm.camera_distance < 3
    h1 = config_for("unitree_h1", 5, "hold")
    assert h1.home_qpos == ROBOTS["unitree_h1"].home_qpos and h1.kp == 1000
    assert config_for(None, 5, "hold").root_body is None


def test_api_accepts_every_robot(run_env):
    client, keys_, *_ = run_env
    key = keys_["alice"]
    r = client.post(
        "/simulate/generate",
        json={"description": "a quadruped near a box", "robot": "atlas"},
        headers={"X-API-Key": key},
    )
    assert r.status_code == 422 and "unitree_go2" in r.text


def test_robot_facts_with_unnamed_free_joint(tmp_path):
    """Go2 and H1 have an unnamed free joint; the facts given to the model must still build (was a TypeError)."""
    from api.robots import robot_facts

    spec = ROBOTS["unitree_go2"]
    (tmp_path / spec.directory).mkdir()
    (tmp_path / spec.directory / spec.model_file).write_text(
        '<mujoco><compiler angle="radian"/><worldbody><body name="base"><freejoint/><geom size=".1"/>'
        '<body name="leg"><joint name="knee" range="-1 1"/><geom size=".05"/></body></body></worldbody>'
        '<actuator><motor joint="knee"/></actuator></mujoco>'
    )
    facts = robot_facts(str(tmp_path), "unitree_go2")
    assert "(unnamed joint 0)" in facts and "knee [-1.00, 1.00]" in facts