#!/usr/bin/env python3
"""Bounded, independently authored MB3 DAT research inspector; not a production codec."""
import argparse
import json
from pathlib import Path
import struct

VERSIONS = {201501280, 201512200}
MAX_ITEMS = 100_000
MAX_BYTES = 64 * 1024 * 1024


class InvalidDAT(ValueError):
    pass


class Reader:
    def __init__(self, data):
        if len(data) > MAX_BYTES:
            raise InvalidDAT("research inspector size limit exceeded")
        self.data, self.pos, self.old = data, 0, False

    def take(self, n):
        if n < 0 or n > len(self.data) - self.pos:
            raise InvalidDAT(f"truncated field at offset {self.pos}, needs {n} bytes")
        result = self.data[self.pos:self.pos + n]
        self.pos += n
        return result

    def scalar(self, fmt):
        return struct.unpack('<' + fmt, self.take(struct.calcsize('<' + fmt)))[0]

    def integer(self, selector):
        return self.scalar('B' if selector == 0 else ('h' if self.old else 'H') if selector == 1 else 'i')

    def count(self, selector=2, minimum_bytes=0):
        n = self.integer(selector)
        if n < 0 or n > MAX_ITEMS:
            raise InvalidDAT(f"invalid/excessive count {n} at offset {self.pos}")
        if minimum_bytes and n > (len(self.data) - self.pos) // minimum_bytes:
            raise InvalidDAT(f"count {n} exceeds remaining payload at offset {self.pos}")
        return n

    def string(self):
        end = self.data.find(b'\0', self.pos)
        if end < 0:
            raise InvalidDAT(f"unterminated ASCII string at offset {self.pos}")
        return self.take(end - self.pos + 1)[:-1].decode('ascii', errors='replace').replace('\ufffd', '?')

    def vectors(self, n):
        if n < 0 or n > MAX_ITEMS or n > (len(self.data) - self.pos) // 12:
            raise InvalidDAT(f"invalid/truncated vector array {n} at offset {self.pos}")
        return [list(v) for v in struct.iter_unpack('<fff', self.take(12 * n))]


def inspect(data):
    r = Reader(data)
    if r.take(8) != b'Mb3World':
        raise InvalidDAT('wrong magic')
    version = r.scalar('i')
    if version not in VERSIONS:
        raise InvalidDAT(f'unsupported version {version}')
    r.old = version == 201501280
    flags = r.scalar('B')
    computed = bool(flags) if r.old else bool(flags & 1)
    sp_width, uv_width = (flags >> 1) & 3, (flags >> 3) & 3
    names = [] if r.old else [r.string() for _ in range(r.count(sp_width, 1))]
    uvnames = [] if r.old else [r.string() for _ in range(r.count(uv_width, 1))]
    brushes = []
    brush_count = r.count(minimum_bytes=2)
    patch_count = r.count() if r.old else None
    for _ in range(brush_count):
        start = r.pos
        b = r.scalar('B')
        planes = []
        for _ in range(r.count(b & 3, 36)):
            plane = {'offset': r.pos, 'strings': {}, 'uv': {}}
            for key, table, width in [('strings', names, sp_width), ('uv', uvnames, uv_width)]:
                if not table:
                    continue
                while True:
                    index = r.integer(width)
                    if index == (255 if width == 0 else 65535 if width == 1 else -1):
                        break
                    if index < 0 or index >= len(table):
                        raise InvalidDAT(f'property index {index} out of range at offset {r.pos}')
                    # Legacy reader overwrites repeated binary property slots.
                    plane[key][str(index)] = r.string() if key == 'strings' else list(struct.unpack('<10f', r.take(40)))
            plane['p0'], plane['u'], plane['v'] = r.vectors(3)
            if computed:
                plane['polygon'] = r.vectors(r.count((b >> 2) & 3, 12))
            planes.append(plane)
        brushes.append({'offset': start, 'flags': b, 'planes': planes})
    if patch_count is None:
        patch_count = r.count(minimum_bytes=5)
    patches = []
    for _ in range(patch_count):
        start = r.pos
        b = r.scalar('B')
        au, av = r.count(b & 3), r.count(b & 3)
        qu, qv = r.count((b >> 2) & 3), r.count((b >> 2) & 3)
        patch = {'offset': start, 'flags': b, 'anchors_u': au, 'anchors_v': av,
                 'quads_u': qu, 'quads_v': qv, 'anchors': r.vectors(au * av)}
        if computed:
            n = (au // 2 * qu + 1) * (av // 2 * qv + 1)
            patch['points'], patch['normals'] = r.vectors(n), r.vectors(n)
        patches.append(patch)
    return {'version': version, 'flags': flags, 'computed': computed,
            'string_properties': names, 'uv_properties': uvnames,
            'brushes': brushes, 'patches': patches,
            'consumed_bytes': r.pos, 'trailing_bytes': len(data) - r.pos}


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('path', type=Path)
    args = p.parse_args()
    try:
        with args.path.open('rb') as stream:
            data = stream.read(MAX_BYTES + 1)
        print(json.dumps(inspect(data), indent=2, allow_nan=False))
    except (OSError, ValueError) as exc:
        p.exit(1, f'{exc}\n')


if __name__ == '__main__':
    main()