"""Tests for spec_config.py."""

import math
from typing import cast

import mujoco
import pytest

from mjlab.utils.spec_config import (
  CameraCfg,
  CollisionCfg,
  GeomCfg,
  LightCfg,
  MaterialCfg,
  MeshCfg,
  TextureCfg,
)


@pytest.fixture
def simple_robot_xml():
  """Minimal robot XML for testing."""
  return """
    <mujoco>
      <worldbody>
        <body name="base" pos="0 0 1">
          <freejoint name="free_joint"/>
          <geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="1.0"/>
          <body name="link1" pos="0 0 0">
            <joint name="joint1" type="hinge" axis="0 0 1" range="0 1.57"/>
            <geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
          </body>
          <body name="link2" pos="0 0 0">
            <joint name="joint2" type="hinge" axis="0 0 1" range="0 1.57"/>
            <geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
          </body>
        </body>
      </worldbody>
    </mujoco>
    """


# Collision Tests


@pytest.fixture
def multi_geom_spec():
  """Spec with multiple geoms for collision testing."""
  spec = mujoco.MjSpec()
  body = spec.worldbody.add_body(name="test_body")
  body.add_geom(
    name="left_foot1_collision", type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1]
  )
  body.add_geom(
    name="right_foot3_collision", type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1]
  )
  body.add_geom(
    name="arm_collision", type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1]
  )
  return spec


def test_collision_basic_properties(multi_geom_spec):
  """CollisionCfg should set basic collision properties."""
  collision_cfg = CollisionCfg(
    geom_names_expr=("arm_collision",), contype=2, conaffinity=3, condim=4, priority=1
  )
  collision_cfg.edit_spec(multi_geom_spec)

  geom = multi_geom_spec.geom("arm_collision")
  assert geom.contype == 2
  assert geom.conaffinity == 3
  assert geom.condim == 4
  assert geom.priority == 1


def test_collision_regex_matching(multi_geom_spec):
  """CollisionCfg should support regex pattern matching."""
  collision_cfg = CollisionCfg(
    geom_names_expr=(r"^(left|right)_foot\d_collision$",),
    contype=1,
    conaffinity=1,
    condim=3,
    priority=1,
    friction=(0.6,),
    disable_other_geoms=False,
  )
  collision_cfg.edit_spec(multi_geom_spec)

  left_foot = multi_geom_spec.geom("left_foot1_collision")
  assert left_foot.condim == 3
  assert left_foot.priority == 1
  assert left_foot.friction[0] == 0.6

  right_foot = multi_geom_spec.geom("right_foot3_collision")
  assert right_foot.condim == 3

  arm = multi_geom_spec.geom("arm_collision")
  assert arm.condim == 3  # Default unchanged.


def test_collision_dict_field_resolution(multi_geom_spec):
  """CollisionCfg should support dict-based field resolution."""
  collision_cfg = CollisionCfg(
    geom_names_expr=(r".*_foot\d_collision$", "arm_collision"),
    contype=1,
    conaffinity=1,
    condim={r".*_foot\d_collision$": 3, "arm_collision": 1},
    priority={r".*_foot\d_collision$": 2, "arm_collision": 0},
  )
  collision_cfg.edit_spec(multi_geom_spec)

  left_foot = multi_geom_spec.geom("left_foot1_collision")
  assert left_foot.condim == 3
  assert left_foot.priority == 2

  arm = multi_geom_spec.geom("arm_collision")
  assert arm.condim == 1
  assert arm.priority == 0


def test_collision_margin_gap_solmix(multi_geom_spec):
  """CollisionCfg should set margin, gap, and solmix on matched geoms."""
  collision_cfg = CollisionCfg(
    geom_names_expr=("arm_collision",),
    contype=1,
    conaffinity=1,
    condim=3,
    priority=0,
    margin=0.01,
    gap=0.005,
    solmix=0.5,
    disable_other_geoms=False,
  )
  collision_cfg.edit_spec(multi_geom_spec)

  geom = multi_geom_spec.geom("arm_collision")
  assert geom.margin == pytest.approx(0.01)
  assert geom.gap == pytest.approx(0.005)
  assert geom.solmix == pytest.approx(0.5)


