"""Build a small synthetic Arkham City folder (WADs, a stream pack) for tests that must not depend on
real game data. Layout rules follow docs/formats/wad.md; WEMs come from ackit's own encoder."""
import os
import struct

import numpy as np

from ackit import encoder

SPEECH_ENTRY = 0xB49CDF48   # CHR_RadioScanner (3 variations here)
SINGLE_ENTRY = 0x3F21770E   # an entry with one WEM
PACK_ID = 0xEDF036D6        # 'Music' stream pack
MEDIA = (0x003CB311, 0x393BC53A, 0x40000001)
AMB_ENTRY = 0x98A82EFF      # AMB_Steel_B3


def tone(seconds, rate=24000, freq=440.0, ch=1):
    t = np.arange(int(seconds * rate)) / rate
    x = (0.3 * np.sin(2 * np.pi * freq * t) * (0.6 + 0.4 * np.sin(2 * np.pi * 2 * t))).astype(np.float32)
    return np.repeat(x[:, None], ch, axis=1), rate


def wem(seconds, freq, rate=24000, ch=1):
    x, sr = tone(seconds, rate, freq, ch)
    return encoder.encode(x, sr, rate, ch)[0]


def _a(x, n):
    return (x + n - 1) // n * n


def write_wad(path, payloads):
    """payloads: {hash: bytes}"""
    s32 = lambda h: struct.unpack("<i", struct.pack("<I", h))[0]
    hashes = sorted(payloads, key=s32)
    n = len(hashes)
    data_start = _a(0x1C + 12 * n, 0x1000)
    pos, table, body = data_start, [], []
    for h in hashes:
        off = _a(pos, 0x1000)
        p = payloads[h] + bytes(_a(len(payloads[h]), 0x800) - len(payloads[h]))
        table.append((h, len(p), off))
        body.append((off, p))
        pos = off + len(p)
    total = _a(pos, 0x1000)
    out = bytearray(total)
    struct.pack_into("<4s6I", out, 0, b"WWAD", 0x01020000, n, 0x1C, 12 * n, data_start, total - data_start)
    for i, t in enumerate(table):
        struct.pack_into("<III", out, 0x1C + 12 * i, *t)
    for off, p in body:
        out[off:off + len(p)] = p
    os.makedirs(os.path.dirname(path), exist_ok=True)
    open(path, "wb").write(bytes(out))


def grouped(wems, junk_gap=False):
    """Several WEMs in one entry, each at the next 0x800 boundary (zero padding)."""
    out = bytearray()
    for w in wems:
        out += bytes(_a(len(out), 0x800) - len(out))
        out += w
    return bytes(out)


def pack(media):
    """-> (bnk bytes, wad payload) for {media id: wem}. One gap gets non-zero bytes, like the real
    'Music' pack, to test byte-exact reverts."""
    ids = sorted(media)
    payload = bytearray()
    table = []
    for i, m in enumerate(ids):
        start = _a(len(payload), 0x800)
        gap = bytearray(start - len(payload))
        if i == 1 and len(gap) > 4:
            gap[1:4] = b"\x11\x22\x33"
        payload += gap
        table.append((m, len(media[m]), len(payload)))
        payload += media[m]
    payload += bytes(_a(len(payload), 0x800) - len(payload))
    bank = b"BKHD" + struct.pack("<I", 8) + struct.pack("<II", 56, PACK_ID)
    n = len(ids)
    hdr = bytearray(0x800)
    struct.pack_into("<4sII", hdr, 0, b"BKPK", 0x01020000, PACK_ID)
    struct.pack_into("<IIIIII", hdr, 0x10, n, 0, 12 * n, 0x48, len(payload), 0)
    struct.pack_into("<I", hdr, 0x2C, 0x48 + 12 * n)
    struct.pack_into("<II", hdr, 0x38, len(bank), 0x800)
    for i, t in enumerate(table):
        struct.pack_into("<III", hdr, 0x48 + 12 * i, *t)
    return bytes(hdr) + bank, bytes(payload)


