#!/usr/bin/env python3
"""
mw4db.py - reader for GameOS "#VBD" resource packages (*.mw4).

Format (from mw4\\Libraries\\stuff\\Database.cpp):

    header (20 bytes, little-endian)
        char[4]  tag         "#VBD"
        u32      format
        u32      contentVersion      (63 for this codebase, VER_CONTENTVERSION)
        u32      indexSize           byte offset of the payload area
        u16      recordCount
        u16      nextId

    then `recordCount` directory records, packed, no alignment:
        s64      FILETIME    when this entry was last (re)packed
        u32      dataLength  decompressed size
        u32      recLength   stored size
        u32      dataOffset  offset from `indexSize`
        u16      recordId
        u8       nameLength
        char[]   name        (latin-1, backslash-separated, may carry
                              {qualifier} / [page] suffixes)

    payload at indexSize + dataOffset, `recLength` bytes.
    Storage rule (Database.cpp:451): recLength == dataLength  ->  raw,
    otherwise LZW-compressed with gos_LZCompress.

The decompressor is a faithful port of gos_LZDecompress from
CoreTech\\Libraries\\GameOS\\FileIO.cpp - variable-width LZW, LSB-first,
9 -> 12 bits, CLEAR = 256, EOF = 257, dictionary starts at 258.
"""
import struct, hashlib, os


def lz_decompress(src: bytes) -> bytes:
    src = bytes(src) + b"\x00\x00\x00\x00"          # u32-read padding margin
    bitpos = 0

    def getcode(width, mask):
        nonlocal bitpos
        i = bitpos >> 3
        v = src[i] | (src[i + 1] << 8) | (src[i + 2] << 16) | (src[i + 3] << 24)
        v = (v >> (bitpos & 7)) & mask
        bitpos += width
        return v

    CLEAR, EOF, FREE = 256, 257, 258
    chain = [0] * 8192
    suffix = [0] * 8192
    width, mask, maxindex, free = 9, 511, 512, FREE
    oldchain = oldsuffix = 0
    out = bytearray()

    while True:
        code = getcode(width, mask)
        if code == EOF:
            break
        if code == CLEAR:
            width, mask, maxindex, free = 9, 511, 512, FREE
            code = getcode(width, mask)              # ClearHash emits one literal
            out.append(code & 0xFF)
            oldchain, oldsuffix = code, code & 0xFF
            continue

        stack = []
        if code >= free:                              # KwKwK special case
            stack.append(oldsuffix)
            suffix[code] = oldsuffix
            chain[code] = oldchain
            walk = oldchain
        else:
            walk = code
        while walk >= 256:
            stack.append(suffix[walk])
            walk = chain[walk]
        stack.append(walk)
        oldsuffix = walk & 0xFF
        out += bytes(reversed(stack))

        suffix[free] = oldsuffix
        chain[free] = oldchain
        free += 1
        oldchain = code
        if free == maxindex and width != 12:
            width += 1
            maxindex <<= 1
            mask = (mask << 1) | 1
    return bytes(out)


def read_index(path):
    """Parse directory only. Fast - no decompression.

    Returns (contentVersion, [(recordId, name, dataLen, recLen, storedBytes), ...])
    or None when the file is not a #VBD package.
    """
    with open(path, "rb") as fh:
        data = fh.read()
    if data[:4] != b"#VBD":
        return None
    _tag, _fmt, content, indexsize, nrec, _nextid = struct.unpack_from("<4sIIIHH", data, 0)
    out, off = [], 20
    for _ in range(nrec):
        off += 8                                            # FILETIME
        dlen, rlen, doff = struct.unpack_from("<III", data, off); off += 12
        rid, nlen = struct.unpack_from("<HB", data, off); off += 3
        name = data[off:off + nlen].decode("latin-1"); off += nlen
        raw = data[indexsize + doff: indexsize + doff + rlen] if dlen else b""
        out.append((rid, name, dlen, rlen, raw))
    return content, out


def read_records(path):
    """Parse and decode every record.

    Returns (contentVersion, [(recordId, name, dataLen, recLen, decodedBytes), ...])
    """
    r = read_index(path)
    if r is None:
        return None
    content, recs = r
    out = []
    for rid, name, dlen, rlen, raw in recs:
        blob = b"" if not dlen else (raw if rlen == dlen else lz_decompress(raw))
        out.append((rid, name, dlen, rlen, blob))
    return content, out


def norm(name: str) -> str:
    """Canonical comparison key for an entry name.

    Package entry names are case-inconsistent between builds (V4H's packer
    lowercased everything) and use backslashes. Always compare through this.
    """
    return name.lower().replace("\\", "/")


def md5(b: bytes) -> str:
    return hashlib.md5(b).hexdigest()


def walk_packages(root):
    """Yield (absolutePath, relativePathWithForwardSlashes) for every *.mw4 under root."""
    for dirpath, _dirs, files in os.walk(root):
        for f in sorted(files):
            if f.lower().endswith(".mw4"):
                p = os.path.join(dirpath, f)
                yield p, os.path.relpath(p, root).replace("\\", "/")