"""High-level operations shared by the CLI and the editor: apply / revert sounds safely.

Safety rules (docs/decisions.md):
  * before the first change to a sound its original bytes go to the originals store;
  * optionally a full copy of each WAD (and pack index) is taken once, as a fallback;
  * each WAD is rewritten atomically; several sounds of one WAD are written in one pass;
  * if a sound in the game is neither its original nor what we applied (Steam update, "verify
    integrity", another tool), we refuse to write it unless forced, and report it.
In-memory sounds (archive BANK) and the prefetch heads of speech lines live in loose banks and .upk
packages (ackit/banks.py); a replacement fills each in-memory slot exactly (wem.fill). Before the first
change to such a file a verified full copy is kept, and the
file is always rebuilt as "original copy + every replacement currently applied", so reverting the last
change restores it byte-for-byte and any later operation repairs an interrupted one.
"""
import hashlib
import os
from collections import defaultdict

from . import audio, banks
from .catalog import Catalog
from .encoder import EncodeError, encode_fit
from .game import Game, GameError
from .inject import inject, read_target, parse_target, pack_bnk_path, SizeLimitError
from .originals import OriginalsStore, sha256, sha1_16
from .pack import Pack
from .texts import Texts
from .paths import user_dir
from .wad import Wad
from .wem import FILL_MIN, fill, fillable, unfill

# sound states
ORIGINAL = "original"      # untouched
MODIFIED = "modified"      # contains what we applied
LOST = "lost"              # we applied something, but the game has the original again (e.g. Steam verify)
CHANGED = "changed"        # neither original nor ours: changed underneath (update / other tool)
FOREIGN = "foreign"        # never touched by us, but differs from the catalog's original checksum


class ConflictError(GameError):
    def __init__(self, sounds):
        self.sounds = sounds
        super().__init__("These sounds were changed outside this tool (Steam update or 'verify', or "
                         "another mod tool): %s.\nRevert them or confirm to overwrite."
                         % ", ".join(s.name for s in sounds[:5]) + (" ..." if len(sounds) > 5 else ""))


def default_store_root(game):
    return os.path.join(user_dir("data"), "stores", game.store_id)


