"""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

Location = namedtuple("Location", "media file offset size kind sound", 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 ---------------------------------------------------------------------------------------
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.packs = defaultdict(dict)       # WAD entry id -> {media id: (offset in entry, size)} (BKPK tables)
        self.issues = []                     # (kind, detail)


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


def bank_languages(data, start, imports, exports, banks, issues=None, rel=""):
    """{bank offset in data: language} for localized banks. A localized AkBank export starts with its list of
    languages (FStrings, e.g. INT FRA ITA DEU ESN) and holds one bank per language, in that order
    (docs/formats/banks.md). Banks of other exports are language-neutral (not in the result)."""
    out = {}
    akbank = {-(i + 1) for i, (c, o, _) in enumerate(imports) if c == "Class" and o == "AkBank"}
    starts = sorted(bk.at for bk in banks)
    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
        inside = [a for a in starts if o0 <= a < o0 + x.size]
        if len(inside) != len(langs):
            if issues is not None:
                issues.append(("languages", "%s: %s has %d banks for %d languages" % (rel, x.name, len(inside), len(langs))))
            continue
        out.update(zip(inside, langs))
    return out


def _bkpk_tables(data, packs):
    """BKPK tables (docs/formats/wad.md) embedded in a package: WAD entry -> media offsets/sizes."""
    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
        entry, n = struct.unpack_from("<I", data, p + 8)[0], struct.unpack_from("<I", data, p + 0x10)[0]
        if n > 100000 or p + 0x48 + 12 * n > len(data):
            continue
        for i in range(n):
            mid, size, off = struct.unpack_from("<III", data, p + 0x48 + 12 * i)
            packs[entry][mid] = (off, size)


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)
            _bkpk_tables(data, sc.packs)
            if visit:
                visit(rel, r, start, data)
            banks = list(find_banks(data))
            lang_of = bank_languages(data, start, imports, exports, banks, sc.issues, rel)
        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))
            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():
                if s.stream != PREFETCH:
                    sc.sound_media[s.object].add(s.media)
    if progress:
        progress(len(files), len(files))
    return sc


# --- names --------------------------------------------------------------------------------------
def event_media(sc):
    """{event name: [media ids in traversal order]} for every event whose name is known."""
    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:
                for m in sorted(sc.sound_media[o]):
                    if m not in order:
                        order.append(m)
            stack.extend(list(kids.get(o, ()))[::-1])
        if order:
            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:
                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"))