"""Reads and writes gzipped NBT files (level.dat) keeping every tag's type, so a file can be changed and saved back.

A tag is (type, value): compounds are dicts of name -> tag, lists are (element type, [values]). Used by
tools/make_machine_world.py to make a world creative with cheats on.
"""
import gzip
import struct

END, BYTE, SHORT, INT, LONG, FLOAT, DOUBLE, BYTE_ARRAY, STRING, LIST, COMPOUND, INT_ARRAY, LONG_ARRAY = range(13)
_SCALARS = {BYTE: ">b", SHORT: ">h", INT: ">i", LONG: ">q", FLOAT: ">f", DOUBLE: ">d"}


class _Reader:
    def __init__(self, data):
        self.data, self.pos = data, 0

    def take(self, n):
        chunk = self.data[self.pos:self.pos + n]
        self.pos += n
        return chunk

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

    def string(self):
        return self.take(self.unpack(">H")).decode("utf-8")

    def payload(self, kind):
        if kind in _SCALARS:
            return self.unpack(_SCALARS[kind])
        if kind == BYTE_ARRAY:
            return list(self.take(self.unpack(">i")))
        if kind == STRING:
            return self.string()
        if kind == LIST:
            element, length = self.unpack(">b"), self.unpack(">i")
            return element, [self.payload(element) for _ in range(length)]
        if kind == COMPOUND:
            result = {}
            while (child := self.unpack(">b")) != END:
                name = self.string()
                result[name] = (child, self.payload(child))
            return result
        if kind == INT_ARRAY:
            return [self.unpack(">i") for _ in range(self.unpack(">i"))]
        if kind == LONG_ARRAY:
            return [self.unpack(">q") for _ in range(self.unpack(">i"))]
        raise ValueError(f"unknown NBT tag type {kind}")


def _string(text):
    raw = text.encode("utf-8")
    return struct.pack(">H", len(raw)) + raw


def _payload(kind, value):
    if kind in _SCALARS:
        return struct.pack(_SCALARS[kind], value)
    if kind == BYTE_ARRAY:
        return struct.pack(">i", len(value)) + bytes(v & 0xFF for v in value)
    if kind == STRING:
        return _string(value)
    if kind == LIST:
        element, items = value
        return struct.pack(">bi", element, len(items)) + b"".join(_payload(element, v) for v in items)
    if kind == COMPOUND:
        return b"".join(struct.pack(">b", k) + _string(name) + _payload(k, v) for name, (k, v) in value.items()) + b"\x00"
    if kind == INT_ARRAY:
        return struct.pack(">i", len(value)) + b"".join(struct.pack(">i", v) for v in value)
    if kind == LONG_ARRAY:
        return struct.pack(">i", len(value)) + b"".join(struct.pack(">q", v) for v in value)
    raise ValueError(f"unknown NBT tag type {kind}")


def load(path):
    """The root compound of a gzipped NBT file, as (name, compound dict)."""
    reader = _Reader(gzip.decompress(path.read_bytes()))
    if reader.unpack(">b") != COMPOUND:
        raise ValueError(f"{path} doesn't start with a compound")
    name = reader.string()
    return name, reader.payload(COMPOUND)


def save(path, name, root):
    path.write_bytes(gzip.compress(struct.pack(">b", COMPOUND) + _string(name) + _payload(COMPOUND, root)))