"""Index of every sound in the game's WAD archives, with names, categories and characters
(sounds stored inside soundbanks are not indexed yet: docs/formats/banks.md).

Sources, in order:
  1. data/catalog/<ARCHIVE>.csv  -- shipped catalog (instant) when the WAD fingerprint matches;
  2. otherwise a header scan of the WAD (a few seconds, cached in the user cache), classified with
     `classify` and labelled with data/catalog/characters.csv (speaker per sound name).
Names follow the community hash list (ackit/data/hashlist.txt): an entry with one WEM is NAME
(alias NAME_1), entries with n variations are NAME_1..NAME_n (alias NAME#i); stream-pack media are
named by their ID or <Pack>_<ID>.
"""
import collections
import csv
import difflib
import json
import os
import re
import struct
from dataclasses import dataclass, asdict

from . import banks, classify, subtitles
from .game import wad_fingerprint
from .inject import pack_bnk_path
from .pack import Pack
from .paths import catalog_dir, package_data, user_dir
from .wad import Wad

FIELDS = ["archive", "target", "name", "aliases", "group", "category", "bank", "language", "character",
          "characters_inferred", "seconds", "rate", "channels", "size", "sha1", "notes"]
SCAN_VERSION = 2


@dataclass
class Sound:
    archive: str
    target: str
    name: str
    aliases: tuple = ()
    group: str = "Unknown"
    category: str = "Unknown"
    bank: str = ""                  # BANK sounds: the soundbank holding it (name); WAD sounds: ''
    language: str = ""              # BANK sounds: language of a localized speech line ('' = any language)
    character: str = ""             # verified speaker(s), '; '-separated (curated catalog)
    characters_inferred: str = ""   # characters/roles mentioned by the name (hint)
    seconds: float = 0.0
    rate: int = 0
    channels: int = 0
    size: int = 0                   # size of the ORIGINAL WEM (bytes) when known from the catalog
    sha1: str = ""                  # first 16 hex digits of the original WEM's SHA-1 ('' if unknown)
    notes: str = ""

    @property
    def key(self):
        return "%s:%s" % (self.archive, self.target)

    @property
    def is_pack_media(self):
        return "/" in self.target

    @property
    def limit(self):
        """Max replacement size in bytes (None = any size)."""
        return None if self.is_pack_media else self.size

    def names(self):
        return (self.name,) + tuple(self.aliases)


# --- static data ---------------------------------------------------------------------------
_names = None


def hash_names():
    global _names
    if _names is None:
        _names = {}
        for ln in open(package_data("hashlist.txt"), encoding="latin-1"):
            p = ln.rstrip("\r\n").split("\t")
            if len(p) >= 2 and p[1]:
                try:
                    _names[int(p[0], 16)] = p[1]
                except ValueError:
                    pass
    return _names


def load_characters(path=None):
    """{sound name (lower): 'Speaker' or 'A; B'} from data/catalog/characters.csv."""
    path = path or os.path.join(catalog_dir(), "characters.csv")
    out = {}
    if os.path.exists(path):
        for row in csv.DictReader(open(path, encoding="utf-8")):
            out[row["name"].lower()] = row["character"]
    return out


def read_csv(path):
    out = []
    for r in csv.DictReader(open(path, encoding="utf-8", newline="")):
        out.append(Sound(archive=r["archive"], target=r["target"], name=r["name"],
                         aliases=tuple(a for a in r["aliases"].split("|") if a), group=r["group"],
                         category=r["category"], bank=r.get("bank") or "", language=r.get("language") or "",
                         character=r["character"],
                         characters_inferred=r["characters_inferred"], seconds=float(r["seconds"] or 0),
                         rate=int(r["rate"] or 0), channels=int(r["channels"] or 0), size=int(r["size"] or 0),
                         sha1=r["sha1"], notes=r["notes"]))
    return out


