#!/usr/bin/env python3
"""Reframe and verify the ten L99 passage-change hotspot captures."""
from __future__ import annotations

import json
import math
import struct
import subprocess
import sys
import time
import zlib
from array import array
from pathlib import Path

ROOT = Path(__file__).resolve().parents[4]
HERE = Path(__file__).resolve().parent
CAPTURES = HERE / "captures"
RUNS = ROOT / "Saved/Linux/l99/captures-fix"
LAYOUT = ROOT / "Saved/Linux/l99/resolved-sw.json"
REACH = ROOT / "Saved/Linux/l99/carve-reach/summary.json"
DATA = ROOT / "Standalone/data/owner_world.json"
RENDERER = ROOT / "Saved/Linux/l99/bin/VoxelForgeOffscreenCapture-l99-clang"
BINARIES = {
    "before-l98": ROOT / "Saved/Linux/l98/bin/VoxelForgeCoreHeadless-l98-clang19",
    "after-l99": ROOT / "Saved/Linux/l99/bin/VoxelForgeCoreHeadless-l99-clang",
}
sys.path.insert(0, str(HERE))
from RunL99Captures import angles, fmt, make_passages, nearest_route, norm, add, mul, sub  # noqa: E402


def screen_projection(camera, target, yaw_degrees, pitch_degrees, width=1600, height=900,
                      horizontal_fov=56.0):
    yaw, pitch = math.radians(yaw_degrees), math.radians(pitch_degrees)
    forward = (math.sin(yaw) * math.cos(pitch), math.cos(yaw) * math.cos(pitch),
               math.sin(pitch))
    right = (math.cos(yaw), -math.sin(yaw), 0.0)
    up = (right[1] * forward[2] - right[2] * forward[1],
          right[2] * forward[0] - right[0] * forward[2],
          right[0] * forward[1] - right[1] * forward[0])
    delta = sub(target, camera)
    depth = sum(delta[i] * forward[i] for i in range(3))
    x = sum(delta[i] * right[i] for i in range(3))
    y = sum(delta[i] * up[i] for i in range(3))
    tan_h = math.tan(math.radians(horizontal_fov) * 0.5)
    tan_v = tan_h / (width / height)
    return (0.5 + x / (2.0 * depth * tan_h),
            0.5 - y / (2.0 * depth * tan_v))


def run(name: str, command: list[str]) -> dict:
    log = RUNS / f"{name}.log"
    wrapped = ["taskset", "-c", "0-11", "prlimit", "--as=6442450944", "--",
               "timeout", "1740s", *map(str, command)]
    start = time.monotonic()
    with log.open("w", encoding="utf-8") as stream:
        result = subprocess.run(wrapped, cwd=ROOT, stdout=stream,
                                stderr=subprocess.STDOUT, check=False, timeout=1745)
    row = {"name": name, "exit_code": result.returncode,
           "wall_s": round(time.monotonic() - start, 3),
           "log": str(log.relative_to(ROOT))}
    print("[L99CaptureFixRun] " + json.dumps(row, sort_keys=True), flush=True)
    if result.returncode != 0:
        raise RuntimeError(f"capture step failed: {name}; see {log}")
    return row


def closest_on_triangle(p, a, b, c):
    """Ericson closest point, with barycentric weights for normal interpolation."""
    ab, ac, ap = sub(b, a), sub(c, a), sub(p, a)
    d1, d2 = sum(ab[i] * ap[i] for i in range(3)), sum(ac[i] * ap[i] for i in range(3))
    if d1 <= 0.0 and d2 <= 0.0:
        return a, (1.0, 0.0, 0.0)
    bp = sub(p, b)
    d3, d4 = sum(ab[i] * bp[i] for i in range(3)), sum(ac[i] * bp[i] for i in range(3))
    if d3 >= 0.0 and d4 <= d3:
        return b, (0.0, 1.0, 0.0)
    vc = d1 * d4 - d3 * d2
    if vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0:
        v = d1 / (d1 - d3)
        return add(a, mul(ab, v)), (1.0 - v, v, 0.0)
    cp = sub(p, c)
    d5, d6 = sum(ab[i] * cp[i] for i in range(3)), sum(ac[i] * cp[i] for i in range(3))
    if d6 >= 0.0 and d5 <= d6:
        return c, (0.0, 0.0, 1.0)
    vb = d5 * d2 - d1 * d6
    if vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0:
        w = d2 / (d2 - d6)
        return add(a, mul(ac, w)), (1.0 - w, 0.0, w)
    va = d3 * d6 - d5 * d4
    if va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0:
        w = (d4 - d3) / ((d4 - d3) + (d5 - d6))
        return add(b, mul(sub(c, b), w)), (0.0, 1.0 - w, w)
    denominator = 1.0 / (va + vb + vc)
    v, w = vb * denominator, vc * denominator
    return add(a, add(mul(ab, v), mul(ac, w))), (1.0 - v - w, v, w)


