"""In-memory sounds: WEMs held inside soundbanks, i.e. the loose SFX\\*.bnk files and the banks embedded
in the .upk packages (docs/formats/banks.md, upk.md). Catalog archive `BANK`, target = media ID.

A media ID can be stored many times (the same bank copied into several packages, or several banks
holding the same media). The location index (data/catalog/BANK_locations.csv.gz) lists every copy:
  media, file (relative to CookedPCConsole, '/'-separated), offset (in the *uncompressed* file), size, kind
kind 'm' = an in-memory sound, 'p' = the in-memory head of a streamed WAD sound (prefetch); for 'p' rows
`sound` is the WAD sound's key (e.g. ITA:2B14E472:3): speech lines have per-language media IDs and heads.
Replacements are written in place into every copy (<= original size, zero-padded), so bank and package
layouts never change.
"""
import csv
import gzip
import hashlib
import io
import os
import re
import struct
from collections import defaultdict, namedtuple

from .bnk import Bank, BankError, find_banks, fnv1, PREFETCH
from .upk import Package, PackageError, is_package

ARCHIVE = "BANK"
LOCATIONS = "BANK_locations.csv.gz"
SOUND, MEMORY, PREFETCH_HEAD = 2, "m", "p"
EVENT, ACTION = 4, 3
CONTAINERS = (5, 6, 7, 9)            # random/sequence, switch, actor-mixer, layer

# lang: language of the bank holding this copy ('' = language-neutral bank); known after a scan only
Location = namedtuple("Location", "media file offset size kind sound lang", defaults=("", ""))


# --- files ----------------------------------------------------------------------------------------
def bank_files(cooked):
    """Relative paths ('/'-separated) of every loose regular bank and every package."""
    out = []
    sfx = os.path.join(cooked, "SFX")
    if os.path.isdir(sfx):
        for f in sorted(os.listdir(sfx)):
            if f.lower().endswith(".bnk"):
                with open(os.path.join(sfx, f), "rb") as fh:
                    if fh.read(4) == b"BKHD":
                        out.append("SFX/" + f)
    for f in sorted(os.listdir(cooked)):
        if f.lower().endswith(".upk") and is_package(os.path.join(cooked, f)):
            out.append(f)
    return out


class _Plain:
    """Uniform read access for a loose bank (offsets = file offsets)."""

    def __init__(self, path):
        self.path = path

    def read(self, off, n):
        with open(self.path, "rb") as f:
            f.seek(off)
            return f.read(n)

    def read_all(self):
        return 0, open(self.path, "rb").read()

    def fingerprint(self):
        return file_fingerprint(self.path)

    def patched(self, edits):
        out = bytearray(open(self.path, "rb").read())
        for off, data in edits:
            out[off:off + len(data)] = data
        return bytes(out)


def file_fingerprint(path):
    """Identity of a loose bank (small files: content hash)."""
    return hashlib.sha1(open(path, "rb").read()).hexdigest()[:16]


_open_cache = {}


def open_file(path):
    """A reader for a bank file or package (cached while the file is unchanged)."""
    st = os.stat(path)
    key = (os.path.normcase(os.path.abspath(path)), st.st_size, st.st_mtime_ns)
    r = _open_cache.get(key)
    if r is None:
        r = Package(path) if path.lower().endswith(".upk") else _Plain(path)
        if len(_open_cache) > 64:
            _open_cache.clear()
        _open_cache[key] = r
    return r


def file_path(cooked, rel):
    return os.path.join(cooked, *rel.split("/"))


def riff_trim(b):
    """A WEM slot without the zero padding after a smaller replacement (RIFF length)."""
    if len(b) >= 8 and b[:4] == b"RIFF":
        n = 8 + struct.unpack_from("<I", b, 4)[0]
        if n <= len(b):
            return b[:n]
    return b


# --- HIRC graph (for names) ---------------------------------------------------------------------
def _hirc_objects(buf, bank):
    """(type, id, payload memoryview) of every HIRC object of a bank."""
    if "HIRC" not in bank.chunks:
        return
    o, n = bank.chunks["HIRC"]
    p = bank.at + o
    end = p + n
    count = struct.unpack_from("<I", buf, p)[0]
    p += 4
    for _ in range(count):
        if p + 9 > end:
            return
        t = buf[p]
        ln = struct.unpack_from("<I", buf, p + 1)[0]
        yield t, struct.unpack_from("<I", buf, p + 5)[0], bytes(buf[p + 5:p + 5 + ln])
        p += 5 + ln


