"""Blender headless probe deformation (mm, object-local).

    blender -b sensors/finger-xynova/r01/source/geometry.blend \\
        -P blender/scripts/probe_deform.py -- \\
        --job <abs/build/probe/finger-xynova/r01/_deform_job.json>

Reads the job written by ``orisys_sensor_design.deform``, resolves the mesh and
vertex group, builds the probe-average orthonormal basis, computes the
proportional falloff ONCE from the probe members' ORIGINAL positions (so the
weight vector is constant across every frame), then writes::

    probe_frames.npz      the deformation/rendering contract
    deform_report.json    measured probe statistics + honesty warnings

THE .BLEND IS NEVER WRITTEN. No ``bpy.ops.wm.save_mainfile`` /
``save_as_mainfile`` call exists in this script; the source file is opened
read-only and only read from. The script does leave Edit mode if the opened file
happens to be in it, but that changes the in-memory session only.

WHY NOT bpy.ops.transform.translate(use_proportional_edit=True)
--------------------------------------------------------------
That operator is viewport-modal. It needs a 3D-view region for the pivot and its
``tool_settings.proportional_size`` is interpreted in SCREEN SPACE (pixels)
relative to that region. Under ``blender -b`` there is no 3D-view region and no
screen-space scale, so the affected vertex set would depend on GUI state and
zoom and could not be reproduced headlessly. Every displacement here is computed
explicitly in millimetres by ``probe_falloff.py``::

    delta_i = weight_i * (axis_vector * step_mm)        (mm, object-local)

``weight_i = 1.0`` for the probe members, otherwise the falloff curve of
``d_i / radius_mm`` where ``d_i`` is the distance from the vertex's ORIGINAL
position to the nearest probe vertex. Do not "fix" this back to the operator.

Exit codes: 0 success, 1 scene/geometry error, 2 config/job error.
"""

# pyright: reportMissingImports=false

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import bpy
import numpy as np

# Blender -P does not consistently include the script directory on sys.path.
sys.path.insert(0, str(Path(__file__).resolve().parent))

from probe_falloff import (  # noqa: E402
    FALLOFF_MODES,
    affected_bands,
    edge_lengths_mm,
    falloff_weights,
    flatness_mm,
    orthonormal_basis,
    percentile,
)

JOB_SCHEMA = "orisys-sensor-design/probe-deform-job-v1"
REPORT_SCHEMA = "orisys-sensor-design/probe-deform-report-v1"
FRAMES_NPZ_FILENAME = "probe_frames.npz"
DEFORM_REPORT_FILENAME = "deform_report.json"

AXIS_ROWS = {"x": 0, "y": 1, "z": 2}
VALID_AXES = (None, "x", "y", "z")


class ConfigError(Exception):
    """Bad or incomplete job file (exit 2)."""


class SceneError(Exception):
    """Blend contents do not match the job (exit 1)."""


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


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(prog="probe_deform.py")
    parser.add_argument(
        "--job",
        type=Path,
        required=True,
        help="Absolute path to the _deform_job.json written by orisys-design deform",
    )
    return parser.parse_args(argv)