def bank(bank_id, media, name=None, extra_sources=(), prefetch=(), events=()):
    """A v56 soundbank holding `media` {id: wem} in memory (DIDX/DATA, 16-byte aligned) with one HIRC
    Sound object per media; extra_sources: [(media id, bank id, off, size)] for sources of other banks;
    media ids in `prefetch` get a prefetch source (their DIDX data is the head of a streamed WEM);
    events: [(event name, streamed media id)] -> event -> action -> streamed Sound object."""
    data = bytearray()
    didx = b""
    for m in sorted(media):
        data += bytes(_a(len(data), 16) - len(data))
        didx += struct.pack("<III", m, len(data), len(media[m]))
        data += media[m]
    bkhd = struct.pack("<II", 56, bank_id) + bytes(8)
    head = b"BKHD" + struct.pack("<I", len(bkhd)) + bkhd + b"DIDX" + struct.pack("<I", len(didx)) + didx
    data_pos = len(head) + 8                       # DATA payload, from the start of the bank
    objs = []
    for i, m in enumerate(sorted(media)):
        off = struct.unpack_from("<I", didx, 12 * i + 4)[0]
        objs.append((m, bank_id if m not in prefetch else 0x1234, data_pos + off, len(media[m]),
                     2 if m in prefetch else 0))
    objs += [x + (0,) for x in extra_sources]
    from ackit.bnk import fnv1
    hirc_objs = []
    for i, (m, b, off, size, stream) in enumerate(objs):
        hirc_objs.append((2, struct.pack("<7I", 0x100 + i, 0x00040001, stream, m, b, off, size) + bytes(6)))
    for i, (ev, m) in enumerate(events):
        sid, aid = 0x300 + i, 0x400 + i
        hirc_objs.append((2, struct.pack("<7I", sid, 0x00040001, 1, m, 0, 0, 0) + bytes(6)))
        hirc_objs.append((3, struct.pack("<III", aid, 0x4011, sid) + bytes(8)))
        hirc_objs.append((4, struct.pack("<III", fnv1(ev), 1, aid)))
    hirc = struct.pack("<I", len(hirc_objs))
    for t, body in hirc_objs:
        hirc += bytes([t]) + struct.pack("<I", len(body)) + body
    out = head + b"DATA" + struct.pack("<I", len(data)) + bytes(data) + b"HIRC" + struct.pack("<I", len(hirc)) + hirc
    if name:
        nm = name.encode()
        stid = struct.pack("<II", 1, 1) + struct.pack("<IB", bank_id, len(nm)) + nm
        out += b"STID" + struct.pack("<I", len(stid)) + stid
    return out