def _children(payload, known):
    """A container's child list: the first `u32 count, count x u32 id` whose ids are all known objects."""
    n = len(payload)
    for p in range(4, n - 7):
        c = struct.unpack_from("<I", payload, p)[0]
        if 0 < c <= 4096 and p + 4 + 4 * c <= n:
            ids = struct.unpack_from("<%dI" % c, payload, p + 4)
            if all(i in known for i in ids):
                return ids
    return ()


# --- scan ---------------------------------------------------------------------------------------
def _by_language():
    return defaultdict(dict)


class Scan:
    """Everything found in the game's banks (built by `scan`)."""

    def __init__(self):
        self.locations = []                  # [Location]
        self.files = {}                      # rel -> fingerprint
        self.media = {}                      # id -> {"size", "sha1": set, "head": bytes, "banks": set}
        self.bank_names = {}                 # bank id -> name
        self.names = set()                   # identifier-like names seen in packages (for event names)
        self.objects = {}                    # HIRC id -> type
        self.containers = {}                 # id -> payload
        self.events = {}                     # id -> [action ids]
        self.actions = {}                    # id -> target id
        self.sound_media = defaultdict(set)  # sound object id -> media ids
        self.prefetch = set()                # media ids whose in-memory copy is a prefetch head
        self.lang_media = defaultdict(set)   # language -> media ids named by Sound objects of that language's banks
        # BKPK tables: WAD entry id -> {language: {media id: (offset in entry, size)}}. A table belongs to the
        # language its localized AkBank export declares ('' = outside such an export). The same media ID can be
        # used by several languages (at a different position in each language's WAD), so tables never merge.
        self.packs = defaultdict(_by_language)
        self.issues = []                     # (kind, detail)


_LANG = re.compile(rb"([\x01-\x10])\x00\x00\x00\x04\x00\x00\x00[A-Z]{3}\x00")


def localized_exports(data, start, imports, exports):
    """[(begin, end, [languages])] of the localized AkBank exports of a package (offsets in data). Such an
    export starts with its list of languages (FStrings, e.g. INT FRA ITA DEU ESN) and then holds one block per
    language, in that order: a bank and/or the BKPK table of the speech WAD entry (docs/formats/banks.md)."""
    out = []
    akbank = {-(i + 1) for i, (c, o, _) in enumerate(imports) if c == "Class" and o == "AkBank"}
    for x in exports:
        if x.cls not in akbank or x.size <= 0:
            continue
        o0 = x.offset - start
        m = _LANG.search(data, o0, min(o0 + 256, o0 + x.size))
        if not m:
            continue
        n = m.group(1)[0]
        langs, o = [], m.start() + 4
        for _ in range(n):
            ln = struct.unpack_from("<i", data, o)[0]
            langs.append(bytes(data[o + 4:o + 3 + ln]).decode("latin-1"))
            o += 4 + ln
        out.append((o0, o0 + x.size, langs, x.name))
    return out


def _assign(positions, regions, issues, rel, what):
    """{position: language} for the blocks (banks or BKPK tables) inside localized exports: the k-th block of
    an export belongs to its k-th language. An export whose block count differs is reported, not guessed."""
    out = {}
    for o0, o1, langs, name in regions:
        inside = [a for a in positions if o0 <= a < o1]
        if not inside:
            continue
        if len(inside) != len(langs):
            if issues is not None:
                issues.append(("languages", "%s: %s has %d %s for %d languages" % (rel, name, len(inside), what, len(langs))))
            continue
        out.update(zip(inside, langs))
    return out


def bkpk_positions(data):
    """Offsets of the BKPK tables (docs/formats/wad.md) embedded in a package."""
    out = []
    for m in re.finditer(b"BKPK", data):
        p = m.start()
        if p + 0x48 > len(data) or struct.unpack_from("<I", data, p + 0x1C)[0] != 0x48:
            continue
        n = struct.unpack_from("<I", data, p + 0x10)[0]
        if n <= 100000 and p + 0x48 + 12 * n <= len(data):
            out.append(p)
    return out


def read_bkpk(data, p):
    """(WAD entry id, {media id: (offset in entry, size)}) of the BKPK table at p."""
    entry, n = struct.unpack_from("<I", data, p + 8)[0], struct.unpack_from("<I", data, p + 0x10)[0]
    table = {}
    for i in range(n):
        mid, size, off = struct.unpack_from("<III", data, p + 0x48 + 12 * i)
        table[mid] = (off, size)
    return entry, table


def _bkpk_tables(data, packs, regions, issues=None, rel=""):
    """Add a package's BKPK tables to `packs` under the language their export declares. Copies of a table
    in several packages must agree (reported otherwise)."""
    pos = bkpk_positions(data)
    lang_at = _assign(pos, regions, issues, rel, "BKPK tables")
    for p in pos:
        entry, table = read_bkpk(data, p)
        lang = lang_at.get(p, "")
        have = packs[entry].get(lang)
        if have and have != table and issues is not None:
            issues.append(("bkpk", "%s: table of %08X (%s) differs from another copy" % (rel, entry, lang or "-")))
        packs[entry][lang] = table