def load_job(path: Path) -> dict[str, Any]:
    """Load and structurally validate the deform job (raises ConfigError)."""
    if not path.is_file():
        raise ConfigError(f"Deform job not found: {path}")
    try:
        job = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise ConfigError(f"Deform job is not readable JSON: {path}: {exc}") from exc
    if not isinstance(job, dict):
        raise ConfigError(f"Deform job must be a JSON object: {path}")
    if job.get("schema") != JOB_SCHEMA:
        raise ConfigError(
            f"Deform job schema must be {JOB_SCHEMA!r}, got {job.get('schema')!r}"
        )

    for key in ("source_blend", "output_dir", "settings", "frames", "artifacts"):
        if key not in job:
            raise ConfigError(f"Deform job is missing {key!r}")

    objects = job.get("objects")
    if not isinstance(objects, dict) or not objects.get("mesh") or not objects.get("group"):
        raise ConfigError("Deform job 'objects' must name a mesh and a vertex group")

    settings = job["settings"]
    if not isinstance(settings, dict):
        raise ConfigError("Deform job 'settings' must be a mapping")
    for key in ("radius_mm", "falloff", "step_mm", "max_mm", "axes"):
        if key not in settings:
            raise ConfigError(f"Deform job settings is missing {key!r}")
    if not isinstance(settings["axes"], list) or not settings["axes"]:
        raise ConfigError("Deform job settings.axes must be a non-empty list")
    radius = float(settings["radius_mm"])
    if not radius > 0.0:
        raise ConfigError(f"Deform job settings.radius_mm must be > 0, got {radius!r}")
    falloff = str(settings["falloff"])
    if falloff not in FALLOFF_MODES:
        raise ConfigError(
            f"Deform job settings.falloff must be one of "
            f"{', '.join(FALLOFF_MODES)}, got {falloff!r}"
        )
    step = float(settings["step_mm"])
    max_mm = float(settings["max_mm"])
    if not step > 0.0 or not max_mm > 0.0:
        raise ConfigError("Deform job settings.step_mm/max_mm must be > 0")

    frames = job["frames"]
    if not isinstance(frames, list) or not frames:
        raise ConfigError("Deform job 'frames' must be a non-empty list")
    for position, frame in enumerate(frames):
        if not isinstance(frame, dict):
            raise ConfigError(f"frames[{position}] must be a mapping")
        for key in ("index", "name", "axis", "step_mm"):
            if key not in frame:
                raise ConfigError(f"frames[{position}] is missing {key!r}")
        if frame["index"] != position:
            raise ConfigError(
                f"frames[{position}].index must be {position}, got {frame['index']!r}"
            )
        if not isinstance(frame["name"], str) or not frame["name"]:
            raise ConfigError(f"frames[{position}].name must be a non-empty string")
        if frame["axis"] not in VALID_AXES:
            raise ConfigError(
                f"frames[{position}].axis must be null/x/y/z, got {frame['axis']!r}"
            )
        if float(frame["step_mm"]) < 0.0:
            raise ConfigError(f"frames[{position}].step_mm must be >= 0")
    if frames[0]["axis"] is not None or float(frames[0]["step_mm"]) != 0.0:
        raise ConfigError("frames[0] must be the zero-displacement reference")

    artifacts = job["artifacts"]
    if not isinstance(artifacts, dict) or not artifacts.get("frames") or not artifacts.get("report"):
        raise ConfigError("Deform job 'artifacts' must name 'frames' and 'report'")

    source_blend = Path(str(job["source_blend"]))
    if not source_blend.is_file():
        raise ConfigError(f"Deform job source_blend not found: {source_blend}")
    return job


def _sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _require_mesh(name: str):
    obj = bpy.data.objects.get(name)
    if obj is None:
        raise SceneError(f"Object not found: {name!r}")
    if obj.type != "MESH":
        raise SceneError(f"Object {name!r} has type {obj.type!r}, expected 'MESH'")
    if obj.data is None or len(obj.data.vertices) == 0:
        raise SceneError(f"Mesh object {name!r} has no vertices")
    return obj


def _ensure_object_mode(obj) -> None:
    """Leave Edit mode if the opened file is in it (in-memory only, never saved)."""
    if obj.mode == "OBJECT":
        return
    try:
        view_layer = bpy.context.view_layer
        if view_layer.objects.active is not obj:
            view_layer.objects.active = obj
        bpy.ops.object.mode_set(mode="OBJECT")
    except Exception as exc:  # pragma: no cover - depends on the opened file
        raise SceneError(f"Could not leave {obj.mode!r} mode for {obj.name!r}: {exc}") from exc


def _probe_members(obj, group_name: str) -> tuple[list[int], list[float]]:
    group = obj.vertex_groups.get(group_name)
    if group is None:
        available = ", ".join(sorted(g.name for g in obj.vertex_groups)) or "<none>"
        raise SceneError(
            f"Vertex group {group_name!r} not found on {obj.name!r} (have: {available})"
        )
    members: list[int] = []
    group_weights: list[float] = []
    for vertex in obj.data.vertices:
        for assignment in vertex.groups:
            if assignment.group == group.index:
                members.append(int(vertex.index))
                group_weights.append(float(assignment.weight))
                break
    if not members:
        raise SceneError(f"Vertex group {group_name!r} on {obj.name!r} is empty")
    return members, group_weights


