"""Metric collision proxy and occlusion-aware virtual range sensor on Worker."""

import hashlib
import json
import math
import struct
from pathlib import Path

import numpy as np


def load_glb(path):
    data = Path(path).read_bytes()
    magic, version, size = struct.unpack_from("<III", data)
    if (magic, version, size) != (0x46546C67, 2, len(data)):
        raise ValueError("Expected a complete glTF 2 binary collision mesh")
    offset, document, binary = 12, None, None
    while offset < len(data):
        length, kind = struct.unpack_from("<II", data, offset)
        chunk = data[offset + 8 : offset + 8 + length]
        if kind == 0x4E4F534A:
            document = json.loads(chunk)
        elif kind == 0x004E4942:
            binary = chunk
        offset += 8 + length
    if document is None or binary is None:
        raise ValueError("Missing mesh buffers")

    def accessor(index):
        spec = document["accessors"][index]
        view = document["bufferViews"][spec["bufferView"]]
        if "sparse" in spec or view.get("buffer", 0) != 0:
            raise ValueError("Unsupported sparse/external collision buffer")
        dtype = {5123: "<u2", 5125: "<u4", 5126: "<f4"}[spec["componentType"]]
        width = {"SCALAR": 1, "VEC3": 3}[spec["type"]]
        item = np.dtype(dtype).itemsize
        start = view.get("byteOffset", 0) + spec.get("byteOffset", 0)
        return np.ndarray(
            (spec["count"], width),
            dtype=dtype,
            buffer=binary,
            offset=start,
            strides=(view.get("byteStride", width * item), item),
        ).copy()

    vertices, faces, base = [], [], 0
    # SplatTransform collision exports bake coordinates; reject transforms so
    # we never silently misregister contact geometry against the Gaussian view.
    for node in document.get("nodes", []):
        if any(k in node for k in ("matrix", "translation", "rotation", "scale")):
            raise ValueError("Collision node transform must be baked")
    for mesh in document["meshes"]:
        for primitive in mesh["primitives"]:
            if primitive.get("mode", 4) != 4:
                raise ValueError("Collision mesh must contain triangles")
            points = accessor(primitive["attributes"]["POSITION"])
            triangles = accessor(primitive["indices"]).reshape(-1, 3).astype(np.int32)
            if not np.isfinite(points).all() or triangles.max() >= len(points):
                raise ValueError("Invalid collision vertices")
            vertices.append(points[:, [0, 2, 1]] * np.array([1, -1, 1], np.float32))
            faces.append(triangles + base)
            base += len(points)
    return np.concatenate(vertices), np.concatenate(faces)


def ground_intersections(vertices, faces, x, y):
    """Vertical mesh intersections, used only to place the initial rigid body."""
    triangles = vertices[faces]
    low, high = triangles[:, :, :2].min(axis=1), triangles[:, :, :2].max(axis=1)
    candidate = triangles[
        (low[:, 0] <= x) & (high[:, 0] >= x) & (low[:, 1] <= y) & (high[:, 1] >= y)
    ]
    if not len(candidate):
        return np.empty(0)
    a, b, c = candidate[:, 0], candidate[:, 1], candidate[:, 2]
    den = (b[:, 1] - c[:, 1]) * (a[:, 0] - c[:, 0]) + (c[:, 0] - b[:, 0]) * (a[:, 1] - c[:, 1])
    valid = np.abs(den) > 1e-8
    a, b, c, den = a[valid], b[valid], c[valid], den[valid]
    u = ((b[:, 1] - c[:, 1]) * (x - c[:, 0]) + (c[:, 0] - b[:, 0]) * (y - c[:, 1])) / den
    v = ((c[:, 1] - a[:, 1]) * (x - c[:, 0]) + (a[:, 0] - c[:, 0]) * (y - c[:, 1])) / den
    inside = (u >= -1e-6) & (v >= -1e-6) & (u + v <= 1 + 1e-6)
    return (u * a[:, 2] + v * b[:, 2] + (1 - u - v) * c[:, 2])[inside]