def scan(cooked, progress=None, cooked_banks=None, files=None, visit=None):
    """Scan every loose bank and package of the game (minutes), or only `files`. Returns a `Scan`.
    `visit(rel, package, start, data)` is called for every package (other scans share the decompression)."""
    sc = Scan()
    files = bank_files(cooked) if files is None else files
    if cooked_banks is None:
        p = os.path.join(cooked, "CookedBanks.txt")
        cooked_banks = [l.strip() for l in open(p, encoding="latin-1")] if os.path.exists(p) else []
    sc.bank_names.update({fnv1(n): n for n in cooked_banks if n})
    for k, rel in enumerate(files):
        if progress and k % 50 == 0:
            progress(k, len(files))
        path = file_path(cooked, rel)
        try:
            r = open_file(path)
            start, data = r.read_all()
        except (PackageError, OSError) as e:
            sc.issues.append(("unreadable", "%s: %s" % (rel, e)))
            continue
        lang_of = {}
        if isinstance(r, Package):
            names, imports, exports = r.tables(data, start)
            sc.names.update(names)                      # base names ...
            sc.names.update(x.name for x in exports)    # ... and numbered object names (Speech_X_04)
            regions = localized_exports(data, start, imports, exports)
            _bkpk_tables(data, sc.packs, regions, sc.issues, rel)
            if visit:
                visit(rel, r, start, data)
            banks = list(find_banks(data))
            lang_of = _assign(sorted(bk.at for bk in banks), regions, sc.issues, rel, "banks")
        else:
            try:
                banks = [Bank(data, 0)]
            except BankError:
                continue
        if banks:
            sc.files[rel] = r.fingerprint()
        for bk in banks:
            sc.bank_names.update(bk.names())
            media = {m.id: m for m in bk.media()}
            prefetched = set()
            for s in bk.sources():
                if s.stream == PREFETCH and s.media in media:
                    prefetched.add(s.media)
            for m in media.values():
                w = bytes(data[bk.at + m.off:bk.at + m.off + m.size])
                kind = PREFETCH_HEAD if m.id in prefetched else MEMORY
                if kind == PREFETCH_HEAD:
                    sc.prefetch.add(m.id)
                e = sc.media.setdefault(m.id, {"size": m.size, "sha1": set(), "head": w[:0x60], "banks": set(),
                                               "kinds": set(), "langs": set()})
                e["langs"].add(lang_of.get(bk.at, ""))
                e["sha1"].add(hashlib.sha1(w).hexdigest()[:16])
                e["banks"].add(bk.id)
                e["kinds"].add(kind)
                sc.locations.append(Location(m.id, rel, start + bk.at + m.off, m.size, kind, "", lang_of.get(bk.at, "")))
            for t, oid, pl in _hirc_objects(data, bk):
                sc.objects[oid] = t
                if t == EVENT and len(pl) >= 8:
                    n = struct.unpack_from("<I", pl, 4)[0]
                    if len(pl) >= 8 + 4 * n:
                        sc.events[oid] = list(struct.unpack_from("<%dI" % n, pl, 8))
                elif t == ACTION and len(pl) >= 12:
                    sc.actions[oid] = struct.unpack_from("<I", pl, 8)[0]
                elif t in CONTAINERS:
                    sc.containers[oid] = pl
            for s in bk.sources():          # every source, prefetched ones too (streamed lines whose start
                sc.sound_media[s.object].add(s.media)          # is kept in the bank: the game-over taunts)
                if lang_of.get(bk.at):
                    sc.lang_media[lang_of[bk.at]].add(s.media)
    if progress:
        progress(len(files), len(files))
    return sc


# --- names --------------------------------------------------------------------------------------
def event_objects(sc):
    """{event name: [sound object ids in traversal order]} for every event whose name is known. Events,
    actions, containers and sound objects have the same IDs in every language (only the media differ), so
    this order is the same in every language."""
    known = set(sc.objects)
    kids = {cid: _children(pl, known) for cid, pl in sc.containers.items()}
    by_hash = {}
    for n in sc.names:
        if n and len(n) < 128:
            by_hash.setdefault(fnv1(n), n)
    out = {}
    for eid, acts in sc.events.items():
        name = by_hash.get(eid)
        if not name:
            continue
        seen, order, stack = set(), [], [sc.actions.get(a) for a in acts][::-1]
        while stack:
            o = stack.pop()
            if o is None or o in seen:
                continue
            seen.add(o)
            if o in sc.sound_media:
                order.append(o)
            stack.extend(list(kids.get(o, ()))[::-1])
        if order:
            out[name] = order
    return out