def test_collision_margin_gap_solmix_dict(multi_geom_spec):
  """CollisionCfg should support dict-based margin, gap, solmix overrides."""
  collision_cfg = CollisionCfg(
    geom_names_expr=(r".*_foot\d_collision$", "arm_collision"),
    contype=1,
    conaffinity=1,
    condim=3,
    priority=0,
    margin={r".*_foot\d_collision$": 0.02, "arm_collision": 0.01},
    gap={r".*_foot\d_collision$": 0.01, "arm_collision": 0.0},
    solmix={r".*_foot\d_collision$": 0.8, "arm_collision": 0.2},
    disable_other_geoms=False,
  )
  collision_cfg.edit_spec(multi_geom_spec)

  left_foot = multi_geom_spec.geom("left_foot1_collision")
  assert left_foot.margin == pytest.approx(0.02)
  assert left_foot.gap == pytest.approx(0.01)
  assert left_foot.solmix == pytest.approx(0.8)

  arm = multi_geom_spec.geom("arm_collision")
  assert arm.margin == pytest.approx(0.01)
  assert arm.solmix == pytest.approx(0.2)


def test_collision_structural_dict_requires_coverage(multi_geom_spec):
  """Dict-valued structural fields must cover every matched geom."""
  collision_cfg = CollisionCfg(
    geom_names_expr=(".*_collision",),
    contype=1,
    conaffinity=1,
    condim={r".*_foot\d_collision$": 3},  # Does not cover arm_collision.
    priority=0,
  )
  with pytest.raises(ValueError, match="condim dict does not match geom"):
    collision_cfg.edit_spec(multi_geom_spec)


def test_collision_structural_fields_are_required():
  """Structural fields cannot be None."""
  cfg = CollisionCfg(
    geom_names_expr=(".*",),
    contype=1,
    conaffinity=1,
    condim=cast(int, None),
    priority=0,
  )
  with pytest.raises(ValueError, match="condim is required"):
    cfg.validate()


def test_collision_disable_other_geoms(multi_geom_spec):
  """CollisionCfg should disable non-matching geoms when requested."""
  collision_cfg = CollisionCfg(
    geom_names_expr=("left_foot1_collision",),
    contype=2,
    conaffinity=1,
    condim=3,
    priority=0,
    disable_other_geoms=True,
  )
  collision_cfg.edit_spec(multi_geom_spec)

  left_foot = multi_geom_spec.geom("left_foot1_collision")
  assert left_foot.contype == 2

  right_foot = multi_geom_spec.geom("right_foot3_collision")
  assert right_foot.contype == 0
  assert right_foot.conaffinity == 0

  arm = multi_geom_spec.geom("arm_collision")
  assert arm.contype == 0


# fmt: off
@pytest.mark.parametrize(
  "param,value,expected_error",
  [
    ("condim", -1, "condim must be one of"),
    ("condim", 2, "condim must be one of"),
    ("contype", -1, "contype must be non-negative"),
    ("conaffinity", -1, "conaffinity must be non-negative"),
    ("priority", -1, "priority must be non-negative"),
    ("margin", -0.1, "margin must be non-negative"),
    ("gap", -0.1, "gap must be non-negative"),
    ("solmix", 1.5, r"solmix must be in \[0, 1\]"),
  ],
)
# fmt: on
def test_collision_validation(param, value, expected_error):
  """CollisionCfg should validate parameters."""
  with pytest.raises(ValueError, match=expected_error):
    cfg = CollisionCfg(
      geom_names_expr=("test",),
      contype=value if param == "contype" else 1,
      conaffinity=value if param == "conaffinity" else 1,
      condim=value if param == "condim" else 3,
      priority=value if param == "priority" else 0,
      margin=value if param == "margin" else None,
      gap=value if param == "gap" else None,
      solmix=value if param == "solmix" else None,
    )
    cfg.validate()


# Mesh Tests


def _sphere_points(n: int) -> list[float]:
  """Deterministic point cloud on a unit sphere (Fibonacci lattice).

  Every point is on the hull, so the unconstrained hull has n vertices.
  """
  phi = math.pi * (3.0 - math.sqrt(5.0))
  verts: list[float] = []
  for i in range(n):
    y = 1.0 - 2.0 * (i + 0.5) / n
    r = math.sqrt(max(0.0, 1.0 - y * y))
    theta = phi * i
    verts += [math.cos(theta) * r, y, math.sin(theta) * r]
  return verts


@pytest.fixture
def multi_mesh_spec():
  """Spec with two mesh assets, each a 100-vertex sphere point cloud."""
  spec = mujoco.MjSpec()
  for name in ("blob_a", "blob_b"):
    mesh = spec.add_mesh()
    mesh.name = name
    mesh.uservert = _sphere_points(100)
  return spec


