"""Read and replace WEMs inside WAD archives.

Target syntax (stable ids used everywhere, also in mod packages):
  'HASH'        WAD entry holding one WEM                        e.g. 800BE4F0
  'HASH:i'      i-th WEM (0-based) of an entry holding several    e.g. B49CDF48:11
  'PACK/MEDIA'  streamed WEM inside a BKPK stream pack (SFX)      e.g. EDF036D6/393BC53A

Rules (see docs/formats/wad.md):
  * entries / variations are indexed from inside the game's compressed .upk packages, so they are
    replaced IN PLACE and must not exceed the original WEM size (`limit`); the WAD layout never moves;
  * stream-pack media may change size: the pack index (SFX\\<PACK>.bnk) is rewritten too.
Writes go to a temporary file which atomically replaces the original (never half-written).
"""
import os

from .pack import Pack
from .wad import Wad, write_wad
from .wem import Wem, split_concatenated


class SizeLimitError(ValueError):
    def __init__(self, target, limit, size):
        super().__init__("the new sound is %d bytes but the original is %d bytes; the game reads exactly the "
                         "original size, so the replacement must not be larger (shorten the audio or lower "
                         "the quality)" % (size, limit))
        self.target, self.limit, self.size = target, limit, size


def parse_target(t: str):
    """-> (entry_hash, index, media_id or None)"""
    if "/" in t:
        p, m = t.split("/", 1)
        return int(p, 16), 0, int(m, 16)
    h, _, idx = t.partition(":")
    return int(h, 16), int(idx or 0), None


def pack_bnk_path(wad_path, pack_id):
    """SFX\\<PACK>.bnk next to the WAD (case-insensitive), or None."""
    d = os.path.dirname(wad_path)
    want = ("%08X.bnk" % pack_id).lower()
    hit = next((f for f in os.listdir(d) if f.lower() == want), None)
    return os.path.join(d, hit) if hit else None


def is_pack(wad_path, entry_hash):
    p = pack_bnk_path(wad_path, entry_hash)
    if not p:
        return False
    with open(p, "rb") as f:
        return f.read(4) == b"BKPK"


def read_target(wad_path: str, target: str, wad: Wad = None) -> bytes:
    """Bytes of one WEM (seeks directly; does not load whole packs)."""
    wad = wad or Wad(wad_path)
    h, idx, mid = parse_target(target)
    if h not in wad.index:
        raise KeyError("entry %08X not in %s" % (h, os.path.basename(wad_path)))
    e = wad.index[h]
    if mid is not None:
        pk = Pack(open(pack_bnk_path(wad_path, h), "rb").read())
        _, size, off = pk.entries[pk.index[mid]]
        with open(wad_path, "rb") as f:
            f.seek(e.offset + off)
            return f.read(size)
    riffs, _ = split_concatenated(wad.read(h))
    if idx >= len(riffs):
        raise IndexError("entry %08X holds %d WEM(s)" % (h, len(riffs)))
    return riffs[idx][1]


def replace_entry_payload(payload: bytes, index: int, new_wem: bytes, limit: int, target="") -> bytes:
    """Replace WEM `index` in place, zero-filled up to `limit` bytes; nothing else moves."""
    riffs, _ = split_concatenated(payload)
    if index >= len(riffs):
        raise IndexError("entry holds %d WEM(s); index %d invalid" % (len(riffs), index))
    if len(new_wem) > limit:
        raise SizeLimitError(target, limit, len(new_wem))
    off = riffs[index][0]
    if off + limit > len(payload):
        raise ValueError("entry is shorter than the original size limit (unexpected WAD state)")
    return payload[:off] + new_wem + bytes(limit - len(new_wem)) + payload[off + limit:]


def inject(wad_path: str, replacements: dict, limits: dict, out_path: str = None, pack_gaps: dict = None):
    """Write several WEMs into one WAD in a single atomic rewrite.

    replacements: {target: wem_bytes}; limits: {target: original size} for non-pack targets.
    out_path: write a copy elsewhere instead of replacing the WAD (tests/dry runs).
    pack_gaps: {pack hash: Pack.gaps()} original padding, so reverts restore packs byte-for-byte.
    Returns the list of files written."""
    src = Wad(wad_path)
    payloads, pack_repl = {}, {}
    for target, data in replacements.items():
        Wem.parse(data)  # must be a game-format WEM
        h, idx, mid = parse_target(target)
        if h not in src.index:
            raise KeyError("entry %08X not in %s" % (h, os.path.basename(wad_path)))
        if mid is not None:
            pack_repl.setdefault(h, {})[mid] = data
            continue
        if limits.get(target) is None:
            raise ValueError("no size limit given for %s" % target)
        base = payloads.get(h) or src.read(h)
        payloads[h] = replace_entry_payload(base, idx, data, limits[target], target)
    new_bnks = {}
    for h, media in pack_repl.items():
        bpath = pack_bnk_path(wad_path, h)
        new_bnks[bpath], payloads[h] = Pack(open(bpath, "rb").read()).rebuild(
            src.read(h), media, (pack_gaps or {}).get(h))

    dst = out_path or wad_path
    bnk_dsts = {b: (b if out_path is None else os.path.join(os.path.dirname(out_path), os.path.basename(b)))
                for b in new_bnks}
    tmps = [dst + ".ackit-tmp"] + [d + ".ackit-tmp" for d in bnk_dsts.values()]
    try:
        write_wad(src, tmps[0], payloads)
        problems = Wad(tmps[0]).check()
        if problems:
            raise RuntimeError("rebuilt WAD failed validation: %s" % problems[:3])
        for b, d in bnk_dsts.items():
            with open(d + ".ackit-tmp", "wb") as f:
                f.write(new_bnks[b])
        # commit: pack indexes first (tiny), then the WAD; each os.replace is atomic
        for d in bnk_dsts.values():
            os.replace(d + ".ackit-tmp", d)
        os.replace(tmps[0], dst)
    finally:
        for p in tmps:
            if os.path.exists(p):
                os.remove(p)
    return [dst] + list(bnk_dsts.values())