import pytest
from conftest import GOOD_SCENE

from api.llm import ScriptedLLM
from api.mjcf_check import extract_mjcf, static_check, validate_mjcf
from api.simulate import generate_scene


def fenced(xml: str) -> str:
    return f"Here is the scene:\n```xml\n{xml}\n```\nNotes: crate is free."


class TestExtract:
    def test_fenced_block(self):
        assert extract_mjcf(fenced(GOOD_SCENE)) == GOOD_SCENE

    def test_bare_document(self):
        assert extract_mjcf(f"Sure! {GOOD_SCENE} done") == GOOD_SCENE

    def test_skips_non_mjcf_blocks(self):
        text = "```python\nprint(1)\n```\n" + fenced(GOOD_SCENE)
        assert extract_mjcf(text) == GOOD_SCENE

    def test_nothing(self):
        assert extract_mjcf("I cannot do that.") is None


class TestStaticCheck:
    def test_good_scene_passes(self):
        assert static_check(GOOD_SCENE, "unitree_g1") is None

    @pytest.mark.parametrize(
        "xml, message",
        [
            ("<mujoco><include file='/etc/passwd'/></mujoco>", "absolute"),
            ("<mujoco><include file='g1.xml'/><asset><mesh file='../../secret.stl'/></asset></mujoco>", "'..'"),
            ("<mujoco><include file='g1.xml'/><compiler meshdir='/tmp'/></mujoco>", "absolute"),
            ("<mujoco><include file='other.xml'/></mujoco>", "not allowed"),
            ("<mujoco><worldbody/></mujoco>", "must contain"),
            ("<robot/>", "root element"),
            ("<mujoco><worldbody></mujoco>", "parse error"),
            ('<!DOCTYPE x [<!ENTITY e SYSTEM "file:///etc/passwd">]><mujoco/>', "DOCTYPE"),
        ],
    )
    def test_rejections(self, xml, message):
        result = static_check(xml, "unitree_g1")
        assert result is not None and not result.valid and result.stage == "static"
        assert message in result.error

    def test_no_robot_forbids_includes(self):
        assert static_check("<mujoco><include file='g1.xml'/></mujoco>", None).error.startswith("<include")
        assert static_check("<mujoco><worldbody/></mujoco>", None) is None

    def test_size_limit(self):
        huge = "<mujoco><include file='g1.xml'/>" + "<!--" + "x" * 400_000 + "--></mujoco>"
        assert "larger than" in static_check(huge, "unitree_g1").error


class TestLoadAndStep:
    def test_valid_scene(self, menagerie):
        r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=0.2)
        assert r.valid and r.stage == "ok", r
        assert r.stats["nbody"] >= 4 and r.stats["simulated_seconds"] == pytest.approx(0.2, abs=0.01)

    def test_mujoco_compiler_error_is_reported(self, menagerie):
        r = validate_mjcf(GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"'), "unitree_g1", menagerie)
        assert not r.valid and r.stage == "load"
        assert "frictoin" in r.error

    def test_robot_root_body_required(self, menagerie):
        (menagerie / "unitree_g1" / "g1.xml").write_text("<mujoco><worldbody/></mujoco>")
        r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=0.1)
        assert not r.valid and "pelvis" in r.error

    def test_missing_menagerie(self, tmp_path):
        with pytest.raises(FileNotFoundError, match="rag.ingest"):
            validate_mjcf(GOOD_SCENE, "unitree_g1", tmp_path / "nope")

    def test_timeout(self, menagerie):
        r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=1e6, timeout=2)
        assert not r.valid and "longer than" in r.error


