#!/usr/bin/env python3
"""Validate and summarize the exact game471 stage-note wire format."""

from __future__ import annotations

import argparse
import collections
import pathlib
import struct


NOTE = struct.Struct(">IBB9hB3fBB4f3IfI5fI")
TYPE_NAMES = (
    "NONE", "NORMAL", "FLICK", "HOLD", "SCRATCH", "BEAT", "MERRY GO ROUND", "HIDDEN",
    "HIDDEN2", "CRITICAL", "SLIDE HOLD", "SLIDE COUNTER", "TURN", "SPIN", "FINISH", "DUAL HOLD",
)


def u32be(blob: bytes, off: int) -> int:
    return struct.unpack_from(">I", blob, off)[0]


def decode(path: pathlib.Path):
    blob = path.read_bytes()
    if len(blob) < 20:
        raise ValueError("short header")
    notes_off = u32be(blob, 12)
    camera_off = u32be(blob, 16)
    if not (0 <= notes_off < camera_off <= len(blob)):
        raise ValueError("invalid note/camera offsets")

    off = notes_off
    name_count = u32be(blob, off)
    off += 4
    names = []
    for _ in range(name_count):
        if off >= camera_off:
            raise ValueError("truncated note name")
        size = blob[off]
        off += 1
        if off + size > camera_off:
            raise ValueError("truncated note name bytes")
        names.append(blob[off : off + size].rstrip(b"\0").decode("utf-8", "replace"))
        off += size

    if off + 4 > camera_off:
        raise ValueError("missing note count")
    count = u32be(blob, off)
    off += 4
    expected = off + count * NOTE.size
    if expected != camera_off:
        raise ValueError(f"size mismatch: count={count}, expected=0x{expected:x}, camera=0x{camera_off:x}")
    return names, [NOTE.unpack_from(blob, off + i * NOTE.size) for i in range(count)]


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=pathlib.Path, help="stage .dat or directory")
    parser.add_argument("--samples", action="store_true", help="show one decoded record per raw type")
    args = parser.parse_args()

    if args.path.is_dir():
        paths = [
            p for p in sorted(args.path.glob("*.dat"))
            if not p.name.endswith(("_ext.dat", "_clip.dat"))
        ]
    else:
        paths = [args.path]

    types = collections.Counter()
    name_counts = collections.Counter()
    override_count = 0
    overrides = collections.Counter()
    note_count = 0
    samples = {}
    failures = []
    for path in paths:
        try:
            names, notes = decode(path)
        except (OSError, ValueError, struct.error) as exc:
            failures.append((path, str(exc)))
            continue
        name_counts[len(names)] += 1
        note_count += len(notes)
        for note in notes:
            raw_type = note[1]
            types[raw_type] += 1
            override_count += note[2] != 0
            if note[2] != 0:
                overrides[(raw_type, note[2])] += 1
            samples.setdefault(raw_type, (path, note))

    print(f"record_size={NOTE.size}")
    print(f"files={len(paths)} valid={len(paths) - len(failures)} invalid={len(failures)} notes={note_count}")
    print("name_counts=" + " ".join(f"{key}:{value}" for key, value in sorted(name_counts.items())))
    print("types=" + " ".join(
        f"0x{key:02x}/{TYPE_NAMES[key] if key < len(TYPE_NAMES) else 'UNKNOWN'}:{value}"
        for key, value in sorted(types.items())
    ))
    print(f"type_override_nonzero={override_count}")
    print("overrides=" + " ".join(
        f"0x{raw:02x}/0x{override:02x}:{count}"
        for (raw, override), count in sorted(overrides.items())
    ))

    if args.samples:
        for raw_type, (path, note) in sorted(samples.items()):
            print(f"sample 0x{raw_type:02x} {path.name}: {note}")
    for path, error in failures[:10]:
        print(f"invalid {path.name}: {error}")
    return 1 if failures else 0


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