#!/usr/bin/env python3
"""Reads and writes level WADs of "Ein Königreich für ein Lama" (Argonaut).

Layout: u32 total length (excluding this field), then chunks:
  4-byte tag (stored reversed, e.g. 'OFNI' = INFO), u32 length, data
up to the chunk 'END '.

  wadtool.py list  level.wad
  wadtool.py name  level.wad new.wad "New level name"
  wadtool.py extract level.wad CHUNK file.bin
  wadtool.py replace level.wad new.wad CHUNK file.bin
"""
import struct
import sys


def read(path):
    d = open(path, 'rb').read()
    total = struct.unpack_from('<I', d, 0)[0]
    if total != len(d) - 4:
        raise ValueError(f'{path}: length field {total} does not match {len(d) - 4}')
    chunks, o = [], 4
    while o + 8 <= len(d):
        tag = d[o:o + 4][::-1].decode('latin1')
        ln = struct.unpack_from('<I', d, o + 4)[0]
        chunks.append([tag, d[o + 8:o + 8 + ln]])
        o += 8 + ln
        if tag == 'END ':
            break
    return chunks


def write(path, chunks):
    body = b''.join(t.encode('latin1')[::-1] + struct.pack('<I', len(c)) + c for t, c in chunks)
    open(path, 'wb').write(struct.pack('<I', len(body)) + body)


def get(chunks, tag):
    for c in chunks:
        if c[0] == tag:
            return c
    raise KeyError(tag)


if __name__ == '__main__':
    cmd, src = sys.argv[1], sys.argv[2]
    ch = read(src)
    if cmd == 'list':
        for t, c in ch:
            extra = ''
            if t == 'NAME':
                extra = repr(c[4:].split(b'\0')[0].decode('latin1'))
            print(f'{t!r:8} {len(c):10d}  {extra}')
    elif cmd == 'name':
        c = get(ch, 'NAME')
        c[1] = c[1][:4] + sys.argv[4].encode('latin1') + b'\0'
        write(sys.argv[3], ch)
    elif cmd == 'extract':
        open(sys.argv[4], 'wb').write(get(ch, sys.argv[3].ljust(4))[1])
    elif cmd == 'replace':
        get(ch, sys.argv[4].ljust(4))[1] = open(sys.argv[5], 'rb').read()
        write(sys.argv[3], ch)
    else:
        sys.exit(__doc__)