#!/usr/bin/env python3
"""Capture matched L98/L99 tunnel change and passage views."""
from __future__ import annotations

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

ROOT = Path(__file__).resolve().parents[4]
OUT = ROOT / "Saved/Linux/l99/captures"
IMAGE_DIR = Path(__file__).resolve().parent / "captures"
DATA = ROOT / "Standalone/data/owner_world.json"
LAYOUT_PATH = ROOT / "Saved/Linux/l99/resolved-sw.json"
REACH_SUMMARY = ROOT / "Saved/Linux/l99/carve-reach/summary.json"
BINARIES = {
    "before-l98": ROOT / "Saved/Linux/l98/bin/VoxelForgeCoreHeadless-l98-clang19",
    "after-l99": ROOT / "Saved/Linux/l99/bin/VoxelForgeCoreHeadless-l99-clang",
}
RENDERER = ROOT / "Saved/Linux/l99/bin/VoxelForgeOffscreenCapture-l99-clang"
VOXEL_Q = 256.0
CORE_UM_PER_VOXEL = 250_000.0


def add(a, b):
    return tuple(x + y for x, y in zip(a, b))


def sub(a, b):
    return tuple(x - y for x, y in zip(a, b))


def mul(a, s):
    return tuple(x * s for x in a)


def dot(a, b):
    return sum(x * y for x, y in zip(a, b))


def norm(a):
    length = math.sqrt(dot(a, a))
    if length < 1.0e-8:
        raise ValueError("degenerate capture direction")
    return tuple(x / length for x in a)


def fmt(values):
    return ",".join(f"{value:.9f}" for value in values)


def angles(camera, target):
    delta = sub(target, camera)
    return (math.degrees(math.atan2(delta[0], delta[1])),
            math.degrees(math.atan2(delta[2], math.hypot(delta[0], delta[1]))))


def interpolate(points, distance):
    remaining = max(0.0, distance)
    for a, b in zip(points, points[1:]):
        delta = sub(b[0], a[0])
        length = math.sqrt(dot(delta, delta))
        if length < 1.0e-8:
            continue
        if remaining <= length:
            t = remaining / length
            return (tuple(a[0][i] + delta[i] * t for i in range(3)),
                    a[1] + (b[1] - a[1]) * t)
        remaining -= length
    return points[-1]


def route_frame(points, distance):
    center, radius = interpolate(points, distance)
    before, _ = interpolate(points, max(0.0, distance - 0.25))
    after, _ = interpolate(points, min(sum(
        math.sqrt(dot(sub(b[0], a[0]), sub(b[0], a[0])))
        for a, b in zip(points, points[1:])), distance + 0.25))
    tangent = norm(sub(after, before))
    horizontal = norm((tangent[0], tangent[1], 0.0))
    side = (-horizontal[1], horizontal[0], 0.0)
    return center, radius, tangent, side


def route_length(points):
    return sum(math.sqrt(dot(sub(b[0], a[0]), sub(b[0], a[0])))
               for a, b in zip(points, points[1:]))


def nearest_route(points, point):
    best = None
    distance_along = 0.0
    for a, b in zip(points, points[1:]):
        delta = sub(b[0], a[0])
        length_sq = dot(delta, delta)
        if length_sq < 1.0e-10:
            continue
        segment_length = math.sqrt(length_sq)
        t = max(0.0, min(1.0, dot(sub(point, a[0]), delta) / length_sq))
        center = add(a[0], mul(delta, t))
        tangent = norm(delta)
        radial = sub(point, center)
        radial = sub(radial, mul(tangent, dot(radial, tangent)))
        distance = math.sqrt(dot(sub(point, center), sub(point, center)))
        if best is None or distance < best[0]:
            radius = a[1] + (b[1] - a[1]) * t
            best = (distance, center, tangent, radial, radius,
                    distance_along + t * segment_length)
        distance_along += segment_length
    if best is None:
        raise ValueError("passage route has no non-degenerate segment")
    return best


def run(name, command, headless):
    log = OUT / f"{name}.log"
    wrapped = ["taskset", "-c", "0-11"]
    if headless:
        wrapped += ["prlimit", "--as=6442450944", "--"]
    wrapped += ["timeout", "1740s", *map(str, command)]
    started = 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() - started, 3),
           "log": str(log.relative_to(ROOT))}
    print("[L99CaptureRun] " + json.dumps(row, sort_keys=True), flush=True)
    if result.returncode != 0:
        raise RuntimeError(f"capture step failed: {name}; see {log}")
    return row