def _hull_numvert(spec: mujoco.MjSpec, mesh_name: str) -> int:
  """Compile and return the convex-hull vertex count for a mesh asset.

  maxhullvert caps the hull stored in mesh_graph (which the narrowphase support
  function walks), not the raw mesh_vert array, so we read it back from the
  compiled graph: the first int at mesh_graphadr is the hull vertex count.
  """
  body = spec.worldbody.add_body()
  geom = body.add_geom()
  geom.type = mujoco.mjtGeom.mjGEOM_MESH
  geom.meshname = mesh_name
  model = spec.compile()
  mesh_id = model.mesh(mesh_name).id
  return int(model.mesh_graph[model.mesh_graphadr[mesh_id]])


def test_mesh_maxhullvert_scalar(multi_mesh_spec):
  """MeshCfg should set maxhullvert on every matched mesh asset."""
  MeshCfg(mesh_names_expr=(".*",), maxhullvert=12).edit_spec(multi_mesh_spec)

  assert multi_mesh_spec.mesh("blob_a").maxhullvert == 12
  assert multi_mesh_spec.mesh("blob_b").maxhullvert == 12


def test_mesh_maxhullvert_regex_matching(multi_mesh_spec):
  """Only meshes matching the regex should be modified; others stay default."""
  MeshCfg(mesh_names_expr=("blob_a",), maxhullvert=8).edit_spec(multi_mesh_spec)

  assert multi_mesh_spec.mesh("blob_a").maxhullvert == 8
  assert multi_mesh_spec.mesh("blob_b").maxhullvert == -1  # Unchanged default.


def test_mesh_maxhullvert_dict_resolution(multi_mesh_spec):
  """A dict should resolve per-mesh values by pattern."""
  MeshCfg(
    mesh_names_expr=(".*",),
    maxhullvert={"blob_a": 6, "blob_b": 20},
  ).edit_spec(multi_mesh_spec)

  assert multi_mesh_spec.mesh("blob_a").maxhullvert == 6
  assert multi_mesh_spec.mesh("blob_b").maxhullvert == 20


def test_mesh_unconstrained_hull_keeps_all_vertices(multi_mesh_spec):
  """Baseline: without a cap, every sphere point is a hull vertex."""
  assert _hull_numvert(multi_mesh_spec, "blob_a") == 100


def test_mesh_maxhullvert_caps_compiled_hull(multi_mesh_spec):
  """The cap should actually shrink the compiled convex hull below the baseline."""
  MeshCfg(mesh_names_expr=("blob_a",), maxhullvert=10).edit_spec(multi_mesh_spec)
  assert _hull_numvert(multi_mesh_spec, "blob_a") == 10


@pytest.mark.parametrize("value", [0, 2, 3, -2])
def test_mesh_maxhullvert_validation(value):
  """maxhullvert must be -1 or greater than 3."""
  with pytest.raises(ValueError, match="maxhullvert must be"):
    MeshCfg(mesh_names_expr=(".*",), maxhullvert=value).validate()


def test_mesh_maxhullvert_dict_validation():
  """Validation should also cover dict-valued maxhullvert."""
  with pytest.raises(ValueError, match="maxhullvert must be.*pattern 'blob_a'"):
    MeshCfg(mesh_names_expr=(".*",), maxhullvert={"blob_a": 2}).validate()


def test_mesh_maxhullvert_allows_unlimited():
  """-1 (unlimited) is valid and is the MuJoCo default."""
  MeshCfg(mesh_names_expr=(".*",), maxhullvert=-1).validate()


# Visual Element Tests


def test_texture_cfg():
  """TextureCfg should add textures to spec."""
  spec = mujoco.MjSpec()
  texture_cfg = TextureCfg(
    name="test_texture",
    type="2d",
    builtin="checker",
    rgb1=(1.0, 0.0, 0.0),
    rgb2=(0.0, 1.0, 0.0),
    width=64,
    height=64,
  )
  texture_cfg.edit_spec(spec)

  texture = spec.texture("test_texture")
  assert texture.name == "test_texture"


