"""Bake deterministic, area-weighted surface samples from a single-mesh GLB.

Output layout is little-endian float32, eight floats per point:
position xyz, normal xyz, random seed, surface rank.
"""

from __future__ import annotations

import argparse
import json
import math
import struct
from bisect import bisect_left
from pathlib import Path

import numpy as np


COMPONENT_DTYPES = {
    5123: np.dtype("<u2"),
    5125: np.dtype("<u4"),
    5126: np.dtype("<f4"),
}


def load_glb(path: Path):
    data = path.read_bytes()
    magic, version, total_length = struct.unpack_from("<III", data, 0)
    if magic != 0x46546C67 or version != 2 or total_length != len(data):
        raise ValueError("Expected a valid GLB 2.0 file")

    json_chunk = None
    binary_chunk = None
    offset = 12
    while offset < len(data):
        length, chunk_type = struct.unpack_from("<II", data, offset)
        offset += 8
        chunk = data[offset : offset + length]
        offset += length
        if chunk_type == 0x4E4F534A:
            json_chunk = json.loads(chunk.rstrip(b"\x00 ").decode("utf-8"))
        elif chunk_type == 0x004E4942:
            binary_chunk = chunk

    if json_chunk is None or binary_chunk is None:
        raise ValueError("GLB must contain JSON and BIN chunks")
    return json_chunk, binary_chunk


def read_accessor(document, binary, accessor_index):
    accessor = document["accessors"][accessor_index]
    view = document["bufferViews"][accessor["bufferView"]]
    dtype = COMPONENT_DTYPES[accessor["componentType"]]
    widths = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4}
    width = widths[accessor["type"]]
    start = view.get("byteOffset", 0) + accessor.get("byteOffset", 0)
    stride = view.get("byteStride", dtype.itemsize * width)
    count = accessor["count"]

    if stride == dtype.itemsize * width:
        result = np.frombuffer(binary, dtype=dtype, count=count * width, offset=start)
        return result.reshape(count, width)

    result = np.empty((count, width), dtype=dtype)
    for index in range(count):
        result[index] = np.frombuffer(
            binary,
            dtype=dtype,
            count=width,
            offset=start + index * stride,
        )
    return result


def sample_surface(positions, normals, indices, count, seed):
    triangles = indices.reshape(-1, 3)
    a = positions[triangles[:, 0]]
    b = positions[triangles[:, 1]]
    c = positions[triangles[:, 2]]
    areas = np.linalg.norm(np.cross(b - a, c - a), axis=1) * 0.5
    cumulative = np.cumsum(areas, dtype=np.float64)
    if cumulative[-1] <= 0:
        raise ValueError("Mesh has no positive-area triangles")

    rng = np.random.default_rng(seed)
    picks = rng.random(count) * cumulative[-1]
    triangle_ids = np.searchsorted(cumulative, picks, side="left")
    chosen = triangles[triangle_ids]

    r1 = np.sqrt(rng.random(count)).astype(np.float32)
    r2 = rng.random(count).astype(np.float32)
    w0 = 1.0 - r1
    w1 = r1 * (1.0 - r2)
    w2 = r1 * r2

    sampled_positions = (
        positions[chosen[:, 0]] * w0[:, None]
        + positions[chosen[:, 1]] * w1[:, None]
        + positions[chosen[:, 2]] * w2[:, None]
    )
    sampled_normals = (
        normals[chosen[:, 0]] * w0[:, None]
        + normals[chosen[:, 1]] * w1[:, None]
        + normals[chosen[:, 2]] * w2[:, None]
    )
    lengths = np.linalg.norm(sampled_normals, axis=1)
    sampled_normals /= np.maximum(lengths[:, None], 1e-7)

    random_seed = rng.random(count, dtype=np.float32)
    # Sorting by a second random value keeps every draw-range density spatially uniform.
    rank = rng.random(count, dtype=np.float32)
    order = np.argsort(rank)
    baked = np.column_stack(
        (
            sampled_positions,
            sampled_normals,
            random_seed,
            rank,
        )
    ).astype("<f4", copy=False)
    return baked[order], float(cumulative[-1]), len(triangles)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--count", type=int, default=220_000)
    parser.add_argument("--seed", type=int, default=240824)
    args = parser.parse_args()

    document, binary = load_glb(args.input)
    primitive = document["meshes"][0]["primitives"][0]
    positions = read_accessor(document, binary, primitive["attributes"]["POSITION"]).astype(
        np.float32
    )
    normals = read_accessor(document, binary, primitive["attributes"]["NORMAL"]).astype(
        np.float32
    )
    indices = read_accessor(document, binary, primitive["indices"]).reshape(-1)
    baked, area, triangles = sample_surface(
        positions, normals, indices, args.count, args.seed
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    baked.tofile(args.output)
    print(
        json.dumps(
            {
                "points": args.count,
                "triangles": triangles,
                "surface_area": area,
                "stride_floats": 8,
                "bytes": args.output.stat().st_size,
            },
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()