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


def _bkpk_tables(data, packs):
    """BKPK tables (docs/formats/wad.md) embedded in a package: WAD entry -> media offsets/sizes."""
    import re
    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
        if isinstance(r, Package):
            names, _, 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))
        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()})
                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_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>)."""
    ev = event_media(sc)
    best = {}
    for name, ms in ev.items():
        rank = (0 if name.lower().startswith("play") else 1, len(ms), name.lower())
        for i, m in enumerate(ms):
            if m in sc.media and (m not in best or rank < best[m][0]):
                best[m] = (rank, name, i, len(ms))
    out = {}
    taken = set()
    for m, (_, name, i, n) in sorted(best.items(), key=lambda x: (x[1][1].lower(), x[1][2])):
        nm = name if n == 1 else "%s_%d" % (name, i + 1)
        if nm.lower() in taken:
            nm = "%s_%08X" % (nm, m)
        taken.add(nm.lower())
        out[m] = nm
    # the rest: by bank, numbered in media id order
    per_bank = defaultdict(list)
    for m, e in sc.media.items():
        bank = sorted(sc.bank_names.get(b, "%08X" % b) for b in e["banks"])[0]
        per_bank[bank].append(m)
        if m in out:
            out[m] = (out[m], bank)
    for bank, ms in per_bank.items():
        rest = sorted(m for m in ms if m not in out)
        for i, m in enumerate(rest):
            nm = "%s_%d" % (bank, i + 1)
            if nm.lower() in taken:
                nm = "%s_%08X" % (bank, m)
            taken.add(nm.lower())
            out[m] = (nm, 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"))