def _write_frames_npz(
    path: Path,
    *,
    base_vertices: np.ndarray,
    vertices: np.ndarray,
    frame_names: list[str],
    frame_axes: list[str],
    frame_steps_mm: list[float],
    weights: np.ndarray,
    origin: np.ndarray,
    basis: np.ndarray,
    probe_indices: np.ndarray,
) -> None:
    if len(vertices) != len(frame_names):
        raise SceneError(
            f"frame count mismatch: {len(vertices)} vertex sets, {len(frame_names)} names"
        )
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez(
        path,
        base_vertices=np.asarray(base_vertices, dtype=np.float64),
        vertices=np.asarray(vertices, dtype=np.float64),
        frame_names=np.array(frame_names),
        frame_axes=np.array(frame_axes, dtype="<U4"),
        frame_steps_mm=np.asarray(frame_steps_mm, dtype=np.float64),
        weights=np.asarray(weights, dtype=np.float64),
        origin=np.asarray(origin, dtype=np.float64),
        basis=np.asarray(basis, dtype=np.float64),
        probe_indices=np.asarray(probe_indices, dtype=np.int32),
    )


def main() -> int:
    args = parse_args(_argv_after_double_dash())
    try:
        job = load_job(args.job.resolve())
    except ConfigError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2

    try:
        mesh_name = str(job["objects"]["mesh"])
        group_name = str(job["objects"]["group"])
        obj = _require_mesh(mesh_name)
        _ensure_object_mode(obj)

        settings = job["settings"]
        radius_mm = float(settings["radius_mm"])
        falloff = str(settings["falloff"])
        members, group_weights = _probe_members(obj, group_name)

        mesh = obj.data
        base_vertices = np.array(
            [vertex.co[:] for vertex in mesh.vertices], dtype=np.float64
        )
        normals = np.array(
            [vertex.normal[:] for vertex in mesh.vertices], dtype=np.float64
        )
        member_array = np.asarray(sorted(members), dtype=np.int64)
        origin, basis, basis_report = orthonormal_basis(
            normals[member_array], origin=base_vertices[member_array].mean(axis=0)
        )
        if not basis_report["orthonormal"]:
            raise SceneError(
                "probe-average basis is not orthonormal within "
                f"{1e-6:g} (error {basis_report['orthonormality_error']:.3g})"
            )

        # Weights come from the ORIGINAL coordinates, computed exactly once, so
        # every frame shares one falloff vector (Blender proportional-edit
        # semantics: distance is measured from the selected vertices' start).
        weights = falloff_weights(base_vertices, member_array, radius_mm, falloff)

        edges = np.array(
            [edge.vertices[:] for edge in mesh.edges], dtype=np.int64
        ).reshape(-1, 2)
        edge_lengths = edge_lengths_mm(base_vertices, edges)
        if edge_lengths.size == 0:
            raise SceneError(f"Mesh object {mesh_name!r} has no edges to measure")
        edge_p50_mm = percentile(edge_lengths, 50.0)
    except SceneError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    except Exception as exc:
        print(f"ERROR: probe deformation failed: {exc}", file=sys.stderr)
        return 1

    frames = job["frames"]
    frame_count = len(frames)
    vertices = np.empty((frame_count, base_vertices.shape[0], 3), dtype=np.float64)
    frame_deltas: list[float] = []
    frame_moved: list[int] = []
    for position, frame in enumerate(frames):
        axis = frame["axis"]
        if axis is None:
            vector = np.zeros(3, dtype=np.float64)
        else:
            vector = basis[AXIS_ROWS[axis]] * float(frame["step_mm"])
        vertices[position] = base_vertices + weights[:, None] * vector[None, :]
        delta = vertices[position] - base_vertices
        frame_deltas.append(float(np.linalg.norm(delta, axis=1).max()))
        # Contract: vertices displaced by a non-zero amount IN THIS FRAME.
        # 0 for the reference, 6 + partial (11 on Xynova) for any non-zero step.
        frame_moved.append(int(np.count_nonzero(np.any(delta != 0.0, axis=1))))

    if not np.array_equal(vertices[0], base_vertices):
        print(
            "ERROR: reference frame differs from the base vertices",
            file=sys.stderr,
        )
        return 1

    members_at_one = int((weights >= 1.0).sum())
    partial = int(((weights > 0.0) & (weights < 1.0)).sum())
    unaffected = int((weights == 0.0).sum())
    warnings: list[str] = []
    if partial < 4 * len(member_array):
        warnings.append(
            f"radius {radius_mm:g} mm reaches only {partial} vertices outside the "
            f"probe group (threshold {4 * len(member_array)} = 4 * "
            f"{len(member_array)} members); the deformation is local to the probe "
            "neighbourhood"
        )
    if radius_mm < edge_p50_mm:
        warnings.append(
            f"radius {radius_mm:g} mm is below the mesh median edge length "
            f"{edge_p50_mm:.6f} mm; affected vertices may be isolated"
        )
    if any(abs(weight - 1.0) > 1e-9 for weight in group_weights):
        warnings.append(
            f"probe group weights are not all 1.0 ({group_weights}); the contract "
            "displaces probe members by the full vector regardless"
        )
    scale = [float(value) for value in obj.matrix_world.to_scale()]
    if max(abs(value - 1.0) for value in scale) > 1e-4:
        warnings.append(
            f"mesh object scale {scale} is not 1; displacement units are "
            "object-local, not world millimetres"
        )

    source_blend = Path(str(job["source_blend"]))
    report = {
        "schema": REPORT_SCHEMA,
        "generated_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
        "source_blend": {
            "path": str(source_blend.resolve()),
            "sha256": _sha256_file(source_blend),
        },
        "mesh": mesh_name,
        "group": group_name,
        "probe_indices": [int(index) for index in member_array],
        "origin": [round(float(value), 9) for value in origin],
        "basis": [[round(float(value), 9) for value in row] for row in basis],
        "basis_orthonormal": bool(basis_report["orthonormal"]),
        "normal_angle_max_deg": round(float(basis_report["normal_angle_max_deg"]), 3),
        "flatness_mm": round(flatness_mm(base_vertices[member_array], origin, basis[2]), 6),
        "radius_mm": radius_mm,
        "falloff": falloff,
        "vertex_count": int(base_vertices.shape[0]),
        "edge_p50_mm": round(edge_p50_mm, 6),
        "weights": {
            "min": float(weights.min()),
            "max": float(weights.max()),
            "members_at_one": members_at_one,
            "partial": partial,
            "unaffected": unaffected,
        },
        "affect_bands": affected_bands(weights),
        "frames": [
            {
                "index": int(frame["index"]),
                "name": str(frame["name"]),
                "axis": frame["axis"],
                "step_mm": float(frame["step_mm"]),
                "max_delta_mm": round(frame_deltas[position], 9),
                # Vertices with a non-zero displacement in this frame: 0 for the
                # reference, 6 + partial (= 11 on Xynova) for any non-zero step.
                "moved_vertices": frame_moved[position],
            }
            for position, frame in enumerate(frames)
        ],
        "warnings": warnings,
    }

    artifacts = job["artifacts"]
    frames_path = Path(str(artifacts["frames"]))
    report_path = Path(str(artifacts["report"]))
    try:
        _write_frames_npz(
            frames_path,
            base_vertices=base_vertices,
            vertices=vertices,
            frame_names=[str(frame["name"]) for frame in frames],
            frame_axes=["null" if frame["axis"] is None else str(frame["axis"]) for frame in frames],
            frame_steps_mm=[float(frame["step_mm"]) for frame in frames],
            weights=weights,
            origin=origin,
            basis=basis,
            probe_indices=member_array.astype(np.int32),
        )
        report_path.parent.mkdir(parents=True, exist_ok=True)
        report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    except (OSError, SceneError) as exc:
        print(f"ERROR: could not write deform artifacts: {exc}", file=sys.stderr)
        return 1

    print(
        f"Probe deform complete: {frame_count} frames, {base_vertices.shape[0]} vertices, "
        f"{members_at_one} probe members, {partial} partial, {unaffected} unaffected"
    )
    for warning in warnings:
        print(f"WARNING: {warning}", file=sys.stderr)
    print(f"  {frames_path}")
    print(f"  {report_path}")
    return 0


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