#!/usr/bin/env python3
"""Convert VoxelForge VFMSH001/VFMSH002 mesh scenes to binary glTF (GLB)."""

import argparse
import json
import math
import struct
import time
from pathlib import Path


HEADER_SIZE = 60
TILE_HEADER_SIZE = 36
MAX_TILES = 1_000_000
MAX_VERTICES_PER_TILE = 10_000_000
MAX_INDICES_PER_TILE = 60_000_000
MAX_MASKS_PER_TILE = 20_000_000


def read_exact(stream, size, label):
    data = stream.read(size)
    if len(data) != size:
        raise ValueError(f"truncated {label}: expected {size} bytes, got {len(data)}")
    return data


def convert(source, destination):
    source = Path(source)
    destination = Path(destination)
    start = time.perf_counter()
    positions = bytearray()
    normals = bytearray()
    indices = bytearray()
    mins = [math.inf, math.inf, math.inf]
    maxs = [-math.inf, -math.inf, -math.inf]
    total_vertices = 0
    total_indices = 0
    minimum_index = 0xFFFFFFFF
    maximum_index = 0
    nonempty_tiles = 0

    with source.open("rb") as stream:
        header = read_exact(stream, HEADER_SIZE, "scene header")
        magic = header[:8]
        if magic not in (b"VFMSH001", b"VFMSH002"):
            raise ValueError(f"{source}: unsupported VF mesh magic {magic!r}")
        version, voxel_um, tile_count = struct.unpack_from("<3I", header, 8)
        expected_version = 1 if magic == b"VFMSH001" else 2
        if version != expected_version or voxel_um == 0 or tile_count > MAX_TILES:
            raise ValueError(f"{source}: invalid scene header")
        camera_xyz = struct.unpack_from("<3d", header, 20)
        yaw, pitch = struct.unpack_from("<2d", header, 44)
        if not all(math.isfinite(value) for value in (*camera_xyz, yaw, pitch)):
            raise ValueError(f"{source}: non-finite camera metadata")

        for tile_number in range(tile_count):
            record = read_exact(stream, TILE_HEADER_SIZE, f"tile {tile_number} header")
            key = struct.unpack_from("<4i", record, 0)
            surface_vertices, skirt_triangles, vertex_count, index_count, mask_count = (
                struct.unpack_from("<5I", record, 16)
            )
            if (key[3] < 0 or key[3] > 8 or vertex_count > MAX_VERTICES_PER_TILE
                    or index_count > MAX_INDICES_PER_TILE or index_count % 3
                    or mask_count > MAX_MASKS_PER_TILE
                    or (mask_count and mask_count != index_count // 3)):
                raise ValueError(f"{source}: malformed tile {tile_number} header")

            position_data = read_exact(stream, vertex_count * 12, f"tile {tile_number} positions")
            normal_data = read_exact(stream, vertex_count * 12, f"tile {tile_number} normals")
            if version == 2:
                read_exact(stream, vertex_count * 10, f"tile {tile_number} materials")
            index_data = read_exact(stream, index_count * 4, f"tile {tile_number} indices")
            read_exact(stream, mask_count, f"tile {tile_number} skirt masks")

            if index_count:
                nonempty_tiles += 1
            for index, in struct.iter_unpack("<I", index_data):
                if index >= vertex_count:
                    raise ValueError(f"{source}: tile {tile_number} index outside vertex array")
                minimum_index = min(minimum_index, index + total_vertices)
                maximum_index = max(maximum_index, index + total_vertices)
                indices.extend(struct.pack("<I", index + total_vertices))

            for (x, y, z), (nx, ny, nz) in zip(
                    struct.iter_unpack("<3f", position_data),
                    struct.iter_unpack("<3f", normal_data)):
                if not all(math.isfinite(value) for value in (x, y, z, nx, ny, nz)):
                    raise ValueError(f"{source}: non-finite vertex in tile {tile_number}")
                # glTF uses Y-up; rotate VoxelForge's Z-up coordinates -90 degrees about X.
                px, py, pz = x, z, -y
                nnx, nny, nnz = nx, nz, -ny
                positions.extend(struct.pack("<3f", px, py, pz))
                normals.extend(struct.pack("<3f", nnx, nny, nnz))
                mins[0] = min(mins[0], px)
                mins[1] = min(mins[1], py)
                mins[2] = min(mins[2], pz)
                maxs[0] = max(maxs[0], px)
                maxs[1] = max(maxs[1], py)
                maxs[2] = max(maxs[2], pz)
            total_vertices += vertex_count
            total_indices += index_count
            if total_vertices > 0xFFFFFFFF:
                raise ValueError(f"{source}: combined vertex count exceeds glTF uint32 indices")
            if surface_vertices > vertex_count or skirt_triangles > index_count // 3:
                raise ValueError(f"{source}: inconsistent tile {tile_number} counts")

        if stream.read(1):
            raise ValueError(f"{source}: trailing bytes after final tile")

    if total_vertices == 0 or total_indices == 0:
        raise ValueError(f"{source}: scene contains no indexed geometry")

    bin_data = bytearray()
    views = []

    def append_view(data, target):
        while len(bin_data) % 4:
            bin_data.append(0)
        offset = len(bin_data)
        bin_data.extend(data)
        view = {"buffer": 0, "byteOffset": offset, "byteLength": len(data)}
        if target is not None:
            view["target"] = target
        views.append(view)
        return len(views) - 1

    position_view = append_view(positions, 34962)
    normal_view = append_view(normals, 34962)
    index_view = append_view(indices, 34963)
    camera = [camera_xyz[0], camera_xyz[2], -camera_xyz[1]]
    document = {
        "asset": {"version": "2.0", "generator": "VoxelForge webviewer converter"},
        "scene": 0,
        "scenes": [{"nodes": [0], "extras": {
            "voxelForge": {
                "source": source.name,
                "voxelSizeMicrons": voxel_um,
                "tiles": tile_count,
                "nonemptyTiles": nonempty_tiles,
                "triangles": total_indices // 3,
                "cameraMeters": camera,
                "yawDegrees": yaw,
                "pitchDegrees": pitch,
                "upAxis": "Y",
                "units": "metres",
            }
        }}],
        "nodes": [{"mesh": 0, "name": source.stem}],
        "meshes": [{"name": source.stem, "primitives": [{
            "attributes": {"POSITION": 0, "NORMAL": 1},
            "indices": 2,
            "material": 0,
            "mode": 4,
        }]}],
        "materials": [{"name": "VoxelForge surface", "pbrMetallicRoughness": {
            "baseColorFactor": [0.82, 0.82, 0.78, 1.0],
            "metallicFactor": 0.0,
            "roughnessFactor": 0.9,
        }}],
        "buffers": [{"byteLength": len(bin_data)}],
        "bufferViews": views,
        "accessors": [
            {"bufferView": position_view, "componentType": 5126, "count": total_vertices,
             "type": "VEC3", "min": mins, "max": maxs},
            {"bufferView": normal_view, "componentType": 5126, "count": total_vertices,
             "type": "VEC3"},
            {"bufferView": index_view, "componentType": 5125, "count": total_indices,
             "type": "SCALAR", "min": [minimum_index], "max": [maximum_index]},
        ],
    }
    json_chunk = json.dumps(document, separators=(",", ":"), allow_nan=False).encode("utf-8")
    json_chunk += b" " * ((-len(json_chunk)) % 4)
    bin_data.extend(b"\0" * ((-len(bin_data)) % 4))
    total_length = 12 + 8 + len(json_chunk) + 8 + len(bin_data)
    glb = (struct.pack("<4sII", b"glTF", 2, total_length)
           + struct.pack("<I4s", len(json_chunk), b"JSON") + json_chunk
           + struct.pack("<I4s", len(bin_data), b"BIN\0") + bin_data)

    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_name(destination.name + ".tmp")
    temporary.write_bytes(glb)
    temporary.replace(destination)
    elapsed = time.perf_counter() - start
    return {
        "source_bytes": source.stat().st_size,
        "output_bytes": destination.stat().st_size,
        "vertices": total_vertices,
        "triangles": total_indices // 3,
        "tiles": tile_count,
        "nonempty_tiles": nonempty_tiles,
        "seconds": elapsed,
    }


def main():
    default_output = Path(__file__).resolve().parent / "meshes"
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("inputs", nargs="+", type=Path, help="one or more .vfmesh files")
    parser.add_argument("--out-dir", type=Path, default=default_output,
                        help=f"output folder (default: {default_output})")
    args = parser.parse_args()

    for source in args.inputs:
        destination = args.out_dir / (source.stem + ".glb")
        result = convert(source, destination)
        print(f"{source} -> {destination}: {result['source_bytes']} -> "
              f"{result['output_bytes']} bytes, {result['vertices']} vertices, "
              f"{result['triangles']} triangles, {result['tiles']} tiles, "
              f"{result['seconds'] * 1000:.3f} ms")


if __name__ == "__main__":
    main()