"""Blender headless nominal compute export (YAML-free).

The host CLI (orisys-design export) reads design.yaml, writes a JSON job file,
and invokes this script. Blender only needs stdlib json — not PyYAML:

  blender -b sensors/finger/r01/source/geometry.blend -P blender/scripts/export.py -- \\
    --job <abs/build/export/finger/r01/_export_job.json>

Manual overrides (optional, win over job values):

  --output DIR --mesh sensing_surface --cam cam_sensor --origin origin_sensor

Writes only the compute payload (mesh.npz, camera_intrinsics.npz, world_to_cam.npy)
under the job's output directory. Metadata such as _export_job.json and
manifest.json live one level above the payload directory.
"""

# pyright: reportMissingImports=false

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path

import bmesh
import bpy
import numpy as np
from mathutils import Matrix

_DEFAULT_OBJECTS = {
    "sensing_surface": "sensing_surface",
    "cam_sensor": "cam_sensor",
    "origin_sensor": "origin_sensor",
}


def _argv_after_double_dash() -> list[str]:
    if "--" not in sys.argv:
        return []
    return sys.argv[sys.argv.index("--") + 1 :]


def _require_object(name: str, expected_type: str | None = None) -> bpy.types.Object:
    obj = bpy.data.objects.get(name)
    if obj is None:
        raise RuntimeError(f"Object not found: {name!r}")
    if expected_type is not None and obj.type != expected_type:
        raise RuntimeError(
            f"Object {name!r} has type {obj.type!r}, expected {expected_type!r}"
        )
    return obj


def _load_job(path: Path) -> tuple[str | None, dict, list[str] | None, dict | None]:
    """Load output path, objects map, housing collections and the optics stack."""
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"Export job must be a JSON object: {path}")
    objects = data.get("objects") or {}
    if not isinstance(objects, dict):
        raise ValueError(f"Export job 'objects' must be a mapping: {path}")
    output = data.get("output")
    housing = data.get("housing_collections")
    if housing is not None:
        if not isinstance(housing, list) or not all(
            isinstance(name, str) for name in housing
        ):
            raise ValueError(
                f"Export job 'housing_collections' must be a list of names: {path}"
            )
    optics = data.get("optics")
    if optics is not None:
        if not isinstance(optics, dict) or not isinstance(optics.get("layers"), list):
            raise ValueError(
                f"Export job 'optics' must be a mapping with 'layers': {path}"
            )
    return (str(output) if output else None), objects, housing, optics


def export_camera_intrinsics(camera_obj: bpy.types.Object, scene: bpy.types.Scene):
    """Return OpenCV-style K (3x3) and render width/height from Blender camera."""
    cam = camera_obj.data
    render = scene.render
    res_x = int(render.resolution_x)
    res_y = int(render.resolution_y)

    if cam.sensor_fit == "HORIZONTAL":
        sensor_width = float(cam.sensor_width)
        sensor_height = sensor_width / (res_x / res_y)
    else:
        sensor_height = float(cam.sensor_height)
        sensor_width = sensor_height * (res_x / res_y)
    focal = float(cam.lens)

    pixel_width = sensor_width / res_x
    pixel_height = sensor_height / res_y
    fx = focal / pixel_width
    fy = focal / pixel_height
    cx = res_x / 2.0
    cy = res_y / 2.0

    K = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=np.float64)
    return K, res_x, res_y


def _triangulated_mesh(
    obj: bpy.types.Object,
) -> tuple[bpy.types.Mesh, Matrix, object]:
    """Return evaluated triangulated mesh, world matrix, and cleanup callback."""
    depsgraph = bpy.context.evaluated_depsgraph_get()
    eval_obj = obj.evaluated_get(depsgraph)
    mesh = eval_obj.to_mesh()
    world = eval_obj.matrix_world.copy()
    bm = bmesh.new()
    try:
        bm.from_mesh(mesh)
        bmesh.ops.triangulate(bm, faces=bm.faces[:])
        bm.to_mesh(mesh)
    finally:
        bm.free()
    mesh.update()

    def _cleanup() -> None:
        eval_obj.to_mesh_clear()

    return mesh, world, _cleanup


