#!/usr/bin/env python3
"""Why a `.slb` bank's XMA data starts at 1392 / 1468 / 1600 / 1728.

Static, disc-only. No emulator, no decoder, no run-generated input: everything
here is read straight out of `dat/sound.pak` + `dat/sound.p00..p04`.

The answer (see docs/re/structures/slb-data-offset.md): the XMA packet grid is
2048-byte aligned *inside each individual `.pNN` segment file*, but the `.pak`
TOC addresses entries in the flat CONCATENATION of those files, at offsets that
are themselves multiples of 2048.  The segment files are not multiples of 2048
long, so the grid phase seen inside an entry is

    X = (cumulative start of the .pNN segment holding the wave) mod 2048

and the four disc-wide values are just the running sums of the segment sizes.

Subcommands
    phases            the segment table and where 1392/1468/1600/1728 come from
    verify            predicted vs measured X over every bank on the disc
    bank <name>       full structural dump of one bank entry
    chain <name> [n]  the tiling arithmetic across n consecutive entries

Usage
    SYLPHEED_DISC=/work/sylph_extract python3 slb_segment_phase.py verify
"""

import bisect
import os
import re
import struct
import sys
import zlib

BLOCK = 2048  # XMA1 packet size and the archive's data-offset alignment unit
WAVE_HDR = 4096  # RIFF + 32-byte `fmt ` + `Dmmy` pad, always exactly two blocks

# --- IPFB name hash (crate `sylpheed_formats::hash`, from `sub_82455C78`) -----
MODULUS = 0x00FF_F9D7
RECIP = 0x8003_1493


def _rotl32(x, n):
    x &= 0xFFFFFFFF
    return ((x << n) | (x >> (32 - n))) & 0xFFFFFFFF


def name_hash(name):
    a = b = 0
    for byte in name.lower().encode("latin-1"):
        c = (byte - 256) & 0xFFFFFFFF if byte >= 0x80 else byte
        a = ((_rotl32(a, 8) & 0xFFFFFF00) + c) & 0xFFFFFFFF
        b = (b + c) & 0xFFFFFFFF
        q = _rotl32(((a * RECIP) >> 32) & 0xFFFFFFFF, 9) & 0x1FF
        a = (a - q * MODULUS) & 0xFFFFFFFF
    return ((b & 0xFF) << 24) | (a & 0xFFFFFF)


class Pak:
    """Minimal IPFB reader: TOC plus raw access to the concatenated segments."""

    def __init__(self, path):
        data = open(path, "rb").read()
        assert data[:4] == b"IPFB", data[:4]
        count, self.block, self.flags = struct.unpack_from(">III", data, 4)
        self.toc = [struct.unpack_from(">III", data, 0x10 + 12 * i) for i in range(count)]
        base, pos, self.segs = path[:-4], 0, []
        i = 0
        while os.path.exists("%s.p%02d" % (base, i)):
            seg = "%s.p%02d" % (base, i)
            n = os.path.getsize(seg)
            self.segs.append((pos, pos + n, seg))
            pos, i = pos + n, i + 1
        self.seg_starts = [s for s, _, _ in self.segs]

    def raw(self, off, size):
        """Bytes [off, off+size) of the flat concatenation of the .pNN files."""
        out = b""
        for s, e, path in self.segs:
            if off < e and off + size > s:
                lo, hi = max(off, s), min(off + size, e)
                with open(path, "rb") as f:
                    f.seek(lo - s)
                    out += f.read(hi - lo)
        return out

    def read(self, idx):
        _, off, csz = self.toc[idx]
        raw = self.raw(off, csz)
        return zlib.decompress(raw[10:]) if raw[:2] == b"Z1" else raw

    def find(self, name):
        h = name_hash(name)
        for i, (hh, _, _) in enumerate(self.toc):
            if hh == h:
                return i
        return None

    def phase_at(self, concat_off):
        """The 2048-grid phase in force at a flat-concatenation offset."""
        k = bisect.bisect_right(self.seg_starts, concat_off) - 1
        return self.seg_starts[k] % BLOCK