def nearest_surface_anchor(mesh, point):
    """Return the nearest surface vertex and outward mesh normal from a VFMSH scene."""
    best_distance_sq = float("inf")
    best_position = best_normal = None
    with mesh.open("rb") as stream:
        magic = stream.read(8)
        has_materials = magic == b"VFMSH002"
        if magic not in (b"VFMSH001", b"VFMSH002"):
            raise RuntimeError(f"unsupported mesh scene header: {mesh}")
        version, voxel_um, tile_count = struct.unpack("<3I", stream.read(12))
        if version != (2 if has_materials else 1):
            raise RuntimeError(f"mesh scene version mismatch: {mesh}")
        stream.seek(40, 1)  # stored camera position, yaw, and pitch
        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))
            extent = (32 << key[3]) * voxel_um / 1_000_000.0
            origin = (key[0] * extent, key[1] * extent, key[2] * extent)
            if any(point[axis] < origin[axis] - 20.0 or
                   point[axis] > origin[axis] + extent + 20.0 for axis in range(3)):
                stream.seek(vertex_count * 24 + (10 * vertex_count if has_materials else 0)
                            + index_count * 4 + mask_count, 1)
                continue
            positions = array("f")
            positions.frombytes(stream.read(vertex_count * 12))
            normals = array("f")
            normals.frombytes(stream.read(vertex_count * 12))
            if has_materials:
                stream.seek(vertex_count * 10, 1)
            stream.seek(index_count * 4 + mask_count, 1)
            if struct.pack("=I", 1) != struct.pack("<I", 1):
                positions.byteswap()
                normals.byteswap()
            for index in range(min(surface_count, vertex_count)):
                offset = index * 3
                # VFMSH stores vertex coordinates in world space; tile origin
                # is metadata for edge rendering and must not be added again.
                position = tuple(float(positions[offset + axis]) for axis in range(3))
                delta = sub(position, point)
                distance_sq = dot(delta, delta)
                if distance_sq >= best_distance_sq:
                    continue
                raw_normal = tuple(float(normals[offset + axis]) for axis in range(3))
                if dot(raw_normal, raw_normal) < 1.0e-10:
                    continue
                best_distance_sq = distance_sq
                best_position = position
                best_normal = norm(raw_normal)
    if best_position is None:
        raise RuntimeError(f"no nearby surface vertices in mesh scene: {mesh}")
    return {"position_m": best_position, "normal": best_normal,
            "distance_m": math.sqrt(best_distance_sq)}


