"""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.
"""
import os
from collections import defaultdict

from . import audio
from .catalog import Catalog
from .encoder import encode_fit
from .game import Game, GameError
from .inject import inject, read_target, parse_target, pack_bnk_path
from .originals import OriginalsStore, sha256, sha1_16
from .pack import Pack
from .paths import user_dir
from .wad import Wad

# 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

    # 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 current_bytes(self, s):
        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
        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 current_audio(self, s):
        return audio.decode_wem(self.current_bytes(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)."""
        return encode_fit(x, src_rate, s.rate or src_rate, s.channels or x.shape[1], self.limit(s),
                          quality, progress)

    def prepare_file(self, s, path, quality="auto", progress=None):
        x, sr = audio.load(path)
        return self.prepare(s, x, sr, quality, progress)

    # 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."""
        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()

    def apply(self, changes, force=False, progress=None):
        """changes: [(sound, wem_bytes, meta dict)]. One atomic rewrite per WAD."""
        conflicts = []
        for s, _, _ in changes:
            if self.state(s) in (CHANGED, FOREIGN):
                conflicts.append(s)
        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 = []
        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)
        self.store.data["entries"] = {}
        self.store.save()
        self.store.gc()
        return out