def entry_names(disc):
    """Every `.slb` path named in any `<lang>\\sounds.tbl` inside tables.pak."""
    tbl = Pak(os.path.join(disc, "dat", "tables.pak"))
    names = set()
    for lang in ("eng", "jpn", "deu", "fra", "esp", "ita"):
        i = tbl.find(lang + "\\sounds.tbl")
        if i is None:
            continue
        blob = tbl.read(i)
        for run in re.finditer(rb"[\x20-\x7e]{6,}", blob):
            for m in re.finditer(r"[A-Za-z0-9_\\.]+\.slb", run.group().decode("latin-1")):
                names.add(m.group())
    return names


# --- structure walkers --------------------------------------------------------
def bank_headers(data, phase):
    """Bank headers inside an entry. They sit on the grid, so only `phase` is scanned."""
    out = []
    for off in range(phase, max(0, len(data) - 56), BLOCK):
        if data[off + 0x18 : off + 0x1C] != b"\x00\x00\x08\x00":
            continue
        if data[off : off + 4] != data[off + 0x20 : off + 0x24]:
            continue
        f = struct.unpack(">14I", data[off : off + 56])
        out.append(
            dict(at=off, id=f[0], data_size=f[7], hdr_blocks=f[9],
                 bits=f[10] >> 16, channels=f[10] & 0xFFFF)
        )
    return out


def waves(data):
    out = []
    for m in re.finditer(b"RIFF", data):
        r = m.start()
        if data[r + 8 : r + 12] != b"WAVE":
            continue
        dp = data.find(b"data", r)
        if dp < 0:
            continue
        out.append(dict(at=r, data_at=dp + 8,
                        data_size=struct.unpack_from("<I", data, dp + 4)[0]))
    return out


def seeks(data):
    """`seek` (XMA dpds) chunks: one u32 LE per emitted packet, so the entry count
    is the wave's packet count and the chunk sits immediately after its data."""
    out = []
    for m in re.finditer(b"seek", data):
        s = m.start()
        if s + 16 > len(data):
            continue
        size, streams, n = struct.unpack_from("<III", data, s + 4)
        if streams == 1 and size == 8 + 4 * n and n:
            out.append(dict(at=s, packets=n, data_start=s - n * BLOCK))
    return out


# --- subcommands --------------------------------------------------------------
def cmd_phases(pak):
    print("segment            size      size%2048   cum_start      PHASE (= cum_start%2048)")
    running = 0
    for s, e, path in pak.segs:
        n = e - s
        print("  %-10s %11d  %8d   %12d   %d" % (os.path.basename(path), n, n % BLOCK, s, s % BLOCK))
        running = s % BLOCK
    print()
    print("The four disc-wide values are the running sums of `size % 2048`:")
    acc = 0
    for s, e, path in pak.segs:
        print("    %-10s phase %4d" % (os.path.basename(path), s % BLOCK))
        acc = (acc + (e - s)) % BLOCK
    _ = running


def cmd_verify(pak, disc):
    names = entry_names(disc)
    ok = bad = 0
    seek_ok = seek_bad = 0
    misses = []
    for name in sorted(names):
        i = pak.find(name)
        if i is None:
            continue
        _, off, _ = pak.toc[i]
        data = pak.read(i)
        w = waves(data)
        if w:
            pred = pak.phase_at(off + w[0]["at"])
            meas = w[0]["at"] % BLOCK
            if pred == meas:
                ok += 1
            else:
                bad += 1
                if len(misses) < 10:
                    misses.append((name, pred, meas))
        else:
            # RIFF-less entry: the prediction is still exact, cross-check on `seek`.
            sk = seeks(data)
            if sk:
                pred = pak.phase_at(off + sk[0]["at"])
                if sk[0]["at"] % BLOCK == pred:
                    seek_ok += 1
                else:
                    seek_bad += 1
    print("banks with a RIFF   : predicted == measured %d, mismatches %d" % (ok, bad))
    print("RIFF-less, via seek : agree %d, disagree %d" % (seek_ok, seek_bad))
    for m in misses:
        print("   MISMATCH", m)
    return 1 if (bad or seek_bad) else 0