def make_passages(layout):
    q_to_m = layout["core_voxel_size_um"] / (VOXEL_Q * 1_000_000.0)
    passages = []
    all_routes = []
    for passage in layout["passages"]:
        controls = [(tuple(float(q) * q_to_m for q in p["position_q"]),
                     float(p["radius_q"]) * q_to_m)
                    for p in passage["control_points"]]
        upper = tuple(float(q) * q_to_m for q in passage["upper_mouth_q"])
        lower = tuple(float(q) * q_to_m for q in passage["lower_mouth_q"])
        ui = min(range(len(controls)), key=lambda i: dot(sub(controls[i][0], upper),
                                                        sub(controls[i][0], upper)))
        li = min(range(len(controls)), key=lambda i: dot(sub(controls[i][0], lower),
                                                        sub(controls[i][0], lower)))
        points = controls[ui:li + 1]
        if len(points) < 2:
            raise RuntimeError(f"invalid mouth path: {passage['stable_id']}")
        upper_offset = math.sqrt(dot(sub(points[0][0], upper), sub(points[0][0], upper)))
        lower_offset = math.sqrt(dot(sub(points[-1][0], lower), sub(points[-1][0], lower)))
        length = route_length(points)
        passage_views = []
        for endpoint, distance, camera_distance, outward_sign in (
                ("upper-mouth", 0.0, min(length, 7.4), -1.0),
                ("lower-mouth", length, max(0.0, length - 5.5), 1.0)):
            center, _, tangent, side = route_frame(points, distance)
            camera_center, _, _, camera_side = route_frame(points, camera_distance)
            camera = (camera_center[0] + camera_side[0] * 2.5,
                      camera_center[1] + camera_side[1] * 2.5, camera_center[2])
            outward = mul(tangent, outward_sign)
            target = (center[0] + outward[0] * 8.0 + side[0] * 2.5,
                      center[1] + outward[1] * 8.0 + side[1] * 2.5, center[2])
            yaw, pitch = angles(camera, target)
            passage_views.append({"name": f"{passage['stable_id']}-{endpoint}",
                "kind": "mouth", "focus": center, "camera": camera, "target": target,
                "yaw": yaw, "pitch": pitch, "clip": None, "fov": 80.0,
                "focus_radius": 20.0, "passage": passage["stable_id"],
                "camera_location": "passage-air"})
        mid = length * 0.5
        center, radius, tangent, side = route_frame(points, mid)
        camera = (center[0] + side[0] * 0.75, center[1] + side[1] * 0.75,
                  center[2] - radius + 1.6)
        target_pos, target_radius = interpolate(points, min(length, mid + 14.0))
        target = (target_pos[0] + side[0] * 2.35, target_pos[1] + side[1] * 2.35,
                  target_pos[2] - target_radius + 1.8)
        yaw, pitch = angles(camera, target)
        passage_views.append({"name": f"{passage['stable_id']}-mid-tunnel",
            "kind": "mid-tunnel", "focus": center, "camera": camera,
            "target": target, "yaw": yaw, "pitch": pitch, "clip": None,
            "fov": 78.0, "focus_radius": 24.0, "passage": passage["stable_id"]})
        passages.append({"stable_id": passage["stable_id"], "route_length_m": length,
                         "control_points": len(points), "upper_mouth_anchor_offset_m": upper_offset,
                         "lower_mouth_anchor_offset_m": lower_offset, "views": passage_views})
        all_routes.append((passage["stable_id"], points))
    return passages, all_routes


def make_hotspot_views(routes, reach_summary):
    views = []
    for series, rows, value_key, location_key, region_key in (
        ("displacement", reach_summary["top_displacement_tiles"],
         "max_displacement_cm", "displacement_location_metres", "displacement_region"),
        ("normal", reach_summary["top_normal_tiles"],
         "max_normal_degrees", "normal_location_metres", "normal_region"),
    ):
        for index, row in enumerate(rows, 1):
            point = tuple(row[location_key])
            candidates = [(nearest_route(route, point), name, route)
                          for name, route in routes]
            (distance, center, tangent, radial, radius, along), nearest_id, route = min(
                candidates, key=lambda item: item[0][0])
            if math.sqrt(dot(radial, radial)) < 0.25:
                horizontal = norm((tangent[0], tangent[1], 0.0))
                radial = (-horizontal[1], horizontal[0], 0.0)
            outward = norm(radial)
            camera = add(point, mul(outward, 20.0))
            target = point
            plane = None
            yaw, pitch = angles(camera, target)
            views.append({"name": f"top-{series}-{index:02d}", "kind": f"top-{series}",
                "focus": point, "camera": camera, "target": target,
                "yaw": yaw, "pitch": pitch, "clip": plane, "fov": 68.0,
                "focus_radius": 18.0, "strate": row["strate"],
                "value": row[value_key], "region": row[region_key],
                "camera_anchor": "nearest-surface-l98-v5", "nearest_passage": nearest_id,
                "route_distance_m": distance})
    return views