def nearest_surface(mesh: Path, point):
    """Find the closest non-skirt LOD2 triangle in the point's tile neighbourhood."""
    best = None
    target_key = tuple(math.floor(value / 32.0) for value in point)
    with mesh.open("rb") as stream:
        magic = stream.read(8)
        materials = magic == b"VFMSH002"
        if magic not in (b"VFMSH001", b"VFMSH002"):
            raise RuntimeError(f"unsupported mesh header {magic!r}: {mesh}")
        version, voxel_um, tile_count = struct.unpack("<3I", stream.read(12))
        if version != (2 if materials else 1):
            raise RuntimeError(f"mesh header version mismatch: {mesh}")
        stream.seek(40, 1)
        lod_counts = {}
        relevant_tiles = []
        for _ in range(tile_count):
            key = struct.unpack("<4i", stream.read(16))
            surface_count, _, vertex_count, index_count, mask_count = struct.unpack(
                "<5I", stream.read(20))
            lod_counts[key[3]] = lod_counts.get(key[3], 0) + 1
            positions = array("f")
            positions.frombytes(stream.read(vertex_count * 12))
            normals = array("f")
            normals.frombytes(stream.read(vertex_count * 12))
            if materials:
                stream.seek(vertex_count * 10, 1)
            indices = array("I")
            indices.frombytes(stream.read(index_count * 4))
            masks = stream.read(mask_count)
            if struct.pack("=I", 1) != struct.pack("<I", 1):
                positions.byteswap()
                normals.byteswap()
                indices.byteswap()
            if key[3] != 2 or any(abs(key[axis] - target_key[axis]) > 1
                                  for axis in range(3)):
                continue
            relevant_tiles.append(key)
            xyz = [tuple(float(positions[i * 3 + axis]) for axis in range(3))
                   for i in range(vertex_count)]
            nrm = [tuple(float(normals[i * 3 + axis]) for axis in range(3))
                   for i in range(vertex_count)]
            triangle_count = index_count // 3
            for tri in range(triangle_count):
                if masks and masks[tri] != 0:
                    continue
                ids = tuple(int(indices[tri * 3 + corner]) for corner in range(3))
                if any(index >= surface_count for index in ids):
                    continue
                hit, weights = closest_on_triangle(point, *(xyz[index] for index in ids))
                distance_sq = sum((hit[axis] - point[axis]) ** 2 for axis in range(3))
                if best is not None and distance_sq >= best[0]:
                    continue
                raw_normal = tuple(sum(weights[k] * nrm[ids[k]][axis] for k in range(3))
                                   for axis in range(3))
                length = math.sqrt(sum(value * value for value in raw_normal))
                if length < 1.0e-8:
                    continue
                best = (distance_sq, hit, tuple(value / length for value in raw_normal), key,
                        ids, weights)
    if best is None:
        raise RuntimeError(f"no LOD2 surface triangles near {point}; lods={lod_counts}")
    return {"distance_m": math.sqrt(best[0]), "position_m": list(best[1]),
            "normal": list(best[2]), "tile": list(best[3]),
            "triangle_vertices": list(best[4]), "barycentric": list(best[5]),
            "lod_counts": lod_counts, "nearby_lod2_tiles": len(relevant_tiles)}


