"""Arkham City 'WWAD' container (SFX and *_LOC_*.wad).

Layout (little endian):
  0x00  'WWAD'
  0x04  u32 version (0x01020000)
  0x08  u32 entry count N
  0x0C  u32 table offset (0x1C)
  0x10  u32 table size (N*12)
  0x14  u32 data start (header area size, 0x1000 or 0x8000; aligned)
  0x18  u32 data size  (file size - data start)
  0x1C  N x {u32 hash, u32 size, u32 offset}   sorted by hash as *signed* int32
Entries are stored in table order, each at a 0x1000-aligned offset; 'size' is the
entry payload size padded to 0x800. Padding bytes are zero; the file ends 0x1000-aligned.
See docs/formats/wad.md.
"""
import struct
from dataclasses import dataclass

MAGIC = b"WWAD"
ENTRY_ALIGN = 0x1000
SIZE_ALIGN = 0x800


def _s32(x):
    return struct.unpack("<i", struct.pack("<I", x))[0]


def align(x, a):
    return (x + a - 1) // a * a


@dataclass
class Entry:
    hash: int
    size: int
    offset: int


class Wad:
    def __init__(self, path):
        self.path = path
        with open(path, "rb") as f:
            hdr = f.read(0x1C)
            (magic, self.version, n, self.table_off, self.table_size,
             self.data_start, self.data_size) = struct.unpack("<4s6I", hdr)
            if magic != MAGIC:
                raise ValueError("not a WWAD file")
            f.seek(self.table_off)
            tbl = f.read(12 * n)
            f.seek(0)
            self.header_area = f.read(self.data_start)
        self.entries = [Entry(*struct.unpack_from("<III", tbl, 12 * i)) for i in range(n)]
        self.index = {e.hash: e for e in self.entries}

    def read(self, h) -> bytes:
        e = self.index[h]
        with open(self.path, "rb") as f:
            f.seek(e.offset)
            return f.read(e.size)

    def check(self):
        """Validate structural invariants; returns list of problems (empty = OK)."""
        probs = []
        hs = [_s32(e.hash) for e in self.entries]
        if hs != sorted(hs):
            probs.append("table not sorted by signed hash")
        if self.table_size != 12 * len(self.entries):
            probs.append("table size mismatch")
        pos = self.data_start
        for e in self.entries:
            if e.offset != align(pos, ENTRY_ALIGN):
                probs.append("entry %08X at %x, expected %x" % (e.hash, e.offset, align(pos, ENTRY_ALIGN)))
            if e.size % SIZE_ALIGN:
                probs.append("entry %08X size not 0x800-aligned" % e.hash)
            pos = e.offset + e.size
        return probs


def write_wad(src: Wad, out_path: str, replacements: dict):
    """Rebuild a WAD, replacing entry payloads (hash -> bytes). Layout rules are reproduced
    exactly so an unmodified rebuild is byte-identical to the original."""
    n = len(src.entries)
    payloads = {}
    for h, data in replacements.items():
        if h not in src.index:
            raise KeyError("hash %08X not in WAD" % h)
        payloads[h] = data + b"\0" * (align(len(data), SIZE_ALIGN) - len(data))
    new_entries = []
    pos = src.data_start
    for e in src.entries:
        size = len(payloads[e.hash]) if e.hash in payloads else e.size
        off = align(pos, ENTRY_ALIGN)
        new_entries.append(Entry(e.hash, size, off))
        pos = off + size
    end = pos
    total = align(end, ENTRY_ALIGN)
    hdr = bytearray(src.header_area)
    struct.pack_into("<4s6I", hdr, 0, MAGIC, src.version, n, src.table_off, 12 * n, src.data_start,
                     total - src.data_start)
    for i, e in enumerate(new_entries):
        struct.pack_into("<III", hdr, src.table_off + 12 * i, e.hash, e.size, e.offset)
    with open(src.path, "rb") as fi, open(out_path, "wb") as fo:
        fo.write(hdr)
        for old, new in zip(src.entries, new_entries):
            fo.write(b"\0" * (new.offset - fo.tell()))
            if old.hash in payloads:
                fo.write(payloads[old.hash])
            else:
                fi.seek(old.offset)
                remaining = old.size
                while remaining:
                    chunk = fi.read(min(remaining, 1 << 22))
                    fo.write(chunk)
                    remaining -= len(chunk)
        fo.write(b"\0" * (total - fo.tell()))
    return new_entries