def export_and_render(view, label, binary):
    safe = view["name"].replace("/", "-")
    stem = f"{safe}-{label}"
    mesh = OUT / f"{stem}.vfmesh"
    image = IMAGE_DIR / f"{safe}-{label}.png"
    camera, target = view["camera"], view["target"]
    yaw, pitch = angles(camera, target)
    export_args = [binary, DATA, "--world-assembler-prototype", "--stream-benchmark",
        "--stream-pipeline", "--workers=8", "--force-generate", "--clip-radius=3",
        "--max-clip-level=0", "--max-level=0", "--render-distance-chunks=0",
        f"--camera={fmt((*camera, yaw, pitch))}", f"--mesh-export={mesh}"]
    export_row = run(f"{stem}-export", export_args, headless=True)
    export_row["scene_bytes"] = mesh.stat().st_size if mesh.exists() else 0
    if export_row["scene_bytes"] <= 0:
        raise RuntimeError(f"empty mesh scene: {mesh}")
    if label == "before-l98" and view.get("camera_anchor") == "nearest-surface-l98-v5":
        anchor = nearest_surface_anchor(mesh, view["focus"])
        camera = add(anchor["position_m"], mul(anchor["normal"], 8.0))
        view["camera"] = camera
        view["target"] = anchor["position_m"]
        view["yaw"], view["pitch"] = angles(camera, view["target"])
        view["clip"] = None
        view["surface_anchor"] = anchor
        yaw, pitch = view["yaw"], view["pitch"]
    render_args = [RENDERER, mesh, image, "--size=1600x900",
        f"--fov-deg={view['fov']:.1f}", f"--camera-m={fmt(camera)}",
        f"--yaw-deg={yaw:.6f}", f"--pitch-deg={pitch:.6f}", "--shade=readable",
        "--headlamp=on", "--height-tint=on", "--tile-edges=off", "--tile-skirts=off",
        f"--height-center-m={view['focus'][2]:.6f}", "--height-range-m=10",
        f"--focus-region={fmt((*view['focus'], view['focus_radius']))}"]
    if view["clip"] is not None:
        render_args.append(f"--clip-plane={fmt(view['clip'])}")
    try:
        render_row = run(f"{stem}-render", render_args, headless=True)
    except RuntimeError:
        if view["clip"] is None:
            raise
        retry_args = [arg for arg in render_args
                      if not (isinstance(arg, str) and arg.startswith("--clip-plane="))]
        render_row = run(f"{stem}-render-no-clip-retry", retry_args, headless=True)
        render_row["clip_fallback"] = True
    render_row["image"] = str(image.relative_to(ROOT))
    render_row["image_bytes"] = image.stat().st_size if image.exists() else 0
    if render_row["image_bytes"] <= 0:
        raise RuntimeError(f"empty image: {image}")
    mesh.unlink()
    return export_row, render_row


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    IMAGE_DIR.mkdir(parents=True, exist_ok=True)
    for path in (*BINARIES.values(), RENDERER, DATA, LAYOUT_PATH, REACH_SUMMARY):
        if not path.is_file():
            raise SystemExit(f"missing input: {path}")
    layout = json.loads(LAYOUT_PATH.read_text(encoding="utf-8"))
    reach_summary = json.loads(REACH_SUMMARY.read_text(encoding="utf-8"))
    passages, routes = make_passages(layout)
    views = [view for passage in passages for view in passage["views"]]
    views += make_hotspot_views(routes, reach_summary)
    partial_path = OUT / "runs-partial.json"
    existing_records = {}
    if partial_path.is_file():
        for item in json.loads(partial_path.read_text(encoding="utf-8")):
            existing_records[item["view"]["name"]] = item
    records = []
    for view in views:
        prior = existing_records.get(view["name"])
        camera_matches = (prior is not None and
                          prior["view"].get("camera_anchor") == view.get("camera_anchor"))
        if camera_matches:
            for key in ("camera", "target", "yaw", "pitch", "clip", "surface_anchor"):
                if key in prior["view"]:
                    view[key] = prior["view"][key]
        if camera_matches and all(version in prior and
                         (ROOT / prior[version]["render"]["image"]).is_file()
                         for version in BINARIES):
            records.append(prior)
            print(f"[L99CaptureSkip] complete_pair={view['name']}", flush=True)
            continue
        pair = {"view": view}
        for label, binary in BINARIES.items():
            export_row, render_row = export_and_render(view, label, binary)
            pair[label] = {"export": export_row, "render": render_row}
            (OUT / "runs-partial.json").write_text(json.dumps(records + [pair], indent=2) + "\n",
                                                   encoding="utf-8")
        records.append(pair)
    summary = {"source_layout_identity": layout["stable_identity"],
        "view_count": len(views), "paired_png_count": len(views) * 2,
        "passages": passages, "views": records}
    (OUT / "summary.json").write_text(json.dumps(summary, indent=2) + "\n",
                                      encoding="utf-8")
    print("[L99CaptureSummary] " + json.dumps({
        "views": len(views), "paired_png_count": len(views) * 2,
        "images": [row[version]["render"]["image"] for row in records
                   for version in BINARIES]}, sort_keys=True), flush=True)
    return 0


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