def export_mesh_in_sensor_frame(
    obj: bpy.types.Object,
    t_sensor_from_world: Matrix,
    filepath: Path,
) -> tuple[int, int]:
    """Export triangulated mesh vertices in the sensor frame."""
    mesh, world, cleanup = _triangulated_mesh(obj)
    try:
        if not mesh.uv_layers.active:
            raise RuntimeError(f"Mesh {obj.name!r} has no active UV layer")

        t_sensor_from_object = t_sensor_from_world @ world

        vertices = np.array(
            [t_sensor_from_object @ v.co for v in mesh.vertices],
            dtype=np.float32,
        )
        faces_list = [tuple(poly.vertices) for poly in mesh.polygons]
        if not faces_list:
            raise RuntimeError(f"Mesh {obj.name!r} has no faces")
        if any(len(face) != 3 for face in faces_list):
            raise RuntimeError(f"Mesh {obj.name!r} is not fully triangulated after export")
        faces = np.asarray(faces_list, dtype=np.int32)

        uv_layer = mesh.uv_layers.active.data
        uv_data = np.array([uv_layer[i].uv[:] for i in range(len(uv_layer))], dtype=np.float32)
        loop_to_vertex = np.array(
            [mesh.loops[i].vertex_index for i in range(len(mesh.loops))],
            dtype=np.int32,
        )

        np.savez_compressed(
            filepath,
            vertices=vertices,
            faces=faces,
            uv_data=uv_data,
            loop_to_vertex=loop_to_vertex,
        )
        print(f"Mesh exported to {filepath} ({len(vertices)} verts, {len(faces)} tris)")
        return int(len(vertices)), int(len(faces))
    finally:
        cleanup()


def _housing_objects(collection_names: list[str]) -> list[bpy.types.Object]:
    """Unique mesh objects in the configured collections, in a stable order."""
    found: dict[str, bpy.types.Object] = {}
    for name in collection_names:
        collection = bpy.data.collections.get(name)
        if collection is None:
            print(
                f"WARNING: housing collection not found: {name!r}",
                file=sys.stderr,
            )
            continue
        for obj in collection.all_objects:
            if obj.type == "MESH" and obj.name not in found:
                found[obj.name] = obj
    return list(found.values())


def export_housing_in_sensor_frame(
    objects: list[bpy.types.Object],
    t_sensor_from_world: Matrix,
    filepath: Path,
) -> tuple[int, int, int]:
    """Export the housing meshes in the sensor frame, keeping part identity.

    Parts are concatenated into one buffer -- the same schema family as
    ``mesh.npz`` -- with ``part_index`` linking every face to a ``part_names``
    entry. Housing never goes into ``mesh.npz``: the SDK unroll reads that file's
    vertices as the sensing surface.
    """
    vertex_blocks: list[np.ndarray] = []
    face_blocks: list[np.ndarray] = []
    part_index: list[int] = []
    part_names: list[str] = []
    offset = 0

    for obj in objects:
        mesh, world, cleanup = _triangulated_mesh(obj)
        try:
            transform = t_sensor_from_world @ world
            faces_list = [tuple(poly.vertices) for poly in mesh.polygons]
            if not faces_list:
                print(f"WARNING: housing object {obj.name!r} has no faces; skipped")
                continue
            if any(len(face) != 3 for face in faces_list):
                raise RuntimeError(
                    f"Housing object {obj.name!r} is not fully triangulated after export"
                )
            vertex_blocks.append(
                np.array([transform @ v.co for v in mesh.vertices], dtype=np.float32)
            )
            faces = np.asarray(faces_list, dtype=np.int32) + offset
            face_blocks.append(faces)
            part_index.extend([len(part_names)] * len(faces))
            part_names.append(obj.name)
            offset += len(mesh.vertices)
        finally:
            cleanup()

    if vertex_blocks:
        vertices = np.concatenate(vertex_blocks)
        faces = np.concatenate(face_blocks)
    else:
        vertices = np.zeros((0, 3), dtype=np.float32)
        faces = np.zeros((0, 3), dtype=np.int32)
    names = (
        np.array(part_names, dtype=np.str_)
        if part_names
        else np.array([""], dtype=np.str_)[:0]
    )

    np.savez_compressed(
        filepath,
        vertices=vertices,
        faces=faces,
        part_index=np.asarray(part_index, dtype=np.int32),
        part_names=names,
    )
    print(
        f"Housing exported to {filepath} "
        f"({len(part_names)} part(s), {len(vertices)} verts, {len(faces)} tris)"
    )
    if not part_names:
        print(
            "WARNING: no housing meshes found; housing.npz is empty",
            file=sys.stderr,
        )
    return len(part_names), int(len(vertices)), int(len(faces))


def _mesh_world(obj: bpy.types.Object):
    """Triangulated mesh plus its world matrix, with a cleanup callback."""
    return _triangulated_mesh(obj)