class TestGenerateScene:
    def test_repair_loop_fixes_compiler_error(self, store, settings):
        bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"')
        llm = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)])
        result = generate_scene("G1 next to a slope and a crate", "unitree_g1", llm, store, settings)
        assert result.valid and len(result.attempts) == 2
        assert result.mjcf == GOOD_SCENE
        repair_prompt = llm.calls[1][-1].content
        assert "frictoin" in repair_prompt and "stage 'load'" in repair_prompt
        assert llm.calls[1][-2].role == "assistant"

    def test_prompt_is_grounded(self, store, settings):
        llm = ScriptedLLM([fenced(GOOD_SCENE)])
        result = generate_scene("a floor with friction and a box", "unitree_g1", llm, store, settings)
        system, user = llm.calls[0]
        assert '<include file="g1.xml"/>' in system.content
        assert "ROBOT FACTS" in user.content and "left_knee_joint" in user.content  # read from the model file
        assert "REFERENCE CONTEXT" in user.content and "doc#geom" in user.content
        assert "doc#geom" in result.sources

    def test_gives_up_after_max_attempts(self, store, settings):
        llm = ScriptedLLM(["no xml here", "still none"])
        result = generate_scene("anything at all", "unitree_g1", llm, store, settings, max_attempts=2)
        assert not result.valid and result.validation.stage == "extract" and len(result.attempts) == 2

    def test_unknown_robot(self, store, settings):
        with pytest.raises(ValueError, match="unknown robot"):
            generate_scene("x" * 20, "atlas", ScriptedLLM([]), store, settings)

    def test_no_robot(self, store, settings):
        scene = "<mujoco><worldbody><geom name='floor' type='plane' size='0 0 1'/></worldbody></mujoco>"
        result = generate_scene("just a floor please", None, ScriptedLLM([fenced(scene)]), store, settings)
        assert result.valid


SCHEMA_RST = """
.. dropdown:: :ref:`mujoco<mujoco>` :octicon:`star`

   .. grid-item::
      :ref:`model<mujoco-model>`

   .. dropdown:: :ref:`(world)body<body>` :octicon:`sync`

         .. grid-item::
            :ref:`name<body-name>`

         .. grid-item::
            :ref:`pos<body-pos>`

      .. dropdown:: :ref:`geom<body-geom>` |*|

            .. grid-item::
               :ref:`type<body-geom-type>`

            .. grid-item::
               :ref:`friction<body-geom-friction>`
"""


def test_parse_mjcf_schema():
    from rag.chunking import parse_mjcf_schema

    schema = parse_mjcf_schema(SCHEMA_RST)
    assert schema["body"] == {"element": "body", "path": "body", "attributes": ["name", "pos"]}
    assert schema["body-geom"]["attributes"] == ["type", "friction"]
    assert schema["mujoco"]["attributes"] == ["model"]


def test_repair_prompt_lists_valid_attributes(store, settings):
    import json

    from rag.chunking import parse_mjcf_schema

    settings.data_dir.mkdir(parents=True, exist_ok=True)
    settings.mjcf_schema_file.write_text(json.dumps(parse_mjcf_schema(SCHEMA_RST)))
    (settings.menagerie_dir / "unitree_g1" / "scene.xml").write_text('<mujoco><include file="g1.xml"/></mujoco>')
    bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"')
    llm = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)])
    result = generate_scene("G1 next to a crate", "unitree_g1", llm, store, settings)
    first_prompt = llm.calls[0][1].content
    assert "MENAGERIE EXAMPLE" in first_prompt and "MJCF SCHEMA" in first_prompt
    assert "<geom> (body/geom): type, friction" in first_prompt
    assert "Valid attributes for that element:\n<geom> (body/geom): type, friction" in llm.calls[1][-1].content
    assert result.valid and result.sources[0] == "mujoco_menagerie:unitree_g1/scene.xml"


def test_validator_ignores_rendering_settings(menagerie, monkeypatch):
    """Validation never renders, so a GL backend that is unusable here must not break it (seen in CI)."""
    monkeypatch.setenv("MUJOCO_GL", "not-a-backend")
    monkeypatch.setenv("PYOPENGL_PLATFORM", "not-a-backend")
    r = validate_mjcf(GOOD_SCENE, "unitree_g1", menagerie, sim_seconds=0.1)
    assert r.valid, r.error