#!/usr/bin/env python3
"""Structural regression checks for Kaptan's deterministic Gazebo SITL world.

This is deliberately an SDF/model-contract check, not a dynamics substitute.
It protects the reviewed boundary that the map is visual-only while the flat
ground and four landing feet provide the sole touchdown geometry.
"""

from __future__ import annotations

from pathlib import Path
import sys
import xml.etree.ElementTree as xml


WORLD = Path(__file__).with_name("worlds") / "kaptan-sitl-quad.sdf.in"


def fail(message: str) -> None:
    raise AssertionError(message)


def model(root: xml.Element, name: str) -> xml.Element:
    found = root.find(f"./world/model[@name='{name}']")
    if found is None:
        fail(f"missing model {name!r}")
    return found


def named(parent: xml.Element, tag: str, name: str) -> xml.Element:
    found = parent.find(f"./{tag}[@name='{name}']")
    if found is None:
        fail(f"missing {tag} {name!r}")
    return found


def pose(element: xml.Element) -> tuple[float, ...]:
    text = element.findtext("pose")
    if text is None:
        fail(f"{element.tag} {element.get('name')!r} has no pose")
    values = tuple(float(value) for value in text.split())
    if len(values) != 6:
        fail(f"{element.tag} {element.get('name')!r} pose does not have six fields")
    return values


def cylinder(element: xml.Element) -> tuple[float, float]:
    radius = element.findtext("geometry/cylinder/radius")
    length = element.findtext("geometry/cylinder/length")
    if radius is None or length is None:
        fail(f"{element.tag} {element.get('name')!r} is not a cylinder")
    return float(radius), float(length)


def vector(element: xml.Element, path: str, size: int) -> tuple[float, ...]:
    text = element.findtext(path)
    if text is None:
        fail(f"{element.tag} {element.get('name')!r} has no {path!r}")
    values = tuple(float(value) for value in text.split())
    if len(values) != size:
        fail(f"{element.tag} {element.get('name')!r} {path!r} has {len(values)} fields")
    return values


def check_world() -> None:
    root = xml.parse(WORLD).getroot()
    world = root.find("./world")
    if world is None:
        fail("missing world")
    model_names = {model_element.get("name") for model_element in world.findall("model")}
    if model_names != {"ground", "sitl_reference_map", "kaptan_sitl_quad"}:
        fail(
            "the declared flat-world model may contain only ground, the visual reference map, "
            f"and the quad; found {sorted(model_names)!r}"
        )

    ground = model(root, "ground")
    ground_collisions = ground.findall(".//collision")
    if [collision.get("name") for collision in ground_collisions] != ["ground_collision"]:
        fail("ground must retain exactly one flat collision plane")
    if ground_collisions[0].find("geometry/plane") is None:
        fail("ground collision must remain a flat plane")
    if vector(ground_collisions[0], "geometry/plane/normal", 3) != (0.0, 0.0, 1.0):
        fail("ground plane must remain horizontal with world +Z normal")
    if vector(ground_collisions[0], "geometry/plane/size", 2) != (100.0, 100.0):
        fail("ground plane must retain its declared 100 m square contact area")

    reference_map = model(root, "sitl_reference_map")
    if reference_map.findall(".//collision"):
        fail("reference map must remain visual-only and cannot add obstacle physics")
    if reference_map.findall(".//transparency"):
        fail("reference map must not use unsupported SDF transparency material children")
    reference_link = named(reference_map, "link", "reference_link")
    actual_marks = {visual.get("name") for visual in reference_link.findall("visual")}
    if actual_marks != {"landing_pad"}:
        fail(f"reference map may contain only the landing pad; found {sorted(actual_marks)!r}")

    quad = model(root, "kaptan_sitl_quad")
    if pose(quad)[:3] != (0.0, 0.0, 0.14):
        fail("quad base_link must start in its reviewed landing-foot rest pose")
    base_link = named(quad, "link", "base_link")
    body_collision = named(base_link, "collision", "body_collision")
    if vector(body_collision, "geometry/box/size", 3) != (0.30, 0.30, 0.08):
        fail("quad body collision must retain its reviewed clearance geometry")
    foot_positions = {
        "front_right": (0.12, -0.12),
        "rear_right": (-0.12, -0.12),
        "rear_left": (-0.12, 0.12),
        "front_left": (0.12, 0.12),
    }
    for foot, expected_xy in foot_positions.items():
        collision = named(base_link, "collision", f"landing_foot_{foot}_collision")
        visual = named(base_link, "visual", f"landing_foot_{foot}_visual")
        if pose(collision)[:3] != (*expected_xy, -0.10):
            fail(f"landing foot {foot!r} has an unexpected contact pose")
        if pose(visual) != pose(collision):
            fail(f"landing foot {foot!r} visual must match its collision geometry")
        if cylinder(collision) != (0.022, 0.08) or cylinder(visual) != (0.022, 0.08):
            fail(f"landing foot {foot!r} must retain its reviewed cylinder dimensions")

    expected_quad_collisions = {
        "body_collision",
        *(f"landing_foot_{foot}_collision" for foot in foot_positions),
    }
    actual_quad_collisions = {collision.get("name") for collision in base_link.findall("collision")}
    if actual_quad_collisions != expected_quad_collisions:
        fail(
            "quad collision geometry must remain its body and four declared landing feet; "
            f"found {sorted(actual_quad_collisions)!r}"
        )

    sensors = quad.findall(".//sensor")
    if [(sensor.get("name"), sensor.get("type")) for sensor in sensors] != [("imu", "imu")]:
        fail("this local-state SITL world must not imply unmodeled visual/range/flow sensors")


if __name__ == "__main__":
    try:
        check_world()
    except (AssertionError, OSError, xml.ParseError, ValueError) as error:
        print(f"SITL world structure check failed: {error}", file=sys.stderr)
        sys.exit(1)
    print("SITL world structure check passed")