def package(path, body, block_size=0x1000, chunk_blocks=3, tables=(0, 0, 0, 0, 0, 0)):
    """A fully LZO-compressed UE3 package (docs/formats/upk.md) whose uncompressed data area is `body`.
    Small blocks so that tests cross block and chunk boundaries. `tables`: (NameCount, NameOffset,
    ExportCount, ExportOffset, ImportCount, ImportOffset) as absolute uncompressed offsets."""
    from ackit import lzo
    folder = b"None\x00"
    nblocks = max(1, _a(len(body), block_size) // block_size)
    nchunks = _a(nblocks, chunk_blocks) // chunk_blocks

    def summary(chunks):
        s = struct.pack("<IHHi", 0x9E2A83C1, 805, 101, 0) + struct.pack("<i", len(folder)) + folder
        s += struct.pack("<12i", 0, *tables, 0, 0, 0, 0, 0) + b"GUID-synthetic!!"
        s += struct.pack("<i3i", 1, 0, 0, 0) + struct.pack("<3I", 7748, 0, 2)
        s += struct.pack("<i", len(chunks)) + b"".join(struct.pack("<4i", *c) for c in chunks)
        return s + struct.pack("<Ii", 0x12345678, 0)
    header_u = len(summary([])) + 16                       # uncompressed header (some spare bytes)
    chunks, blobs = [], []
    coff = len(summary([(0, 0, 0, 0)] * nchunks)) + 16
    u = header_u
    for c in range(nchunks):
        part = body[(u - header_u):(u - header_u) + block_size * chunk_blocks]
        blocks = [part[i:i + block_size] for i in range(0, len(part), block_size)] or [b""]
        comp = [lzo.compress_stored(b) for b in blocks]
        hdr = struct.pack("<4I", 0x9E2A83C1, block_size, sum(map(len, comp)), len(part))
        hdr += b"".join(struct.pack("<2I", len(cb), len(b)) for cb, b in zip(comp, blocks))
        blob = hdr + b"".join(comp)
        chunks.append((u, len(part), coff, len(blob)))
        blobs.append(blob)
        coff += len(blob)
        u += len(part)
    s = summary(chunks)
    out = s + bytes(16) + b"".join(blobs)
    assert len(s) + 16 == chunks[0][2]
    os.makedirs(os.path.dirname(path), exist_ok=True)
    open(path, "wb").write(out)
    return header_u


HEADER_U = None   # uncompressed summary size of synthetic packages (fixed; computed on first use)


def _header_u():
    global HEADER_U
    if HEADER_U is None:
        import tempfile
        with tempfile.TemporaryDirectory() as d:
            HEADER_U = package(os.path.join(d, "x.upk"), b"\0" * 16)
    return HEADER_U


def ue_package(path, objects, prefix=b""):
    """A package with name/import/export tables. objects: [(class name, object name, bytes or callable)]:
    a callable gets {name: index} and the export index map and returns the object's bytes.
    Returns {object name: (absolute offset, size)}."""
    classes = sorted({c for c, _, _ in objects})
    names = ["None", "Core", "Class", "LocalizedSubtitles", "Subtitles"] + classes + [n for _, n, _ in objects]
    idx = {n: i for i, n in enumerate(names)}
    exp_index = {n: i + 1 for i, (_, n, _) in enumerate(objects)}
    blobs = [o(idx, exp_index) if callable(o) else o for _, _, o in objects]
    nt = b"".join(struct.pack("<i", len(n) + 1) + n.encode() + b"\0" + bytes(8) for n in names)
    it = b"".join(struct.pack("<7i", idx["Core"], 0, idx["Class"], 0, 0, idx[c], 0) for c in classes)
    base = _header_u() + len(prefix)
    name_off, imp_off = base, base + len(nt)
    exp_off = imp_off + len(it)
    data_off = exp_off + 72 * len(objects)
    et, pos, offs = b"", data_off, {}
    for (c, n, _), b in zip(objects, blobs):
        et += struct.pack("<7i", -(classes.index(c) + 1), 0, 0, idx[n], 0, 0, 0) + bytes(8)
        et += struct.pack("<ii", len(b), pos) + bytes(4 + 4 + 16 + 4)
        offs[n] = (pos, len(b))
        pos += len(b)
    body = prefix + nt + it + et + b"".join(blobs) + bytes(64)
    package(path, body, tables=(len(names), name_off, len(objects), exp_off, len(classes), imp_off))
    return offs


def _fstr(text, wide=False):
    if wide:
        return struct.pack("<i", -(len(text) + 1)) + text.encode("utf-16-le") + b"\0\0"
    return struct.pack("<i", len(text) + 1) + text.encode("cp1252") + b"\0"


def dialogue(speech_object, texts, wide=()):
    """An RDialogueEvent export (compact tagged properties, docs/formats/subtitles.md) referencing the
    export `speech_object`; texts: {lang: [line, ...]}; languages in `wide` are stored as UTF-16."""
    def build(idx, exp_index):
        content = struct.pack("<I", len(texts))
        for lang, lines in texts.items():
            sub = struct.pack("<I", len(lines))
            for ln in lines:
                sub += struct.pack("<HH", 0x0D, 0) + _fstr(ln, lang in wide) + struct.pack("<HHf", 4, 0x0C, 1.5)
                sub += struct.pack("<H", 0)
            content += struct.pack("<HH", 0x0D, 0) + _fstr(lang)
            content += struct.pack("<HHiiii", 9, 0x0C, idx["Subtitles"], 0, len(sub), 0) + sub + struct.pack("<H", 0)
        b = bytes(4) + struct.pack("<HHi", 0x10, 0x2C, exp_index[speech_object])
        b += struct.pack("<HH", 0x0D, 0x34) + _fstr("Batman")
        b += struct.pack("<HHiiii", 9, 0x9C, idx["LocalizedSubtitles"], 0, len(content), 0) + content
        return b + struct.pack("<H", 0)
    return build


def make_game(root):
    """Create the folder; returns the dict of original WEMs by target."""
    cooked = os.path.join(root, "BmGame", "CookedPCConsole")
    v = [wem(1.0, 300), wem(1.6, 500), wem(0.7, 700)]
    single = wem(2.0, 350)
    write_wad(os.path.join(cooked, "English(US)", "28CCF006_LOC_INT.wad"),
              {SPEECH_ENTRY: grouped(v), SINGLE_ENTRY: single})
    music = {MEDIA[0]: wem(1.5, 600, 48000, 2), MEDIA[1]: wem(3.0, 250, 48000, 2), MEDIA[2]: wem(0.8, 900, 48000, 2)}
    bnk, payload = pack(music)
    amb = grouped([wem(1.2, 200, 32000, 2), wem(0.9, 260, 32000, 2)])
    write_wad(os.path.join(cooked, "SFX", "17705D3E.wad"), {PACK_ID: payload, AMB_ENTRY: amb})
    open(os.path.join(cooked, "SFX", "%08X.bnk" % PACK_ID), "wb").write(bnk)
    os.makedirs(os.path.join(root, "Binaries", "Win32"), exist_ok=True)
    open(os.path.join(root, "Binaries", "Win32", "BatmanAC.exe"), "wb").write(b"MZ")
    originals = {"%08X:%d" % (SPEECH_ENTRY, i): w for i, w in enumerate(v)}
    originals["%08X" % SINGLE_ENTRY] = single
    originals.update({"%08X/%08X" % (PACK_ID, m): w for m, w in music.items()})
    originals.update(make_banks(cooked, v))
    make_localization(root)
    return originals


# soundbanks (docs/formats/banks.md): media ids of the in-memory sounds
FOOT = (0x0A000001, 0x0A000002)       # loose bank CHR_SynthFootsteps_Media
FOLEY = (0x0B000001, 0x0B000002)      # bank CHR_SynthFoley, copied into both packages
GUI = (0x0C000001,)                   # bank GUI_Synth (package B only)
LINE_MEDIA = (0x0D000000, 0x0D000001, 0x0D000002)   # INT media ids of the SPEECH_ENTRY variations
HEAD = 0x200                          # prefetch head size of variation 1


def make_banks(cooked, speech_wems):
    """Loose bank + two packages (one bank in both) + a prefetch head of INT SPEECH_ENTRY variation 1."""
    from ackit.bnk import fnv1
    names = ["Init", "CHR_SynthFootsteps_Media", "CHR_SynthFoley", "GUI_Synth", "Speech_Synth"]
    with open(os.path.join(cooked, "CookedBanks.txt"), "w", newline="\r\n") as f:
        f.write("\n".join(names) + "\n")
    foot = {FOOT[0]: wem(0.4, 800), FOOT[1]: wem(0.3, 900)}
    foley = {FOLEY[0]: wem(0.5, 1000), FOLEY[1]: wem(0.6, 1100, 32000)}
    gui = {GUI[0]: wem(0.5, 1200)}
    open(os.path.join(cooked, "SFX", "%08X.bnk" % fnv1(names[1])), "wb").write(
        bank(fnv1(names[1]), foot, name=names[1]))
    # BKPK table of SPEECH_ENTRY for INT (media ids differ per language; offsets as in the WAD entry)
    offs, pos = [], 0
    for w in speech_wems:
        pos = _a(pos, 0x800)
        offs.append(pos)
        pos += len(w)
    tbl = bytearray(0x48)
    struct.pack_into("<4sII", tbl, 0, b"BKPK", 0x01020000, SPEECH_ENTRY)
    struct.pack_into("<IIIIII", tbl, 0x10, 3, 0, 36, 0x48, _a(pos, 0x800), 0)
    tbl += b"".join(struct.pack("<III", m, len(w), o) for m, w, o in zip(LINE_MEDIA, speech_wems, offs))
    speech = bank(fnv1(names[4]), {LINE_MEDIA[1]: speech_wems[1][:HEAD]}, prefetch=(LINE_MEDIA[1],))
    filler = bytes(range(256)) * 9
    synth_a = filler + bank(fnv1(names[2]), foley) + filler[:777] + bytes(tbl) + speech + filler
    synth_b = filler[:1000] + bank(fnv1(names[3]), gui) + filler + bank(fnv1(names[2]), foley) + filler
    package(os.path.join(cooked, "Synth_A.upk"), synth_a)
    package(os.path.join(cooked, "Synth_B.upk"), synth_b)
    # dialogue (subtitles) in two packages; the first one also holds a bank with an in-memory sound and the
    # event that links the dialogue's speech event to INT variation 1 of SPEECH_ENTRY
    dlg = dialogue(SPEECH_EVENT, SUBTITLES, wide=("RUS",))
    dlg_wem = wem(0.3, 1400)
    ev_bank = bank(fnv1("NPC_SynthDialogue"), {DLG_MEDIA: dlg_wem}, name="NPC_SynthDialogue",
                   events=[(SPEECH_EVENT, LINE_MEDIA[1])])
    objs = [("AkEvent", SPEECH_EVENT, bytes(8)), ("RDialogueEvent", "Radio_02", dlg)]
    ue_package(os.path.join(cooked, "Synth_Dialogue.upk"), objs + [("AkBank", "NPC_SynthDialogue", ev_bank)])
    ue_package(os.path.join(cooked, "Synth_Dialogue2.upk"), objs, prefix=bytes(3000))
    with open(os.path.join(cooked, "CookedBanks.txt"), "a", newline="\r\n") as f:
        f.write("NPC_SynthDialogue\n")
    out = {"%08X" % m: w for m, w in list(foot.items()) + list(foley.items()) + list(gui.items())}
    out["%08X" % DLG_MEDIA] = dlg_wem
    return out


SPEECH_EVENT = "Speech_Synth_Radio_02"
SUBTITLES = {"INT": ["Scanner line two.", "Second part."], "ITA": ["Riga due dello scanner."],
             "RUS": ["Строка два"]}
DLG_MEDIA = 0x0E000001
MENU = {"INT": ("FRONTEND.int", "\r\n[Menu]\r\nStart=Press any key\r\nQuit=\"Quit game\"\r\n; comment\r\n"
                                "[Tips]\r\nTip=Use the grapnel\r\nTip=Glide over the city\r\n"),
        "ITA": ("FRONTEND.ITA", "\r\n[Menu]\r\nStart=Premi un tasto\r\nQuit=\"Esci\"\r\n"
                                "[Tips]\r\nTip=Usa il rampino\r\nTip=Plana sulla citta\r\n")}


def make_localization(root):
    for lang, (fname, text) in MENU.items():
        d = os.path.join(root, "BmGame", "Localization", lang)
        os.makedirs(d, exist_ok=True)
        open(os.path.join(d, fname), "wb").write(b"\xff\xfe" + text.encode("utf-16-le"))
    open(os.path.join(root, "BmGame", "Localization", "INT", "Version.int"), "wb").write(b"[ID]\r\nUnreal=SYNTH\r\n")