def test_texture_cfg_image_fields():
  """TextureCfg should set file, grid, nchannel, flip, and random fields."""
  spec = mujoco.MjSpec()
  texture_cfg = TextureCfg(
    name="test_texture",
    type="cube",
    mark="random",
    random=0.25,
    file="sky.png",
    gridsize=(3, 4),
    gridlayout=".U..LFRB.D..",
    nchannel=4,
    hflip=True,
    vflip=True,
  )
  texture_cfg.edit_spec(spec)

  texture = spec.texture("test_texture")
  assert texture.file == "sky.png"
  assert tuple(texture.gridsize) == (3, 4)
  assert texture.nchannel == 4
  assert texture.hflip
  assert texture.vflip
  assert texture.random == pytest.approx(0.25)


@pytest.mark.parametrize(
  "param,value,expected_error",
  [
    ("width", 0, "width must be positive"),
    ("width", -1, "width must be positive"),
    ("height", 0, "height must be positive"),
    ("height", -1, "height must be positive"),
  ],
)
def test_texture_cfg_validation(param, value, expected_error):
  """TextureCfg should validate explicitly set width and height."""
  with pytest.raises(ValueError, match=expected_error):
    TextureCfg(
      name="test_texture",
      type="2d",
      builtin="checker",
      width=value if param == "width" else 64,
      height=value if param == "height" else 64,
    ).validate()


def test_texture_requires_source():
  """A texture with no builtin, file, or cubefiles has nothing to render."""
  with pytest.raises(ValueError, match="must specify a builtin pattern"):
    TextureCfg(name="test_texture", type="2d").validate()


@pytest.mark.parametrize("gridsize", [None, (1, 1)])
def test_texture_single_file_skybox_validation(gridsize):
  """A single-image skybox cannot be sampled as six faces by the Warp renderer."""
  with pytest.raises(ValueError, match="must declare a gridsize"):
    TextureCfg(
      name="sky", type="skybox", file="sky.png", gridsize=gridsize
    ).validate()


def test_texture_skybox_allows_grid_and_cubefiles():
  """A larger grid or explicit cubefiles cover all six faces."""
  TextureCfg(name="sky", type="skybox", file="sky.png", gridsize=(6, 1)).validate()
  TextureCfg(
    name="sky", type="skybox", cubefiles=("a", "b", "c", "d", "e", "f")
  ).validate()


def test_texture_skybox_builtin_overrides_single_file():
  """A builtin pattern takes precedence over the file, so the grid is irrelevant."""
  TextureCfg(
    name="sky", type="skybox", builtin="gradient", file="sky.png"
  ).validate()


def test_material_cfg():
  """MaterialCfg should add materials to spec."""
  spec = mujoco.MjSpec()
  material_cfg = MaterialCfg(
    name="test_material",
    texuniform=True,
    texrepeat=(2, 2),
    reflectance=0.5,
  )
  material_cfg.edit_spec(spec)

  material = spec.material("test_material")
  assert material.name == "test_material"


def test_light_cfg():
  """LightCfg should add lights to spec."""
  spec = mujoco.MjSpec()
  light_cfg = LightCfg(
    name="test_light",
    body="world",
    type="spot",
    pos=(1.0, 2.0, 3.0),
    dir=(0.0, 0.0, -1.0),
  )
  light_cfg.edit_spec(spec)

  light = spec.light("test_light")
  assert light.name == "test_light"


def test_light_cfg_color_and_attenuation():
  """LightCfg should set diffuse, specular, ambient, active, attenuation."""
  spec = mujoco.MjSpec()
  light_cfg = LightCfg(
    name="test_light",
    diffuse=(1.0, 0.0, 0.0),
    specular=(0.0, 1.0, 0.0),
    ambient=(0.0, 0.0, 1.0),
    active=False,
    attenuation=(1.0, 0.5, 0.25),
  )
  light_cfg.edit_spec(spec)

  light = spec.light("test_light")
  assert tuple(light.diffuse) == pytest.approx((1.0, 0.0, 0.0))
  assert tuple(light.specular) == pytest.approx((0.0, 1.0, 0.0))
  assert tuple(light.ambient) == pytest.approx((0.0, 0.0, 1.0))
  assert not light.active
  assert tuple(light.attenuation) == pytest.approx((1.0, 0.5, 0.25))


def test_camera_cfg():
  """CameraCfg should add cameras to spec."""
  spec = mujoco.MjSpec()
  camera_cfg = CameraCfg(
    name="test_camera", body="world", fovy=60.0, pos=(0.0, 0.0, 5.0)
  )
  camera_cfg.edit_spec(spec)

  camera = spec.camera("test_camera")
  assert camera.name == "test_camera"