def decode_rgba(path: Path) -> tuple[int, int, bytes]:
    command = ["taskset", "-c", "0-11", "prlimit", "--as=6442450944", "--",
               "timeout", "1740s", "ffmpeg", "-v", "error", "-i", str(path),
               "-frames:v", "1", "-f", "rawvideo", "-pix_fmt", "rgba", "pipe:1"]
    result = subprocess.run(command, cwd=ROOT, stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE, check=False, timeout=1745)
    if result.returncode != 0:
        raise RuntimeError(f"PNG decode failed for {path}: {result.stderr.decode(errors='replace')}")
    # ffmpeg rawvideo has no dimensions. Read the PNG IHDR separately for an exact size.
    with path.open("rb") as stream:
        if stream.read(8) != b"\x89PNG\r\n\x1a\n":
            raise RuntimeError(f"invalid PNG signature: {path}")
        stream.seek(16)
        width, height = struct.unpack(">II", stream.read(8))
    if len(result.stdout) != width * height * 4:
        raise RuntimeError(f"decoded RGBA size mismatch for {path}: {len(result.stdout)}")
    return width, height, result.stdout


def write_png(path: Path, width: int, height: int, rgba: bytes):
    def chunk(kind: bytes, data: bytes) -> bytes:
        payload = kind + data
        return struct.pack(">I", len(data)) + payload + struct.pack(">I", zlib.crc32(payload) & 0xffffffff)
    rows = bytearray()
    stride = width * 4
    for y in range(height):
        rows.append(0)
        rows.extend(rgba[y * stride:(y + 1) * stride])
    png = (b"\x89PNG\r\n\x1a\n"
           + chunk(b"IHDR", struct.pack(">2I5B", width, height, 8, 6, 0, 0, 0))
           + chunk(b"IDAT", zlib.compress(bytes(rows), 9))
           + chunk(b"IEND", b""))
    path.write_bytes(png)


def pixel_diff(before: Path, after: Path, output: Path, amplify: int = 64):
    width, height, a = decode_rgba(before)
    width_b, height_b, b = decode_rgba(after)
    if (width, height) != (width_b, height_b):
        raise RuntimeError(f"image dimensions differ: {before} and {after}")
    diff = bytearray(len(a))
    changed = 0
    channel_sum = 0
    max_delta = 0
    for offset in range(0, len(a), 4):
        dr = abs(a[offset] - b[offset])
        dg = abs(a[offset + 1] - b[offset + 1])
        db = abs(a[offset + 2] - b[offset + 2])
        if dr or dg or db:
            changed += 1
        channel_sum += dr + dg + db
        max_delta = max(max_delta, dr, dg, db)
        diff[offset] = min(255, dr * amplify)
        diff[offset + 1] = min(255, dg * amplify)
        diff[offset + 2] = min(255, db * amplify)
        diff[offset + 3] = 255
    write_png(output, width, height, bytes(diff))
    return {"width": width, "height": height, "changed_pixels": changed,
            "total_pixels": width * height, "changed_percent": round(100.0 * changed / (width * height), 5),
            "sum_abs_rgb": channel_sum, "max_channel_delta": max_delta,
            "diff_amplification": amplify}


def surface_description(region, outward_normal, route_height_delta, route_radius,
                        route_distance):
    nz = outward_normal[2]
    if region == 2:
        return "passage mouth surface"
    if region == 1:
        return "passage floor"
    if region == 0:
        if route_distance > route_radius + 2.0:
            if nz < -0.55:
                return "rock overhang / underside near passage"
            if abs(nz) < 0.35:
                return "other strata rock wall near passage"
            return "sloping strata surface near passage"
        if route_height_delta > 0.55 * route_radius:
            return "passage upper wall / roof"
        if route_height_delta < -0.55 * route_radius:
            return "passage floor"
        return "passage wall"
    if nz < -0.55:
        return "rock overhang / island underside outside the passage envelope"
    if nz > 0.7:
        return "upper terrain surface outside the passage envelope"
    if abs(nz) < 0.35:
        return "other strata rock wall outside the passage envelope"
    return "sloping strata surface outside the passage envelope"


