#!/usr/bin/env python
"""
extract-mw4.py - dump the records of a GameOS .mw4 package (Stuff::Database / "#VBD")
to human-readable text. Decompresses the per-record LZW (gos_LZDecompress) payloads.

Faithful port of gos_LZDecompress (variable-width LZW, 9->12 bit, CLEAR=256, EOF=257)
from CoreTech\\Libraries\\GameOS\\FileIO.cpp, plus the Database index format from
mw4\\Libraries\\stuff\\Database.cpp.

Usage:  python extract-mw4.py "Arctic Wolf 1.mw4" [-o outdir]
        (with no -o it prints every text record to stdout)
"""
import struct, sys, os, argparse

def lz_decompress(src):
    """Port of gos_LZDecompress: variable-width LZW, LSB-first, CLEAR=256 EOF=257."""
    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      # word: previous code
    suffix = [0] * 8192      # byte: suffix char
    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

        # --- expand the code into a stack (walk chain back to a literal) ---
        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)                            # terminal literal
        oldsuffix = walk & 0xff                        # first char of expansion
        out += bytes(reversed(stack))

        # --- add new dictionary entry {prev_code, firstchar} every code ---
        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 parse_db(path):
    data = open(path, "rb").read()
    tag, fmt, content, indexsize, nrec, nextid = struct.unpack_from("<4sIIIHH", data, 0)
    if tag != b"#VBD":
        raise SystemExit(f"{path}: not a #VBD package (got {tag!r})")
    recs, off = [], 20
    for _ in range(nrec):
        mtime, = struct.unpack_from("<q", data, off); off += 8
        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
        blob = b""
        if dlen:
            raw = data[indexsize + doff : indexsize + doff + rlen]
            blob = lz_decompress(raw) if rlen < dlen else raw
            if len(blob) != dlen:
                blob += b""  # tolerate; report below
        recs.append((rid, name, dlen, rlen, blob))
    return content, recs


def safe(name):
    return "".join(c if c.isalnum() or c in " ._-{}[]" else "_" for c in name).strip()

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("files", nargs="+")
    ap.add_argument("-o", "--outdir")
    args = ap.parse_args()
    for f in args.files:
        content, recs = parse_db(f)
        print(f"\n===== {os.path.basename(f)}  (contentVersion={content}, {len(recs)} records) =====")
        for rid, name, dlen, rlen, blob in recs:
            tag = "raw" if rlen == dlen else f"lzw {rlen}->{dlen}"
            print(f"  [{rid}] {name!r}  {dlen} bytes ({tag})" +
                  ("" if len(blob) == dlen else f"  !! decoded {len(blob)}"))
        if args.outdir:
            base = os.path.join(args.outdir, safe(os.path.splitext(os.path.basename(f))[0]))
            os.makedirs(base, exist_ok=True)
            for rid, name, dlen, rlen, blob in recs:
                fn = os.path.join(base, f"{rid:02d}_{safe(name) or 'unnamed'}.txt")
                open(fn, "wb").write(blob)
            print(f"  -> wrote {len(recs)} records to {base}")
        else:
            for rid, name, dlen, rlen, blob in recs:
                printable = sum(1 for b in blob if 9 <= b <= 13 or 32 <= b < 127)
                if blob and printable / max(1, len(blob)) > 0.85:
                    print(f"\n----- record [{rid}] {name} -----")
                    sys.stdout.write(blob.decode("latin-1"))
                    if not blob.endswith(b"\n"): print()

if __name__ == "__main__":
    main()