def test_material_cfg_geom_assignment():
  """MaterialCfg.geom_names_expr assigns material to matching geoms."""
  spec = mujoco.MjSpec()
  body = spec.worldbody.add_body(name="test_body")
  body.add_geom(
    name="link1_visual",
    type=mujoco.mjtGeom.mjGEOM_BOX,
    size=[0.1, 0.1, 0.1],
  )
  body.add_geom(
    name="link2_visual",
    type=mujoco.mjtGeom.mjGEOM_BOX,
    size=[0.1, 0.1, 0.1],
  )
  body.add_geom(
    name="arm_collision",
    type=mujoco.mjtGeom.mjGEOM_BOX,
    size=[0.1, 0.1, 0.1],
  )

  mat_cfg = MaterialCfg(
    name="my_mat",
    texuniform=True,
    texrepeat=(1, 1),
    geom_names_expr=(r".*_visual$",),
  )
  mat_cfg.edit_spec(spec)

  assert spec.geom("link1_visual").material == "my_mat"
  assert spec.geom("link2_visual").material == "my_mat"
  assert spec.geom("arm_collision").material == ""


def test_geom_group_basic(multi_geom_spec):
  """GeomCfg should set the group of matching geoms."""
  group_cfg = GeomCfg(
    geom_names_expr=(r"^(left|right)_foot\d_collision$",), group=3
  )
  group_cfg.edit_spec(multi_geom_spec)

  assert multi_geom_spec.geom("left_foot1_collision").group == 3
  assert multi_geom_spec.geom("right_foot3_collision").group == 3

  arm = multi_geom_spec.geom("arm_collision")
  assert arm.group == 0  # Default unchanged.


def test_geom_group_none_is_noop(multi_geom_spec):
  """GeomCfg should leave the group alone when it is None."""
  multi_geom_spec.geom("arm_collision").group = 2
  GeomCfg(geom_names_expr=(".*",)).edit_spec(multi_geom_spec)

  assert multi_geom_spec.geom("arm_collision").group == 2


def test_geom_group_dict_field_resolution(multi_geom_spec):
  """GeomCfg should support dict-based field resolution."""
  group_cfg = GeomCfg(
    geom_names_expr=(r".*_foot\d_collision$", "arm_collision"),
    group={r".*_foot\d_collision$": 3, "arm_collision": 4},
  )
  group_cfg.edit_spec(multi_geom_spec)

  assert multi_geom_spec.geom("left_foot1_collision").group == 3
  assert multi_geom_spec.geom("arm_collision").group == 4


def test_geom_group_unnamed_geoms():
  """GeomCfg should update every matching geom, including unnamed ones."""
  spec = mujoco.MjSpec()
  body = spec.worldbody.add_body(name="test_body")
  for _ in range(3):
    body.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1])

  group_cfg = GeomCfg(geom_names_expr=(".*",), group=3)
  group_cfg.edit_spec(spec)

  assert [g.group for g in spec.geoms] == [3, 3, 3]


@pytest.mark.parametrize("value", [-1, mujoco.mjNGROUP, {"arm_collision": -1}])
def test_geom_group_validation(value):
  """GeomCfg should validate the group."""
  with pytest.raises(ValueError, match="group must be in"):
    cfg = GeomCfg(geom_names_expr=("test",), group=value)
    cfg.validate()


def test_geom_patches_collision_attrs(multi_geom_spec):
  """GeomCfg should patch collision attributes, leaving unset ones alone."""
  multi_geom_spec.geom("left_foot1_collision").contype = 2
  GeomCfg(
    geom_names_expr=(r"^(left|right)_foot\d_collision$",),
    condim=6,
    friction=(1.0, 5e-3, 5e-4),
  ).edit_spec(multi_geom_spec)

  left_foot = multi_geom_spec.geom("left_foot1_collision")
  assert left_foot.condim == 6
  assert left_foot.friction[0] == pytest.approx(1.0)
  assert left_foot.contype == 2  # Unset attribute untouched.

  arm = multi_geom_spec.geom("arm_collision")
  assert arm.condim == 3  # Unmatched geom untouched.


def test_geom_collision_attr_validation():
  """GeomCfg should validate collision attributes like CollisionCfg does."""
  with pytest.raises(ValueError, match="condim must be one of"):
    GeomCfg(geom_names_expr=("test",), condim=2).validate()
  with pytest.raises(ValueError, match="solmix must be in"):
    GeomCfg(geom_names_expr=("test",), solmix=1.5).validate()