def event_media(sc):
    """{event name: [media ids in traversal order]} for every event whose name is known."""
    out = {}
    for name, objs in event_objects(sc).items():
        order = []
        for o in objs:
            for m in sorted(sc.sound_media[o]):
                if m not in order:
                    order.append(m)
        out[name] = order
    return out


def media_language(sc, m):
    """Language of a localized in-memory media ('' = language-neutral)."""
    langs = sc.media[m].get("langs", set()) - {""}
    return next(iter(langs)) if len(langs) == 1 else ""


def media_names(sc):
    """{media id: (display name, primary bank name)}: named after the most specific 'Play' event that
    reaches the media, otherwise after its bank (<Bank>_<n>). Localized speech has one media per language
    for the same line: they share the name, variations are numbered within a language. Names are unique
    among the language-neutral media plus any one language (what a catalog shows at a time)."""
    ev = event_media(sc)
    best = {}
    for name, ms in ev.items():
        groups = defaultdict(list)
        for m in ms:
            if m in sc.media and PREFETCH_HEAD not in sc.media[m]["kinds"]:     # heads belong to WAD lines
                groups[media_language(sc, m)].append(m)
        for g in groups.values():
            rank = (0 if name.lower().startswith("play") else 1, len(g), name.lower())
            for i, m in enumerate(g):
                if m not in best or rank < best[m][0]:
                    best[m] = (rank, name, i, len(g))
    out = {}
    taken = set()                       # (language, lower-case name)

    def free(lang, nm):
        k = nm.lower()
        if lang:
            return (lang, k) not in taken and ("", k) not in taken
        return not any(n == k for _, n in taken)

    def take(lang, nm, m, base):
        if not free(lang, nm):
            nm = "%s_%08X" % (base, m)
        taken.add((lang, nm.lower()))
        return nm

    for m, (_, name, i, n) in sorted(best.items(), key=lambda x: (x[1][1].lower(), x[1][2], x[0])):
        out[m] = take(media_language(sc, m), name if n == 1 else "%s_%d" % (name, i + 1), m, name)
    # the rest: by bank (and language), numbered in media id order
    groups = defaultdict(list)
    for m, e in sc.media.items():
        bank = sorted(sc.bank_names.get(b, "%08X" % b) for b in e["banks"])[0]
        if m in out:
            out[m] = (out[m], bank)
        else:
            groups[(bank, media_language(sc, m))].append(m)
    for (bank, lang), ms in sorted(groups.items()):
        for i, m in enumerate(sorted(ms)):
            out[m] = (take(lang, "%s_%d" % (bank, i + 1), m, bank), bank)
    return out


# --- location index -------------------------------------------------------------------------------
def write_locations(path, locs):
    raw = io.StringIO()
    w = csv.writer(raw, lineterminator="\n")
    w.writerow(["media", "file", "offset", "size", "kind", "sound"])
    for l in sorted(locs, key=lambda l: (l.media, l.file, l.offset)):
        w.writerow(["%08X" % l.media, l.file, "%X" % l.offset, l.size, l.kind, l.sound])
    data = raw.getvalue().encode("utf-8")
    with open(path + ".tmp", "wb") as f:
        # mtime=0: identical content gives identical bytes (no spurious git changes)
        with gzip.GzipFile(fileobj=f, mode="wb", mtime=0) as gz:
            gz.write(data)
    os.replace(path + ".tmp", path)


def read_locations(path):
    """{media id: [Location]}"""
    out = defaultdict(list)
    with gzip.open(path, "rt", encoding="utf-8", newline="") as f:
        for r in csv.DictReader(f):
            m = int(r["media"], 16)
            out[m].append(Location(m, r["file"], int(r["offset"], 16), int(r["size"]), r["kind"], r["sound"]))
    return dict(out)


def game_fingerprint(cooked):
    """Cheap identity of the game's bank set: CookedBanks.txt + names/sizes of the loose banks.
    (Each file is checked again before it is written: `file_fingerprint` / Package.fingerprint.)"""
    h = hashlib.sha1()
    p = os.path.join(cooked, "CookedBanks.txt")
    if os.path.exists(p):
        h.update(open(p, "rb").read())
    sfx = os.path.join(cooked, "SFX")
    if os.path.isdir(sfx):
        for f in sorted(os.listdir(sfx)):
            if f.lower().endswith(".bnk"):
                h.update(("%s:%d;" % (f.upper(), os.path.getsize(os.path.join(sfx, f)))).encode())
    return h.hexdigest()[:16]


def has_banks(cooked):
    return os.path.exists(os.path.join(cooked, "CookedBanks.txt"))