"""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]:
    """Load output path and objects map from a JSON export job."""
    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")
    return (str(output) if output else None), objects


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 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)",
    )
    return parser.parse_args(argv)


def _resolve_config(args: argparse.Namespace) -> tuple[Path, str, str, str]:
    """Merge --job JSON with explicit flag overrides. Explicit flags win."""
    job_output: str | None = None
    job_objects: dict = {}
    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 = _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"]
    )
    return Path(output_raw).resolve(), str(mesh), str(cam), str(origin)


def main() -> int:
    args = parse_args(_argv_after_double_dash())
    try:
        output_dir, mesh_name, cam_name, origin_name = _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"
    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)
        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"  {intrinsics_path.name}")
    print(f"  {pose_path.name}")
    return 0


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