def install_terrain(stage, manifest_path, world):
    from pxr import UsdGeom, UsdPhysics, UsdShade
    from rover_profile import PROFILE

    from k1link.simulation.ai_polygon.terrain_contract import terrain_matches

    manifest = json.loads(Path(manifest_path).read_text(encoding="utf-8-sig"))
    if not terrain_matches(manifest, world):
        raise ValueError("Terrain does not match the world calibration")
    path = Path(manifest["collider"])
    if hashlib.sha256(path.read_bytes()).hexdigest() != manifest["collider_sha256"]:
        raise ValueError("Collision mesh identity changed")
    points, triangles = load_glb(path)
    settings = world["settings"]
    heights = ground_intersections(points, triangles, *settings["spawn_xy"])
    candidates = heights[np.abs(heights - settings["ground_z"]) < 0.8]
    if not len(candidates):
        raise ValueError("No reconstructed ground at the configured spawn")
    # Place the whole footprint above nearby ground, without embedding a wheel
    # in a stone beside the centre ray. This is setup, never a planner input.
    angle = math.radians(settings["heading_degrees"])
    support, offsets = [], []
    for dx in (-0.5, 0, 0.5):
        for dy in (-0.5, 0, 0.5):
            x = settings["spawn_xy"][0] + dx * math.cos(angle) - dy * math.sin(angle)
            y = settings["spawn_xy"][1] + dx * math.sin(angle) + dy * math.cos(angle)
            intersections = ground_intersections(points, triangles, x, y)
            nearby = intersections[np.abs(intersections - settings["ground_z"]) < 0.8]
            if not len(nearby):
                raise ValueError("Rover footprint has no reconstructed support")
            support.append(float(nearby.max()))
            offsets.append([x - settings["spawn_xy"][0], y - settings["spawn_xy"][1], 1])
    plane = np.linalg.lstsq(np.asarray(offsets), np.asarray(support), rcond=None)[0]
    residual = np.asarray(support) - np.asarray(offsets) @ plane
    normal = np.array([-plane[0], -plane[1], 1.0])
    normal /= np.linalg.norm(normal)
    if np.max(np.abs(residual)) > 0.12 or normal[2] < math.cos(math.radians(25)):
        raise ValueError("Configured spawn is too uneven for the metre-wide rover")
    from spawn_clearance import obstructing_triangles

    if obstructing_triangles(
        points,
        triangles,
        settings["spawn_xy"],
        settings["heading_degrees"],
        plane,
        step=PROFILE["max_step_m"],
    ):
        raise ValueError("Configured spawn contains an obstacle inside the rover footprint")
    # Align to the local support plane before gravity settles the suspensionless
    # lab chassis. Wheel centres start 4 cm above the highest residual contact.
    wheel_center_z = float(
        plane[2] + (PROFILE["wheel_radius_m"] + 0.04) / normal[2] + max(0, residual.max())
    )
    mesh = UsdGeom.Mesh.Define(stage, "/World/Terrain")
    mesh.CreatePointsAttr(points.tolist())
    mesh.CreateFaceVertexCountsAttr([3] * len(triangles))
    mesh.CreateFaceVertexIndicesAttr(triangles.reshape(-1).tolist())
    mesh.CreateSubdivisionSchemeAttr(UsdGeom.Tokens.none)
    mesh.CreateDoubleSidedAttr(True)
    mesh.CreateVisibilityAttr(UsdGeom.Tokens.invisible)
    UsdPhysics.CollisionAPI.Apply(mesh.GetPrim())
    UsdPhysics.MeshCollisionAPI.Apply(mesh.GetPrim()).CreateApproximationAttr("none")
    material = UsdShade.Material.Define(stage, "/World/Materials/Terrain")
    surface = UsdPhysics.MaterialAPI.Apply(material.GetPrim())
    surface.CreateStaticFrictionAttr(0.9)
    surface.CreateDynamicFrictionAttr(0.8)
    surface.CreateRestitutionAttr(0)
    UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim()).Bind(
        material, UsdShade.Tokens.weakerThanDescendants, "physics"
    )
    return wheel_center_z, dict(
        manifest,
        vertex_count=len(points),
        triangle_count=len(triangles),
        initial_contact_z=float(plane[2]),
        initial_ground_normal=normal.tolist(),
        support_residual_m=float(np.max(np.abs(residual))),
    )


class RangeSensor:
    """Front 360-degree LiDAR plus rear near-field fan; first physical hit.

    Separate mounts observe the ground beyond each bumper without seeing
    through the chassis. The rear fan adds 660 rays to the 2160 front rays.
    """

    def __init__(self):
        import omni.physx

        self.query = omni.physx.get_physx_scene_query_interface()
        elevations = (-80, -75, -70, -60, -45, -35, -28, -22, -18, -14, -10, -5, 0, 5, 15)
        self.directions = [
            (math.cos(e) * math.cos(a), math.cos(e) * math.sin(a), math.sin(e))
            for e in np.radians(elevations)
            for a in np.radians(np.arange(-180, 180, 2.5))
        ]
        self.rear_directions = [
            (math.cos(e) * math.cos(a), math.cos(e) * math.sin(a), math.sin(e))
            for e in np.radians(np.arange(-80, 16, 5))
            for a in np.radians(np.arange(100, 261, 5))
        ]

    def capture(self, origin, rotation, rear_origin=None):
        from pxr import Gf

        points = []
        mounts = [(origin, self.directions)]
        if rear_origin is not None:
            mounts.append((rear_origin, self.rear_directions))
        for mount, directions in mounts:
            for direction in directions:
                ray = rotation.TransformDir(Gf.Vec3d(*direction)).GetNormalized()
                hit = self.query.raycast_closest(tuple(mount), tuple(ray), 8.0, bothSides=True)
                if hit["hit"] and not str(hit.get("rigidBody", "")).startswith("/World/Rover"):
                    points.append(hit["position"])
        return np.asarray(points, dtype=np.float32).reshape(-1, 3)