def cmd_bank(pak, name):
    i = pak.find(name)
    if i is None:
        sys.exit("no such entry: " + name)
    _, off, csz = pak.toc[i]
    data = pak.read(i)
    phase = pak.phase_at(off)
    print("%s\n  pak offset %d (%%2048=%d)  stored %d bytes  phase %d"
          % (name, off, off % BLOCK, csz, phase))
    events = []
    for b in bank_headers(data, phase):
        events.append((b["at"], "BANK  id=%d data_size=%d hdr=%d blocks (%d B) %dbit/%dch -> wave at %d"
                       % (b["id"], b["data_size"], b["hdr_blocks"], b["hdr_blocks"] * BLOCK,
                          b["bits"], b["channels"], b["at"] + b["hdr_blocks"] * BLOCK)))
    for w in waves(data):
        over = w["data_at"] + w["data_size"] - len(data)
        events.append((w["at"], "RIFF  data@%d size=%d (%d packets)%s"
                       % (w["data_at"], w["data_size"], w["data_size"] // BLOCK,
                          "  OVERRUNS window by %d" % over if over > 0 else "")))
    for s in seeks(data):
        events.append((s["at"], "seek  %d packets -> its data began at %d%s"
                       % (s["packets"], s["data_start"],
                          "  (BEFORE this window)" if s["data_start"] < 0 else "")))
    for at, text in sorted(events):
        print("   %8d  %%2048=%-5d %s" % (at, at % BLOCK, text))


def cmd_chain(pak, name, count):
    """Show that entry K's overrunning wave lands exactly on entry K+1's leading `seek`."""
    i = pak.find(name)
    if i is None:
        sys.exit("no such entry: " + name)
    order = sorted(range(len(pak.toc)), key=lambda k: pak.toc[k][1])
    start = order.index(i)
    prev = None
    for k in order[start : start + count]:
        _, off, csz = pak.toc[k]
        data = pak.read(k)
        padded = (len(data) + BLOCK - 1) // BLOCK * BLOCK
        sk = seeks(data)
        lead = sk[0] if sk and sk[0]["data_start"] < 0 else None
        w = waves(data)
        tail = w[-1] if w and w[-1]["data_at"] + w[-1]["data_size"] > len(data) else None
        line = "off=%-11d len=%-7d padded=%-7d" % (off, len(data), padded)
        if lead:
            line += "  leading seek@%-6d (%d packets)" % (lead["at"], lead["packets"])
        if tail:
            line += "  tail wave ends at %d" % (tail["data_at"] + tail["data_size"])
        print(line)
        if prev and lead:
            predicted = prev[0] - prev[1]  # tail end minus previous padded length
            mark = "OK " if predicted == lead["at"] else "!! "
            print("      %s predecessor's wave ends %d past its padded window; leading seek is at %d"
                  % (mark, predicted, lead["at"]))
        prev = ((tail["data_at"] + tail["data_size"]) if tail else None, padded)
        if prev[0] is None:
            prev = None


def main():
    disc = os.environ.get("SYLPHEED_DISC")
    if not disc:
        sys.exit("set SYLPHEED_DISC to the extracted disc root")
    pak = Pak(os.path.join(disc, "dat", "sound.pak"))
    argv = sys.argv[1:] or ["phases"]
    cmd = argv[0]
    if cmd == "phases":
        cmd_phases(pak)
    elif cmd == "verify":
        sys.exit(cmd_verify(pak, disc))
    elif cmd == "bank":
        cmd_bank(pak, argv[1])
    elif cmd == "chain":
        cmd_chain(pak, argv[1], int(argv[2]) if len(argv) > 2 else 4)
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main()