def write_csv(path, sounds):
    with open(path, "w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, FIELDS, lineterminator="\n")
        w.writeheader()
        for s in sounds:
            d = asdict(s)
            d["aliases"] = "|".join(s.aliases)
            d["seconds"] = "%.3f" % s.seconds
            w.writerow({k: d[k] for k in FIELDS})


# --- scanning ------------------------------------------------------------------------------
def _hdr(b):
    """(rate, channels, seconds) from the first bytes of a WEM (fmt is always the first chunk)."""
    if len(b) >= 0x30 and b[:4] == b"RIFF" and b[12:16] == b"fmt ":
        ch, rate = struct.unpack_from("<HI", b, 22)
        frames = struct.unpack_from("<I", b, 44)[0]
        return rate, ch, (frames / float(rate) if rate else 0.0)
    return 0, 0, 0.0


def _walk_entry(f, offset, size):
    """Offsets/lengths of the WEMs inside one WAD entry, reading only headers."""
    out = []
    pos = 0
    while pos + 12 <= size:
        f.seek(offset + pos)
        h = f.read(0x30)
        if h[:4] != b"RIFF" or h[8:12] != b"WAVE":
            break
        ln = 8 + struct.unpack_from("<I", h, 4)[0]
        out.append((pos, ln, h))
        # next WEM starts at the first non-zero byte after this one (zero padding in between)
        cand = pos + ln
        f.seek(offset + cand)
        nxt = None
        while cand < size and nxt is None:
            chunk = f.read(min(0x10000, size - cand))
            if not chunk:
                break
            stripped = chunk.lstrip(b"\0")
            if stripped:
                nxt = cand + len(chunk) - len(stripped)
            cand += len(chunk)
        if nxt is None:
            break
        pos = nxt
    return out


def scan_archive(archive, wad_path, characters=None, progress=None):
    """Header scan of one WAD -> [Sound]."""
    names = hash_names()
    characters = characters if characters is not None else load_characters()
    wad = Wad(wad_path)
    out = []
    with open(wad_path, "rb") as f:
        for n, e in enumerate(wad.entries):
            if progress and n % 200 == 0:
                progress(n, len(wad.entries))
            bnk = pack_bnk_path(wad_path, e.hash)
            if bnk and open(bnk, "rb").read(4) == b"BKPK":
                pk = Pack(open(bnk, "rb").read())
                pname = names.get(e.hash, "%08X" % e.hash)
                for mid, size, off in pk.entries:
                    f.seek(e.offset + off)
                    rate, ch, secs = _hdr(f.read(0x30))
                    out.append(Sound(archive, "%08X/%08X" % (e.hash, mid), names.get(mid, "%s_%08X" % (pname, mid)),
                                     seconds=secs, rate=rate, channels=ch, size=size))
                    out[-1].group, out[-1].category = classify.classify(out[-1].name, pack=pname)
                continue
            wems = _walk_entry(f, e.offset, e.size)
            base = names.get(e.hash, "%08X" % e.hash)
            for i, (off, ln, h) in enumerate(wems):
                rate, ch, secs = _hdr(h)
                if len(wems) == 1:
                    s = Sound(archive, "%08X" % e.hash, base, (base + "_1",))
                else:
                    s = Sound(archive, "%08X:%d" % (e.hash, i), "%s_%d" % (base, i + 1), ("%s#%d" % (base, i + 1),))
                s.seconds, s.rate, s.channels, s.size = secs, rate, ch, ln
                s.group, s.category = classify.classify(s.name)
                out.append(s)
    # a variation named NAME_i can collide with an entry whose own name is NAME_i: the variation is
    # shown as NAME#i and keeps NAME_i as alias (both stay findable, display names stay unique)
    taken = collections.Counter(s.name.lower() for s in out)
    for s in out:
        if taken[s.name.lower()] > 1 and ":" in s.target:
            s.name, s.aliases = s.aliases[0], (s.name,)
    for s in out:
        s.characters_inferred = "; ".join(classify.infer_characters(s.name))
        s.character = next((characters[a.lower()] for a in s.names() if a.lower() in characters), "")
    return out


_CAMEL = re.compile(r"(?<=[a-z])(?=[A-Z])")


def bank_sounds(sc, characters=None):
    """[Sound] of the in-memory sounds found by `banks.scan` (prefetch heads belong to WAD sounds)."""
    characters = characters if characters is not None else load_characters()
    names = banks.media_names(sc)
    first = {}
    for loc in sorted(sc.locations, key=lambda l: (l.media, l.file, l.offset)):
        first.setdefault(loc.media, loc)
    out = []
    for m, e in sc.media.items():
        if banks.PREFETCH_HEAD in e["kinds"]:
            continue
        name, bank = names[m]
        rate, ch, secs = _hdr(e["head"])
        s = Sound(banks.ARCHIVE, "%08X" % m, name, ("%s/%08X" % (bank, m),), bank=bank,
                  language=banks.media_language(sc, m), seconds=secs, rate=rate, channels=ch,
                  size=min(l.size for l in sc.locations if l.media == m) if len(e["sha1"]) > 1 else e["size"],
                  sha1=next(iter(e["sha1"])) if len(e["sha1"]) == 1 else "")
        s.group, s.category = classify.classify(name, bank=bank)
        s.characters_inferred = "; ".join(classify.infer_characters(_CAMEL.sub("_", name + "_" + bank)))
        s.character = next((characters[a.lower()] for a in s.names() if a.lower() in characters), "")
        out.append(s)
    out.sort(key=lambda s: (s.bank.lower(), s.name.lower()))
    return out


def head_owners(game, sc, archives=None):
    """{prefetch media id: WAD sound key} -- the speech line (one language) each prefetch head belongs to."""
    return media_owners(game, sc, {m for m, e in sc.media.items() if banks.PREFETCH_HEAD in e["kinds"]}, archives)


def media_owners(game, sc, want, archives=None):
    """{media id: WAD sound key} for streamed speech media, found through the per-language BKPK tables of
    the packages (docs/formats/banks.md): every language has its own media IDs for the same WAD entry."""
    out = {}
    for a in archives or [x for x in game.archives if x != "SFX"]:
        wad = Wad(game.wad(a))
        by_hash = {e.hash: e for e in wad.entries}
        with open(game.wad(a), "rb") as f:
            for ent, table in sc.packs.items():
                mine = [m for m in table if m in want and m not in out]
                e = by_hash.get(ent)
                if not mine or not e:
                    continue
                wems = _walk_entry(f, e.offset, e.size)
                for m in mine:
                    off, size = table[m]
                    for i, (o, ln, _) in enumerate(wems):
                        if (o, ln) == (off, size):
                            out[m] = "%s:%08X" % (a, ent) + (":%d" % i if len(wems) > 1 else "")
    return out


def scan_banks(game, progress=None, characters=None, visit=None):
    """Full scan of the game's banks -> ([Sound], [Location] with WAD keys on prefetch heads, Scan)."""
    sc = banks.scan(game.cooked, progress, visit=visit)
    owners = head_owners(game, sc)
    locs = [l._replace(sound=owners.get(l.media, "")) if l.kind == banks.PREFETCH_HEAD else l for l in sc.locations]
    return bank_sounds(sc, characters), locs, sc


def subtitle_links(game, sc, speech_names):
    """{speech event name: {language: sound key}}: event -> actions -> containers -> sounds -> the
    per-language media of the line (exact, no name matching). The media is a streamed WAD sound, or an
    in-memory BANK sound of a localized bank."""
    ev = banks.event_media(sc)
    want = {m for n in speech_names for m in ev.get(n, ())}
    owners = {m: (k.split(":")[0], k) for m, k in media_owners(game, sc, want).items()}   # m -> (lang, key)
    for m in want - set(owners):
        lang = banks.media_language(sc, m) if m in sc.media else ""
        if lang:
            owners[m] = (lang, "%s:%08X" % (banks.ARCHIVE, m))
    out = {}
    for n in speech_names:
        keys = {}
        for m in ev.get(n, ()):
            if m in owners:
                lang, key = owners[m]
                keys.setdefault(lang, key)
        if keys:
            out[n] = keys
    return out


def scan_text(game, progress=None, characters=None):
    """One pass over every package: (bank sounds, bank locations, Scan, {lang: [SubLine]}, [SubLoc], issues)."""
    col = subtitles.Collector()
    snd, locs, sc = scan_banks(game, progress, characters, visit=col)
    links = subtitle_links(game, sc, list(col.dialogues))
    texts, sublocs = subtitles.rows(col, links)
    return snd, locs, sc, texts, sublocs, col.issues


def load_subtitles(game, lang, progress=None):
    """([SubLine], {SUB key: [SubLoc]}) of one language: shipped catalog when the game matches it, else a
    cached or fresh scan of the packages (several minutes)."""
    fp = banks.game_fingerprint(game.cooked)
    shipped = _shipped_index().get("SUB", {})
    t = os.path.join(catalog_dir(), "SUB_%s.csv" % lang)
    loc = os.path.join(catalog_dir(), subtitles.LOCATIONS)
    if shipped.get("fingerprint") == fp and os.path.exists(t) and os.path.exists(loc):
        return subtitles.read_texts(t, lang), subtitles.read_locations(loc, lang)
    cache = os.path.join(user_dir("cache"), "scan_SUB_%s_v%d" % (fp, SCAN_VERSION))
    if not os.path.exists(cache + ".locations.csv.gz"):
        _, _, _, texts, sublocs, _ = scan_text(game, progress)
        for lg, rows in texts.items():
            subtitles.write_texts(cache + "_%s.csv" % lg, rows)
        subtitles.write_locations(cache + ".locations.csv.gz", sublocs)
    t = cache + "_%s.csv" % lang
    return (subtitles.read_texts(t, lang) if os.path.exists(t) else []), \
        subtitles.read_locations(cache + ".locations.csv.gz", lang)


def load_banks(game, progress=None):
    """([Sound], {media: [Location]}) of the in-memory sounds: shipped catalog when the game's bank set
    matches it, else a cached or fresh scan (several minutes)."""
    fp = banks.game_fingerprint(game.cooked)
    shipped = _shipped_index().get(banks.ARCHIVE, {})
    csv_path = os.path.join(catalog_dir(), "%s.csv" % banks.ARCHIVE)
    loc_path = os.path.join(catalog_dir(), banks.LOCATIONS)
    if shipped.get("fingerprint") == fp and os.path.exists(csv_path) and os.path.exists(loc_path):
        return read_csv(csv_path), banks.read_locations(loc_path)
    cache = os.path.join(user_dir("cache"), "scan_%s_%s_v%d" % (banks.ARCHIVE, fp, SCAN_VERSION))
    if not (os.path.exists(cache + ".csv") and os.path.exists(cache + ".locations.csv.gz")):
        sounds, locs, _ = scan_banks(game, progress)
        for s in sounds:
            s.sha1 = ""              # unknown build: no trusted original checksums
        banks.write_locations(cache + ".locations.csv.gz", locs)
        write_csv(cache + ".csv", sounds)
    return read_csv(cache + ".csv"), banks.read_locations(cache + ".locations.csv.gz")


# --- catalog -------------------------------------------------------------------------------
def _shipped_index():
    p = os.path.join(catalog_dir(), "index.json")
    try:
        return json.load(open(p))
    except (OSError, ValueError):
        return {}


def is_shipped(game, archive):
    """True if the game's archive matches the shipped catalog (not a scan of another build)."""
    shipped = _shipped_index().get(archive, {})
    if archive == banks.ARCHIVE:
        return shipped.get("fingerprint") == banks.game_fingerprint(game.cooked)
    return shipped.get("fingerprint") == wad_fingerprint(game.wad(archive))


def load_archive(game, archive, progress=None):
    """[Sound] for one archive of a game (shipped CSV, user-cache scan, or fresh scan)."""
    path = game.wad(archive)
    fp = wad_fingerprint(path)
    shipped = _shipped_index().get(archive, {})
    csv_path = os.path.join(catalog_dir(), "%s.csv" % archive)
    if shipped.get("fingerprint") == fp and os.path.exists(csv_path):
        return read_csv(csv_path)
    cache = os.path.join(user_dir("cache"), "scan_%s_%s_v%d.csv" % (archive, fp, SCAN_VERSION))
    if os.path.exists(cache):
        return read_csv(cache)
    sounds = scan_archive(archive, path, progress=progress)
    # unknown build: no original checksums (states then rely on the originals store only)
    write_csv(cache, sounds)
    return sounds


def available_archives(game):
    """WAD archives of the game plus BANK when it has soundbanks."""
    return list(game.archives) + ([banks.ARCHIVE] if banks.has_banks(game.cooked) else [])


class Catalog:
    """All sounds of the selected archives (default: the game's speech language, SFX and BANK)."""

    def __init__(self, game, archives=None, progress=None):
        self.game = game
        avail = available_archives(game)
        self.archives = archives or [a for a in (game.language(), "SFX", banks.ARCHIVE) if a in avail]
        self.sounds = []
        self.locations = {}          # BANK media id -> [Location]
        self.heads = {}              # WAD sound key -> [Location] of its prefetch heads
        bank_sounds_ = []
        if banks.ARCHIVE in avail:
            # always loaded: speech lines may have prefetch heads that must be written with them
            bank_sounds_, self.locations = load_banks(game, progress)
            for ls in self.locations.values():
                for l in ls:
                    if l.kind == banks.PREFETCH_HEAD and l.sound:
                        self.heads.setdefault(l.sound, []).append(l)
        # localized in-memory speech: only the languages loaded here (all of them if none is)
        langs = [a for a in self.archives if a in game.languages] or None
        bank_sounds_ = [s for s in bank_sounds_ if not s.language or not langs or s.language in langs]
        for a in self.archives:
            self.sounds += bank_sounds_ if a == banks.ARCHIVE else load_archive(game, a, progress)
        self.by_key = {s.key: s for s in self.sounds}
        self._by_name = {}
        for s in self.sounds:
            for n in s.names():
                self._by_name.setdefault(n.lower(), []).append(s)

    def get(self, key):
        return self.by_key[key]

    def resolve(self, name):
        """One sound by name, alias, key ('ITA:B49CDF48:11') or file name ('X_3.wav')."""
        if name in self.by_key:
            return self.by_key[name]
        key = os.path.splitext(os.path.basename(name))[0].lower() if re.search(r"\.\w{2,4}$", name) else name.lower()
        hits = self._by_name.get(key, [])
        uniq = list({id(s): s for s in hits}.values())
        if len(uniq) == 1:
            return uniq[0]
        if len(uniq) > 1:
            raise LookupError("'%s' is ambiguous; use one of: %s" % (name, ", ".join(
                next((a for a in s.names() if a.lower() != key), s.key) for s in uniq)))
        near = [s.name for s in self.search(key, limit=8)] or self.suggest(key)
        raise LookupError("no sound named '%s'.%s" % (name, ("\n  did you mean: " + ", ".join(near)) if near else ""))

    def search(self, text, limit=None):
        """All words must appear in the name, category or character (case-insensitive)."""
        words = [w for w in re.split(r"[\s_]+", text.lower()) if w]
        res = [s for s in self.sounds
               if all(w in s.name.lower() or w in s.category.lower() or w in s.character.lower() for w in words)]
        res.sort(key=lambda s: (len(s.name), s.name))
        return res[:limit] if limit else res

    def suggest(self, text, n=8):
        names = {s.name.lower(): s.name for s in self.sounds}
        return [names[k] for k in difflib.get_close_matches(text.lower(), list(names), n=n, cutoff=0.6)]