"""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 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]:
    """Load output path, objects map, and optional housing collections."""
    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}"
            )
    return (str(output) if output else None), objects, housing


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 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]]:
    """Merge --job JSON with explicit flag overrides. Explicit flags win."""
    job_output: str | None = None
    job_objects: dict = {}
    job_housing: list[str] | 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 = _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


def main() -> int:
    args = parse_args(_argv_after_double_dash())
    try:
        output_dir, mesh_name, cam_name, origin_name, housing_collections = (
            _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
        )
        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())