class Workspace:
    def __init__(self, game_dir, store_root=None, archives=None, full_backup=True, progress=None):
        self.game = game_dir if isinstance(game_dir, Game) else Game(game_dir)
        self.store = OriginalsStore(store_root or default_store_root(self.game), self.game.root)
        self.catalog = Catalog(self.game, archives, progress)
        self.full_backup = full_backup
        self.texts = Texts(self)

    # reading -----------------------------------------------------------------------------------
    def sound(self, name_or_key):
        return self.catalog.resolve(name_or_key)

    def wad_path(self, s):
        return self.game.wad(s.archive)

    def bank_path(self, rel):
        return banks.file_path(self.game.cooked, rel)

    def locations(self, s_or_key):
        """Bank copies written together with a sound: every copy of a BANK sound, or the prefetch heads
        of a WAD speech line (usually none)."""
        key = s_or_key if isinstance(s_or_key, str) else s_or_key.key
        archive, _, target = key.partition(":")
        if archive == banks.ARCHIVE:
            return self.catalog.locations.get(int(target, 16), [])
        return self.catalog.heads.get(key, [])

    def copies(self, s):
        """Other places of the same recording in the WAD (written and reverted together with `s`)."""
        return self.catalog.copies.get(s.key, [])

    def files(self, s):
        """Absolute paths of the game files that writing `s` changes."""
        out = []
        if s.archive != banks.ARCHIVE:
            out.append(self.wad_path(s))
            if s.is_pack_media:
                out.append(pack_bnk_path(self.wad_path(s), parse_target(s.target)[0]))
        out += sorted({self.bank_path(l.file) for x in [s] + self.copies(s) for l in self.locations(x)})
        return out

    def current_bytes(self, s):
        if s.archive == banks.ARCHIVE:
            locs = self.locations(s)
            if not locs:
                raise GameError("no copy of %s was found in the game's banks" % s.name)
            l = locs[0]
            return unfill(banks.riff_trim(banks.open_file(self.bank_path(l.file)).read(l.offset, l.size)))
        return read_target(self.wad_path(s), s.target)

    def original_bytes(self, s):
        return self.store.original(s.key) or self.current_bytes(s)

    def state(self, s, current=None):
        cur = current if current is not None else self.current_bytes(s)
        e = self.store.get(s.key)
        if e:
            h = sha256(cur)
            if e.get("applied") and h == e["applied"]:
                return MODIFIED
            if h == e["orig"]:
                return LOST if e.get("applied") else ORIGINAL
            return CHANGED
        if s.sha1 and sha1_16(cur) != s.sha1:
            return FOREIGN
        return ORIGINAL

    def modified_keys(self):
        return [k for k, e in self.store.entries.items() if e.get("applied")]

    def limit(self, s):
        """Max size in bytes for a replacement (None = any size)."""
        if s.is_pack_media:
            return None
        if s.archive == banks.ARCHIVE:
            locs = self.locations(s)
            return min(l.size for l in locs) if locs else s.size
        e = self.store.get(s.key)
        if e:
            return e["orig_size"]
        if s.sha1 and s.size:
            return s.size
        return len(self.current_bytes(s))

    def original_audio(self, s):
        return audio.decode_wem(self.original_bytes(s))

    def fill_sizes(self, s):
        """Slot sizes a replacement of `s` must fill exactly: its in-memory copies, and for a WAD recording the
        bank copies that hold it whole (wem.fill: the game plays a padded in-memory WEM as silence)."""
        if s.archive == banks.ARCHIVE:
            return sorted({l.size for l in self.locations(s) if l.kind == banks.MEMORY})
        orig = self.limit(s)
        return sorted({l.size for l in self.locations(s) if orig is not None and l.size >= orig})

    def fits(self, s, wem):
        """True if `wem` can be written for `s`: within its size limit and able to fill its in-memory slots."""
        lim = self.limit(s)
        return (lim is None or len(wem) <= lim) and all(fillable(len(wem), n) for n in self.fill_sizes(s))

    # encoding ----------------------------------------------------------------------------------
    def prepare(self, s, x, src_rate, quality="auto", progress=None):
        """Encode PCM for sound `s` (its rate/channels, fitted to its size limit) -> (wem, info). A sound with
        in-memory slots gets a WEM that fills them exactly (`fits`): at least FILL_MIN bytes smaller, and an
        even number of bytes smaller; the end is shortened by a few samples until the size allows that."""
        sizes = self.fill_sizes(s)
        lim = self.limit(s)
        if sizes:
            lim = min(([lim] if lim is not None else []) + sizes) - FILL_MIN
        rate, ch = s.rate or src_rate, s.channels or x.shape[1]
        for k in range(32):
            wem, info = encode_fit(x[:len(x) - 16 * k] if k else x, src_rate, rate, ch, lim, quality, progress)
            if self.fits(s, wem):
                return wem, info
        raise EncodeError("could not encode a size that fills this sound's slots exactly; shorten the audio")

    # writing -----------------------------------------------------------------------------------
    def _backup(self, wad, pack_hashes):
        if not self.full_backup:
            return
        self.store.full_backup(wad, self.game.rel(wad))
        for h in pack_hashes:
            b = pack_bnk_path(wad, h)
            self.store.full_backup(b, self.game.rel(b))

    def _write(self, archive, items, progress=None):
        """items: [(sound, new_bytes, meta, is_revert)] all in one archive."""
        if archive == banks.ARCHIVE:
            return self._write_banks(items, progress)
        heads = sorted({self.bank_path(l.file) for s, _, _, _ in items for l in self.locations(s)})
        for path in heads:
            self._keep_original(path)
        wad = self.game.wad(archive)
        self.game.check_writable(wad)
        packs = {parse_target(s.target)[0] for s, _, _, _ in items if s.is_pack_media}
        self._backup(wad, packs)
        w = Wad(wad)
        for h in packs:
            if self.store.pack_gaps(h) is None:
                self.store.remember_pack_gaps(h, Pack(open(pack_bnk_path(wad, h), "rb").read()).gaps(w.read(h)))
        repl, limits = {}, {}
        for s, data, meta, is_revert in items:
            if not is_revert:
                self.store.remember_original(s.key, s.archive, s.target, self.game.rel(wad),
                                             read_target(wad, s.target, w))
            repl[s.target] = data
            limits[s.target] = self.limit(s)
        self.store.save()  # originals are safe on disk before the game file changes
        if progress:
            progress("writing %s" % os.path.basename(wad))
        inject(wad, repl, limits, pack_gaps={h: self.store.pack_gaps(h) for h in packs})
        for s, data, meta, is_revert in items:
            if is_revert:
                self.store.forget(s.key)
            else:
                self.store.record_applied(s.key, data, meta)
        self.store.save()
        self.sync_files(heads, progress)

    # derived files: bank files and packages (in-memory sounds, prefetch heads, subtitles) and the
    # localization files (menu text). Each is rebuilt from its original copy + every edit applied now.
    def _derived(self):
        return self.store.data.setdefault("derived", {})

    def _fingerprint(self, path):
        """Content hash of a file (whole file: any change by another tool is detected)."""
        h = hashlib.sha1()
        with open(path, "rb") as f:
            for c in iter(lambda: f.read(1 << 22), b""):
                h.update(c)
        return h.hexdigest()

    def _keep_original(self, path):
        """Before the first change to a derived file: verified full copy + its fingerprint."""
        rel = self.game.rel(path)
        rec = self._derived().get(rel)
        if rec:
            return rec
        self.game.check_writable(path)
        self.store.full_backup(path, rel)
        rec = self._derived()[rel] = {"orig_fp": self._fingerprint(path), "written_fp": None}
        self.store.save()
        return rec

    def changed_files(self, paths):
        """Derived files changed since we last wrote them (Steam update/verify, another tool)."""
        out = []
        for path in sorted(set(paths)):
            rec = self._derived().get(self.game.rel(path))
            if rec and self._fingerprint(path) not in (rec["orig_fp"], rec["written_fp"]):
                out.append(path)
        return out

    def _audio_edits(self, rel):
        """[(offset, bytes)] that the sound replacements applied now need in bank file `rel` (cooked-relative)."""
        out = []
        for key, e in self.store.entries.items():
            if not e.get("applied"):
                continue
            locs = [l for l in self.locations(key) if l.file == rel]
            if not locs:
                continue
            data = unfill(self.store.get_blob(e["applied"]))
            for l in locs:
                whole = l.kind == banks.MEMORY or l.size >= e["orig_size"]    # a whole WEM in memory
                if whole and fillable(len(data), l.size):
                    payload = fill(data, l.size)
                elif whole:     # recorded by a version before 3.2.2 (zero padding: the game plays silence)
                    payload = data + bytes(l.size - len(data))
                else:           # prefetch head: the start of the WAD recording as it is stored there
                    full = data + bytes(max(0, e["orig_size"] - len(data)))
                    payload = (full + bytes(l.size))[:l.size]
                out.append((l.offset, payload))
        return out

    def _build(self, path, orig):
        """Content of derived file `path` from its original copy `orig`."""
        if self.texts.is_menu_file(path):
            return self.texts.menu_content(path, open(orig, "rb").read())
        rel = os.path.relpath(path, self.game.cooked).replace(os.sep, "/")
        edits = self._audio_edits(rel) + self.texts.subtitle_edits(rel)
        return banks.open_file(orig).patched(edits) if edits else open(orig, "rb").read()

    def sync_files(self, paths, progress=None):
        """Rebuild each derived file from its original copy and the edits applied now (atomic writes)."""
        for path in sorted(set(paths)):
            rec = self._keep_original(path)
            orig = self.store.full_copy(self.game.rel(path))
            if orig is None:
                raise GameError("the original copy of %s is missing from the store" % self.game.rel(path))
            if progress:
                progress("writing %s" % os.path.basename(path))
            data = self._build(path, orig)
            self.game.check_writable(path)
            tmp = path + ".ackit-tmp"
            with open(tmp, "wb") as f:
                f.write(data)
            os.replace(tmp, path)
            same = data == open(orig, "rb").read()
            rec["written_fp"] = None if same else self._fingerprint(path)
        self.store.save()

    def _write_banks(self, items, progress=None):
        for s, data, meta, is_revert in items:
            locs = self.locations(s)
            if not locs:
                raise GameError("no copy of %s was found in the game's banks" % s.name)
            if not is_revert and len(data) > min(l.size for l in locs):
                raise SizeLimitError(s.target, min(l.size for l in locs), len(data))
        files = sorted({self.bank_path(l.file) for s, _, _, _ in items for l in self.locations(s)})
        for path in files:
            self._keep_original(path)
        for s, data, meta, is_revert in items:
            if is_revert:
                self.store.forget(s.key)
            else:
                self.store.remember_original(s.key, s.archive, s.target, self.locations(s)[0].file,
                                             self.current_bytes(s))
                self.store.record_applied(s.key, data, meta)
        self.store.save()  # the store says what the files must contain before they change
        self.sync_files(files, progress)

    def apply(self, changes, force=False, progress=None):
        """changes: [(sound, wem_bytes, meta dict)]. One atomic rewrite per WAD / bank file. A recording stored
        in two places of a WAD gets the same bytes in both (Workspace.copies)."""
        keys = {s.key for s, _, _ in changes}
        changes = [(s, unfill(data), meta) for s, data, meta in changes]
        for s, data, _ in changes:
            if not self.fits(s, data) and len(data) <= (self.limit(s) or len(data)):
                raise GameError("%s: a %d-byte WEM cannot fill this sound's in-memory slots (%s bytes) exactly; "
                                "encode it for this sound (Workspace.prepare / the editor / ackit replace)"
                                % (s.name, len(data), ", ".join(map(str, self.fill_sizes(s)))))
        changes = list(changes) + [(c, data, meta) for s, data, meta in changes for c in self.copies(s)
                                   if c.key not in keys]
        conflicts = []
        for s, _, _ in changes:
            if self.state(s) in (CHANGED, FOREIGN):
                conflicts.append(s)
        changed = set(self.changed_files([self.bank_path(l.file) for s, _, _ in changes for l in self.locations(s)]))
        conflicts += [s for s, _, _ in changes
                      if s not in conflicts and {self.bank_path(l.file) for l in self.locations(s)} & changed]
        if conflicts and not force:
            raise ConflictError(conflicts)
        by_arch = defaultdict(list)
        for s, data, meta in changes:
            by_arch[s.archive].append((s, data, meta, False))
        for a, items in by_arch.items():
            self._write(a, items, progress)

    def revert(self, sounds, progress=None):
        """Put the stored originals back (sounds without a record are skipped). Returns reverted."""
        by_arch = defaultdict(list)
        done = []
        keys = {s.key for s in sounds}
        sounds = list(sounds) + [c for s in sounds for c in self.copies(s) if c.key not in keys]
        for s in sounds:
            orig = self.store.original(s.key)
            if orig is None:
                continue
            by_arch[s.archive].append((s, orig, None, True))
            done.append(s)
        for a, items in by_arch.items():
            self._write(a, items, progress)
        self.store.gc()
        return done

    def revert_all(self, progress=None):
        sounds = [self.catalog.by_key[k] for k in list(self.store.entries) if k in self.catalog.by_key]
        missing = [k for k in self.store.entries if k not in self.catalog.by_key]
        if missing:
            # records of archives not loaded in this catalog (e.g. another language)
            other = Catalog(self.game, sorted({k.split(":")[0] for k in missing}))
            sounds += [other.by_key[k] for k in missing if k in other.by_key]
        return self.revert(sounds, progress)

    def restore_from_full_backup(self):
        """Last-resort fallback: copy back every full backup taken by this store."""
        out = []
        for rel in self.store.data.get("full", {}):
            src = self.store.full_copy(rel)
            dst = os.path.join(self.game.root, rel)
            self.game.check_writable(dst)
            tmp = dst + ".ackit-tmp"
            import shutil
            shutil.copy2(src, tmp)
            os.replace(tmp, dst)
            out.append(rel)
        for rec in self._derived().values():
            rec["written_fp"] = None
        self.store.data["texts"] = {}
        self.store.data["entries"] = {}
        self.store.save()
        self.store.gc()
        return out