#!/usr/bin/env python3
"""Summarize VAND navigation-graph blobs from streamingsector JSON files."""

from __future__ import annotations

import argparse
import base64
import collections
import dataclasses
import json
import math
import struct
from pathlib import Path
from typing import Any, Iterator


HEADER_SIZE = 0x64
POINT_SIZE = 0x14
COORD_SIZE = 0x0C


@dataclasses.dataclass
class VandBlob:
    path: Path
    ordinal: int
    size: int
    version: int
    tile_x: int
    tile_y: int
    point_count: int
    coord_count: int
    aux_count: int
    point_count_2: int
    aux_count_2: int
    point_count_3: int
    class_counts: collections.Counter[int]
    raw_class_counts: collections.Counter[int]
    mask_counts: collections.Counter[int]
    bounds: tuple[float, float, float, float, float, float] | None


def walk_json(value: Any) -> Iterator[str]:
    if isinstance(value, dict):
        blob = value.get("Bytes")
        if isinstance(blob, str):
            yield blob
        for child in value.values():
            yield from walk_json(child)
    elif isinstance(value, list):
        for child in value:
            yield from walk_json(child)


def finite_bounds(coords: list[tuple[float, float, float]]) -> tuple[float, float, float, float, float, float] | None:
    finite = [coord for coord in coords if all(math.isfinite(part) for part in coord)]
    if not finite:
        return None

    xs = [coord[0] for coord in finite]
    ys = [coord[1] for coord in finite]
    zs = [coord[2] for coord in finite]
    return min(xs), min(ys), min(zs), max(xs), max(ys), max(zs)


def parse_vand(path: Path, ordinal: int, data: bytes) -> VandBlob | None:
    if len(data) < HEADER_SIZE or not data.startswith(b"VAND"):
        return None

    version = struct.unpack_from("<I", data, 0x04)[0]
    tile_x = struct.unpack_from("<i", data, 0x08)[0]
    tile_y = struct.unpack_from("<i", data, 0x0C)[0]
    point_count = struct.unpack_from("<I", data, 0x18)[0]
    coord_count = struct.unpack_from("<I", data, 0x1C)[0]
    aux_count = struct.unpack_from("<I", data, 0x20)[0]
    point_count_2 = struct.unpack_from("<I", data, 0x24)[0]
    aux_count_2 = struct.unpack_from("<I", data, 0x28)[0]
    point_count_3 = struct.unpack_from("<I", data, 0x2C)[0]

    coord_offset = HEADER_SIZE
    point_offset = coord_offset + coord_count * COORD_SIZE
    point_end = point_offset + point_count * POINT_SIZE
    if point_end > len(data):
        return VandBlob(
            path=path,
            ordinal=ordinal,
            size=len(data),
            version=version,
            tile_x=tile_x,
            tile_y=tile_y,
            point_count=point_count,
            coord_count=coord_count,
            aux_count=aux_count,
            point_count_2=point_count_2,
            aux_count_2=aux_count_2,
            point_count_3=point_count_3,
            class_counts=collections.Counter({-1: point_count}),
            raw_class_counts=collections.Counter(),
            mask_counts=collections.Counter(),
            bounds=None,
        )

    coords: list[tuple[float, float, float]] = []
    for index in range(coord_count):
        coords.append(struct.unpack_from("<fff", data, coord_offset + index * COORD_SIZE))

    class_counts: collections.Counter[int] = collections.Counter()
    raw_class_counts: collections.Counter[int] = collections.Counter()
    mask_counts: collections.Counter[int] = collections.Counter()
    for index in range(point_count):
        offset = point_offset + index * POINT_SIZE
        mask = struct.unpack_from("<H", data, offset + 0x10)[0]
        raw_class = data[offset + 0x13]
        class_counts[raw_class & 0x3F] += 1
        raw_class_counts[raw_class] += 1
        mask_counts[mask] += 1

    return VandBlob(
        path=path,
        ordinal=ordinal,
        size=len(data),
        version=version,
        tile_x=tile_x,
        tile_y=tile_y,
        point_count=point_count,
        coord_count=coord_count,
        aux_count=aux_count,
        point_count_2=point_count_2,
        aux_count_2=aux_count_2,
        point_count_3=point_count_3,
        class_counts=class_counts,
        raw_class_counts=raw_class_counts,
        mask_counts=mask_counts,
        bounds=finite_bounds(coords),
    )


def iter_vand_blobs(path: Path) -> Iterator[VandBlob]:
    document = json.loads(path.read_text(encoding="utf-8"))
    ordinal = 0
    for encoded in walk_json(document):
        try:
            data = base64.b64decode(encoded)
        except ValueError:
            continue
        blob = parse_vand(path, ordinal, data)
        if blob:
            yield blob
            ordinal += 1


def fmt_counter(counter: collections.Counter[int], limit: int) -> str:
    return " ".join(f"{key}:{value}" for key, value in counter.most_common(limit))


def fmt_bounds(bounds: tuple[float, float, float, float, float, float] | None) -> str:
    if not bounds:
        return "<none>"
    return (
        f"({bounds[0]:.1f},{bounds[1]:.1f},{bounds[2]:.1f}).."
        f"({bounds[3]:.1f},{bounds[4]:.1f},{bounds[5]:.1f})"
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("json_files", nargs="+", type=Path)
    parser.add_argument("--detail", type=int, default=12)
    parser.add_argument("--counter-limit", type=int, default=12)
    args = parser.parse_args()

    blobs: list[VandBlob] = []
    for path in args.json_files:
        blobs.extend(iter_vand_blobs(path))

    total_points = sum(blob.point_count for blob in blobs)
    total_coords = sum(blob.coord_count for blob in blobs)
    class_counts: collections.Counter[int] = collections.Counter()
    raw_counts: collections.Counter[int] = collections.Counter()
    mask_counts: collections.Counter[int] = collections.Counter()
    versions: collections.Counter[int] = collections.Counter()
    files: collections.Counter[str] = collections.Counter()
    for blob in blobs:
        class_counts.update(blob.class_counts)
        raw_counts.update(blob.raw_class_counts)
        mask_counts.update(blob.mask_counts)
        versions[blob.version] += 1
        files[str(blob.path)] += 1

    print(f"files={len(files)} blobs={len(blobs)} points={total_points} coords={total_coords}")
    print(f"versions={dict(sorted(versions.items()))}")
    print(f"classes={fmt_counter(class_counts, args.counter_limit)}")
    print(f"raw13={fmt_counter(raw_counts, args.counter_limit)}")
    print(f"masks={fmt_counter(mask_counts, args.counter_limit)}")
    print("files:")
    for path, count in files.most_common():
        print(f"  {path}: {count}")

    if args.detail:
        print("largest blobs:")
        for blob in sorted(blobs, key=lambda item: item.point_count, reverse=True)[: args.detail]:
            print(
                f"  {blob.path.name}#{blob.ordinal} size={blob.size} tile=({blob.tile_x},{blob.tile_y}) "
                f"points={blob.point_count} coords={blob.coord_count} aux={blob.aux_count} "
                f"classes={fmt_counter(blob.class_counts, args.counter_limit)} "
                f"raw13={fmt_counter(blob.raw_class_counts, args.counter_limit)} "
                f"bounds={fmt_bounds(blob.bounds)}"
            )

    return 0


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