def _measure_layer_thickness(
    outer_obj: bpy.types.Object, inner_obj: bpy.types.Object
) -> dict[str, float | int | None]:
    """Thickness of one layer, by casting from the outer surface onto the inner.

    Each outer vertex is sent inward along its negated vertex normal and the
    distance to the inner surface is recorded, which measures the real layer
    rather than the nearest-neighbour gap between two point clouds (that
    overestimates wherever the surfaces are oblique to each other).
    """
    from mathutils.bvhtree import BVHTree

    inner_mesh, inner_world, inner_cleanup = _mesh_world(inner_obj)
    try:
        inner_vertices = [inner_world @ vertex.co for vertex in inner_mesh.vertices]
        inner_faces = [tuple(poly.vertices) for poly in inner_mesh.polygons]
        if not inner_faces:
            return {"count": 0, "p05": None, "p50": None, "p95": None}
        tree = BVHTree.FromPolygons(inner_vertices, inner_faces, all_triangles=True)
    finally:
        inner_cleanup()

    outer_mesh, outer_world, outer_cleanup = _mesh_world(outer_obj)
    try:
        rotation = outer_world.to_3x3()
        distances: list[float] = []
        for vertex in outer_mesh.vertices:
            origin = outer_world @ vertex.co
            normal = rotation @ vertex.normal
            if normal.length <= 0.0:
                continue
            normal.normalize()
            hit = tree.ray_cast(origin - normal * 1e-4, -normal)
            location = hit[0]
            if location is None:
                continue
            distances.append(float((location - origin).length))
    finally:
        outer_cleanup()

    if not distances:
        return {"count": 0, "p05": None, "p50": None, "p95": None}
    distances.sort()

    def pick(fraction: float) -> float:
        position = (len(distances) - 1) * fraction
        low = int(math.floor(position))
        high = int(math.ceil(position))
        if low == high:
            return distances[low]
        return distances[low] + (distances[high] - distances[low]) * (position - low)

    return {
        "count": len(distances),
        "p05": pick(0.05),
        "p50": pick(0.50),
        "p95": pick(0.95),
    }


def export_optics_in_sensor_frame(
    layers: list[dict],
    t_sensor_from_world: Matrix,
    filepath: Path,
) -> int:
    """Export the refracting interfaces, in the sensor frame, one pair per layer."""
    payload: dict = {"layer_names": np.array([layer["name"] for layer in layers], dtype=np.str_)}
    thickness_p50: list[float] = []
    ior_values: list[float] = []

    for index, layer in enumerate(layers):
        if layer.get("solid"):
            targets = [("solid", _require_object(str(layer["solid"]), "MESH"))]
        else:
            targets = [
                ("outer", _require_object(str(layer["outer"]), "MESH")),
                ("inner", _require_object(str(layer["inner"]), "MESH")),
            ]
        for prefix, obj in targets:
            mesh, world, cleanup = _mesh_world(obj)
            try:
                transform = t_sensor_from_world @ world
                payload[f"layer{index}_{prefix}_vertices"] = np.array(
                    [transform @ vertex.co for vertex in mesh.vertices], dtype=np.float32
                )
                payload[f"layer{index}_{prefix}_faces"] = np.asarray(
                    [tuple(poly.vertices) for poly in mesh.polygons], dtype=np.int32
                )
            finally:
                cleanup()
        if layer.get("solid"):
            # Thickness of a solid: enter one wall and leave by the opposite one.
            measured = _measure_layer_thickness(targets[0][1], targets[0][1])
        else:
            measured = _measure_layer_thickness(targets[0][1], targets[1][1])
        declared = layer.get("thickness_mm")
        thickness = float(declared) if declared is not None else (measured["p50"] or 0.0)
        thickness_p50.append(thickness)
        ior_values.append(float(layer["ior"]))
        print(
            f"  optics layer {layer['name']!r}: ior={float(layer['ior']):g} "
            f"thickness={thickness:.4f} mm "
            f"({'declared' if declared is not None else 'measured p50'}), "
            f"measured p05={measured['p05']} p95={measured['p95']}"
        )

    payload["ior"] = np.asarray(ior_values, dtype=np.float64)
    payload["thickness_mm"] = np.asarray(thickness_p50, dtype=np.float64)
    np.savez_compressed(filepath, **payload)
    print(f"Optics exported to {filepath} ({len(layers)} layer(s))")
    return len(layers)