def main():
    RUNS.mkdir(parents=True, exist_ok=True)
    CAPTURES.mkdir(parents=True, exist_ok=True)
    for path in (*BINARIES.values(), RENDERER, DATA, LAYOUT, REACH):
        if not path.is_file():
            raise SystemExit(f"missing input: {path}")
    layout = json.loads(LAYOUT.read_text(encoding="utf-8"))
    reach = json.loads(REACH.read_text(encoding="utf-8"))
    _, routes = make_passages(layout)
    sites = []
    for kind, rows, value_key, point_key, region_key in (
        ("displacement", reach["top_displacement_tiles"], "max_displacement_cm",
         "displacement_location_metres", "displacement_region"),
        ("normal", reach["top_normal_tiles"], "max_normal_degrees",
         "normal_location_metres", "normal_region"),
    ):
        for index, row in enumerate(rows, 1):
            point = tuple(float(v) for v in row[point_key])
            route, passage_id = min(((nearest_route(route_points, point), stable_id)
                                     for stable_id, route_points in routes),
                                    key=lambda candidate: candidate[0][0])
            sites.append({"name": f"top-{kind}-{index:02d}", "kind": kind,
                          "point": point, "strate": row["strate"],
                          "value": float(row[value_key]), "region": int(row[region_key]),
                          "nearest_passage": passage_id,
                          "passage_distance_m": float(route[0]),
                          "route_tangent": list(route[2]),
                          "route_height_delta_m": point[2] - route[1][2],
                          "route_radius_m": float(route[4]),
                          "row": row})

    records = []
    for site in sites:
        point = site["point"]
        # Place the export camera far enough away that the audited location is in the LOD2 shell.
        # The offscreen renderer then moves back to the exact changed surface for a close view.
        direction = -1.0 if point[0] >= 0.0 else 1.0
        export_camera = (point[0] + direction * 80.0, point[1], point[2])
        yaw_export, pitch_export = 0.0, 0.0
        meshes = {}
        exports = {}
        for label, binary in BINARIES.items():
            mesh = RUNS / f"{site['name']}-{label}.vfmesh"
            exports[label] = run(f"{site['name']}-{label}-export", [binary, DATA,
                "--world-assembler-prototype", "--stream-benchmark", "--stream-pipeline",
                "--workers=8", "--force-generate", "--clip-radius=3",
                "--max-clip-level=2", "--max-level=2", "--render-distance-chunks=0",
                f"--camera={fmt((*export_camera, yaw_export, pitch_export))}",
                f"--mesh-export={mesh}"])
            if not mesh.is_file() or mesh.stat().st_size <= 0:
                raise RuntimeError(f"empty exported scene for {site['name']} {label}")
            meshes[label] = mesh
        anchor = nearest_surface(meshes["before-l98"], point)
        after_anchor = nearest_surface(meshes["after-l99"], point)
        expected_tile = list(site["row"]["tile"])
        if anchor["tile"] != expected_tile or after_anchor["tile"] != expected_tile:
            raise RuntimeError(f"change tile not represented by the LOD2 export for {site['name']}: "
                               f"expected={expected_tile} before={anchor['tile']} "
                               f"after={after_anchor['tile']}")
        if anchor["distance_m"] > 0.1 or after_anchor["distance_m"] > 0.1:
            raise RuntimeError(f"change point too far from exported surface for {site['name']}: "
                               f"before={anchor['distance_m']:.3f}m "
                               f"after={after_anchor['distance_m']:.3f}m")
        camera_distance = 3.25
        # Marching-cubes normals point into the solid in these exports, so the air-side view
        # starts along their negative direction.
        air_normal = mul(tuple(anchor["normal"]), -1.0)
        camera_direction = air_normal
        if site["kind"] == "displacement":
            route_tangent = tuple(site["route_tangent"])
            tangent = sub(route_tangent, mul(air_normal,
                sum(route_tangent[i] * air_normal[i] for i in range(3))))
            tangent = norm(tangent)
            camera_direction = norm(add(mul(air_normal, 0.70710678),
                                        mul(tangent, 0.70710678)))
        camera = add(tuple(anchor["position_m"]), mul(camera_direction, camera_distance))
        yaw, pitch = angles(camera, point)
        actual_camera_distance = math.sqrt(sum((camera[axis] - point[axis]) ** 2
                                               for axis in range(3)))
        if actual_camera_distance > 5.0 or actual_camera_distance < 2.0:
            raise RuntimeError(f"changed point outside close framing range for {site['name']}: "
                               f"camera-to-point={actual_camera_distance:.3f}m; "
                               f"surface-offset={anchor['distance_m']:.3f}m")
        target_uv = screen_projection(camera, point, yaw, pitch)
        projection_error_px = math.hypot((target_uv[0] - 0.5) * 1600,
                                         (target_uv[1] - 0.5) * 900)
        if projection_error_px > 0.01:
            raise RuntimeError(f"changed point is off-axis for {site['name']}: {target_uv}")
        view = {**site, "focus": list(point), "camera": list(camera),
                "target": list(point), "yaw": yaw, "pitch": pitch,
                "target_screen_uv": list(target_uv),
                "target_projection_error_px": projection_error_px,
                "change_tile_lod2": expected_tile,
                "surface_match_distance_m": anchor["distance_m"],
                "camera_to_surface_m": camera_distance,
                "camera_to_changed_point_m": actual_camera_distance,
                "camera_side": "air side, 45-degree passage-tangent view"
                    if site["kind"] == "displacement" else "air side (negative exported normal)",
                "surface_anchor": anchor, "after_surface_anchor": after_anchor,
                "surface_type": surface_description(site["region"],
                    tuple(-value for value in anchor["normal"]),
                    site["route_height_delta_m"], site["route_radius_m"],
                    site["passage_distance_m"]),
                "export_camera": list(export_camera), "export_lod_counts": anchor["lod_counts"]}
        render_records = {}
        render_shade = "normal" if site["kind"] == "normal" else "readable"
        view["render_shade"] = render_shade
        for label in BINARIES:
            suffix = "before-l98" if label == "before-l98" else "after-l99"
            image = CAPTURES / f"{site['name']}-{suffix}.png"
            render_args = [RENDERER, meshes[label], image, "--size=1600x900",
                "--fov-deg=56", f"--camera-m={fmt(camera)}", f"--yaw-deg={yaw:.8f}",
                f"--pitch-deg={pitch:.8f}", f"--shade={render_shade}", "--headlamp=on",
                "--height-tint=on", "--tile-edges=off", "--tile-skirts=off",
                f"--height-center-m={point[2]:.6f}", "--height-range-m=8",
                f"--focus-region={fmt((*point, 10.0))}"]
            render_records[label] = run(f"{site['name']}-{label}-render", render_args)
            render_records[label]["image"] = str(image.relative_to(ROOT))
            render_records[label]["image_bytes"] = image.stat().st_size
        for mesh in meshes.values():
            mesh.unlink()
        diff_path = CAPTURES / f"{site['name']}-diff-amplified.png"
        pixel_stats = pixel_diff(CAPTURES / f"{site['name']}-before-l98.png",
                                 CAPTURES / f"{site['name']}-after-l99.png", diff_path)
        record = {"view": view, "before": render_records["before-l98"],
                  "after": render_records["after-l99"],
                  "diff": str(diff_path.relative_to(ROOT)), "pixel_diff": pixel_stats,
                  "exports": exports}
        records.append(record)
        print("[L99CaptureFixSite] " + json.dumps({"name": site["name"],
              "lod_counts": anchor["lod_counts"], "surface_distance_m": anchor["distance_m"],
              "camera_distance_m": camera_distance,
              "changed_pixels": pixel_stats["changed_pixels"],
              "surface_type": view["surface_type"], "passage": site["nearest_passage"],
              "passage_distance_m": site["passage_distance_m"]}, sort_keys=True), flush=True)
    summary = {"source_layout_identity": layout["stable_identity"],
               "capture_lod": 2, "image_size": [1600, 900],
               "camera_surface_distance_m": 3.25,
               "sites": records}
    (RUNS / "fixed-summary.json").write_text(json.dumps(summary, indent=2) + "\n",
                                              encoding="utf-8")
    print("[L99CaptureFixSummary] " + json.dumps({
        "sites": len(records),
        "changed_pixels": {r["view"]["name"]: r["pixel_diff"]["changed_pixels"]
                           for r in records},
        "fixed_summary": str((RUNS / "fixed-summary.json").relative_to(ROOT))},
        sort_keys=True), flush=True)


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