import struct
from functools import partial
from dataclasses import dataclass
import sys

with open(sys.argv[1], "rb") as f:
    buf = memoryview(f.read())

class Reader:
    def __init__(self, buf):
        self.index = 0
        self.mem = memoryview(buf)

    def read(self, n, t):
        v = struct.unpack(t, self.mem[self.index:self.index+n])
        self.index += n
        return v

    def read_char4(self):
        v = self.read(4, "<cccc")
        return b"".join(v)

    def read_int(self):
        v, = self.read(4, "<i")
        return v

    def read_xyzi(self):
        v = self.read(4, "<BBBB")
        return v

    def read_memory(self, n):
        mem = self.mem[self.index:self.index+n]
        self.index += n
        return mem

@dataclass
class Chunk:
    id: str
    chunk: bytes
    children: list[Chunk]

@dataclass
class SizeChunk:
    x: int
    y: int
    z: int

@dataclass
class XYZIChunk:
    voxels: tuple[int, int, int, int]

@dataclass
class RGBAChunk:
    colors: tuple[int, int, int, int]

@dataclass
class PackChunk:
    model_count: int

def validate_chunk(chunk):
    r = Reader(chunk.chunk)
    if chunk.id == b"SIZE":
        x = r.read_int()
        y = r.read_int()
        z = r.read_int()
        assert r.index == len(chunk.chunk)
        return SizeChunk(x, y, z)
    elif chunk.id == b"XYZI":
        voxel_count = r.read_int()
        voxels = [r.read_xyzi() for _ in range(voxel_count)]
        assert r.index == len(chunk.chunk)
        return XYZIChunk(voxels)
    elif chunk.id == b"RGBA":
        colors = [r.read_xyzi() for _ in range(256)]
        assert r.index == len(chunk.chunk)
        return RGBAChunk(colors)
    elif chunk.id == b"PACK":
        model_count = r.read_int()
        assert r.index == len(chunk.chunk)
        return PackChunk(model_count)
    else:
        assert False, chunk.id

def parse_chunk(r: Reader):
    global index
    chunk_id = r.read_char4()
    chunk_size = r.read_int()
    children_size = r.read_int()
    chunk = r.read_memory(chunk_size)

    children_start = r.index
    children = []
    while (r.index - children_start) < children_size:
        child_chunk = parse_chunk(r)
        children.append(child_chunk)
    assert (r.index - children_start) == children_size

    return Chunk(chunk_id, chunk, children)

reader = Reader(buf)

magic = reader.read_char4()
assert magic == b"VOX ", magic
version = reader.read_int()
assert version == 150, version

main = parse_chunk(reader)
assert reader.index == len(buf)
chunks = []
for chunk in main.children:
    chunk = validate_chunk(chunk)
    chunks.append(chunk)

with open(sys.argv[2], 'wb') as f:
    size_chunk = None
    for chunk in chunks:
        if type(chunk) is SizeChunk:
            print('size')
            size_chunk = chunk
            f.write(struct.pack("<III", chunk.x, chunk.y, chunk.z))
            print(chunk.x, chunk.y, chunk.z)
        elif type(chunk) is XYZIChunk:
            print('xyzi')
            assert size_chunk is not None
            #print(size_chunk)
            print(len(chunk.voxels))
            f.write(struct.pack("<I", len(chunk.voxels)))
            for voxel in chunk.voxels:
                #print(voxel)
                f.write(struct.pack("<BBBB", *voxel))
            size_chunk = None
        elif type(chunk) is RGBAChunk:
            print('rgba')
            for color in chunk.colors:
                f.write(struct.pack("<ffff", *color))
        elif type(chunk) is PackChunk:
            print("model count", chunk.model_count)
            pass
        else:
            assert False, type(chunk)