def export_camera_pose_sensor_frame(
    camera_obj: bpy.types.Object,
    t_sensor_from_world: Matrix,
    filepath: Path,
) -> None:
    """Save world_to_cam.npy where 'world' means the sensor frame (origin_sensor)."""
    t_sensor_from_cam = t_sensor_from_world @ camera_obj.matrix_world
    cam_from_sensor = np.linalg.inv(np.array(t_sensor_from_cam, dtype=np.float64))
    np.save(filepath, cam_from_sensor.astype(np.float32))
    print(f"Camera pose exported to {filepath}")


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(prog="export.py")
    parser.add_argument(
        "--job",
        type=Path,
        default=None,
        help="JSON job written by orisys-design (output + object names)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Absolute output directory (overrides job)",
    )
    parser.add_argument(
        "--mesh",
        default=None,
        help="Sensing surface mesh object name (overrides job)",
    )
    parser.add_argument(
        "--cam",
        default=None,
        help="Sensor camera object name (overrides job)",
    )
    parser.add_argument(
        "--origin",
        default=None,
        help="Sensor-frame empty object name (overrides job)",
    )
    parser.add_argument(
        "--housing-collections",
        dest="housing_collections",
        default=None,
        help="Comma-separated collections exported as housing (overrides job)",
    )
    return parser.parse_args(argv)


def _resolve_config(
    args: argparse.Namespace,
) -> tuple[Path, str, str, str, list[str], dict | None]:
    """Merge --job JSON with explicit flag overrides. Explicit flags win."""
    job_output: str | None = None
    job_objects: dict = {}
    job_housing: list[str] | None = None
    job_optics: dict | None = None
    if args.job is not None:
        job_path = args.job.resolve()
        if not job_path.is_file():
            raise RuntimeError(f"Export job not found: {job_path}")
        job_output, job_objects, job_housing, job_optics = _load_job(job_path)

    output_raw = args.output or job_output
    if not output_raw:
        raise RuntimeError("Provide --job <path.json> and/or --output <dir>")

    mesh = (
        args.mesh
        or job_objects.get("sensing_surface")
        or _DEFAULT_OBJECTS["sensing_surface"]
    )
    cam = args.cam or job_objects.get("cam_sensor") or _DEFAULT_OBJECTS["cam_sensor"]
    origin = (
        args.origin
        or job_objects.get("origin_sensor")
        or _DEFAULT_OBJECTS["origin_sensor"]
    )
    housing: list[str] = []
    if args.housing_collections:
        housing = [
            name.strip()
            for name in str(args.housing_collections).split(",")
            if name.strip()
        ]
    elif job_housing is not None:
        housing = list(job_housing)
    return (
        Path(output_raw).resolve(),
        str(mesh),
        str(cam),
        str(origin),
        housing,
        job_optics,
    )


def main() -> int:
    args = parse_args(_argv_after_double_dash())
    try:
        output_dir, mesh_name, cam_name, origin_name, housing_collections, optics = (
            _resolve_config(args)
        )
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2

    output_dir.mkdir(parents=True, exist_ok=True)

    try:
        mesh_obj = _require_object(mesh_name, "MESH")
        cam_obj = _require_object(cam_name, "CAMERA")
        origin_obj = _require_object(origin_name)
    except RuntimeError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1

    t_sensor_from_world = origin_obj.matrix_world.inverted()
    scene = bpy.context.scene

    mesh_path = output_dir / "mesh.npz"
    housing_path = output_dir / "housing.npz"
    intrinsics_path = output_dir / "camera_intrinsics.npz"
    pose_path = output_dir / "world_to_cam.npy"

    try:
        export_mesh_in_sensor_frame(mesh_obj, t_sensor_from_world, mesh_path)
        export_housing_in_sensor_frame(
            _housing_objects(housing_collections), t_sensor_from_world, housing_path
        )
        if optics and optics.get("layers"):
            export_optics_in_sensor_frame(
                list(optics["layers"]), t_sensor_from_world, output_dir / "optics.npz"
            )
        K, width, height = export_camera_intrinsics(cam_obj, scene)
        np.savez(
            intrinsics_path,
            K=K,
            width=np.int32(width),
            height=np.int32(height),
        )
        print(f"Intrinsics exported to {intrinsics_path} ({width}x{height})")
        export_camera_pose_sensor_frame(cam_obj, t_sensor_from_world, pose_path)
    except Exception as exc:
        print(f"ERROR: export failed: {exc}", file=sys.stderr)
        return 1

    print(f"Export complete: {output_dir}")
    print(f"  {mesh_path.name}")
    print(f"  {housing_path.name}")
    print(f"  {intrinsics_path.name}")
    print(f"  {pose_path.name}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())