UnheardCity/ackit/modpack.py

"""Portable mod packages (*.acmod): share replaced sounds between machines and languages.

A package is a zip:
  manifest.json      {"format": 1, "name", "author", "description", "created", "items": [...]}
  items/NNN.wem      the exact WEM that was applied (used as-is when the target matches)
  items/NNN.flac     the same audio, decoded (used to re-encode when the target differs, e.g. another
                     language whose original has a different size)
Each item records archive, target, name, group, the SHA-1 of the original it replaced and, for a voice line,
its line identity (the same in every language, ackit/lines.py): that is how another language finds it.
  "texts": [{"key", "text", "original"}]   edited subtitles / menu texts (keys: ackit/texts.py); optional,
                                           so packages without texts stay valid
"""
import io
import json
import os
import time
import zipfile
from dataclasses import dataclass

import soundfile as sf

from . import audio
from .originals import sha1_16

FORMAT = 1


def export_mod(ws, path, keys, name="", author="", description="", texts=()):
    """Write the applied versions of sound `keys` and text keys `texts` (from the originals store) into a
    package. Returns the number of entries."""
    items = []
    recs = ws.texts.records()
    text_items = []
    for k in texts:
        r = recs.get(k)
        if not r or r.get("applied") is None:
            raise ValueError("%s is not modified" % k)
        text_items.append({"key": k, "text": r["applied"], "original": r["orig"]})
    with zipfile.ZipFile(path + ".tmp", "w", zipfile.ZIP_DEFLATED) as z:
        for i, key in enumerate(keys):
            e = ws.store.get(key)
            if not e or not e.get("applied"):
                raise ValueError("%s is not modified" % key)
            s = ws.sound_for_key(key)               # also a change made in another language (its name and line)
            wem = ws.store.get_blob(e["applied"])
            x, sr = audio.decode_wem(wem)
            buf = io.BytesIO()
            sf.write(buf, x, sr, format="FLAC")
            base = "items/%03d" % i
            z.writestr(base + ".wem", wem)
            z.writestr(base + ".flac", buf.getvalue())
            items.append({"archive": e["archive"], "target": e["target"], "name": s.name if s else "",
                          "line": s.line if s else "",
                          "group": s.group if s else "", "original_sha1": e["orig_sha1"],
                          "wem": base + ".wem", "audio": base + ".flac", "meta": e.get("meta", {})})
        z.writestr("manifest.json", json.dumps({"format": FORMAT, "name": name, "author": author,
                                                "description": description, "created": time.time(),
                                                "game": "Batman: Arkham City (PC)", "items": items,
                                                "texts": text_items}, indent=1, ensure_ascii=False))
    os.replace(path + ".tmp", path)
    return len(items) + len(text_items)


@dataclass
class ImportItem:
    item: dict
    sound: object = None      # matched Sound or None
    how: str = ""             # 'exact' | 'same line' | 'by name' | 'missing'
    wem: bytes = None         # ready-to-apply bytes (after `prepare_import`)
    note: str = ""


def read_mod(path):
    z = zipfile.ZipFile(path)
    man = json.loads(z.read("manifest.json"))
    if man.get("format") != FORMAT:
        raise ValueError("unsupported mod package format %r" % man.get("format"))
    return z, man


def match_mod(ws, path):
    """Match package items to sounds of this workspace. -> (manifest, [ImportItem])"""
    z, man = read_mod(path)
    out = []
    by_line = {}
    for c in ws.catalog.sounds:
        if c.line:
            by_line.setdefault(c.line, []).append(c)
    for it in man["items"]:
        key = "%s:%s" % (it["archive"], it["target"])
        s = ws.catalog.by_key.get(key)
        # the very same recording: same position and the same name or the same original bytes
        if s and (s.name == it.get("name") or (s.sha1 and s.sha1 == it.get("original_sha1"))):
            out.append(ImportItem(it, s, "exact"))
            continue
        if it.get("line"):                    # a voice line: its recording in the language loaded here
            cands = sorted(by_line.get(it["line"], []), key=lambda c: (c.name != it.get("name"), c.key))
            if cands:                         # several only for one recording stored twice (written together)
                out.append(ImportItem(it, cands[0], "same line"))
            else:
                out.append(ImportItem(it, None, "missing", note="the line %r has no recording in the language "
                                      "loaded here" % it.get("name")))
            continue
        try:
            cands = [c for c in ws.catalog.search(it["name"]) if c.name == it["name"]] if it.get("name") else []
        except LookupError:
            cands = []
        if len(cands) == 1 and cands[0].line:
            # a voice line from a package without line identities (older version): names alone are not proof
            out.append(ImportItem(it, None, "missing", note="%r is a voice line from a package of an older "
                                  "version; export the mod again to move it to another language" % it["name"]))
        elif len(cands) == 1:
            out.append(ImportItem(it, cands[0], "by name"))
        else:
            out.append(ImportItem(it, None, "missing", note="no sound named %r in this game/language" % it.get("name")))
    z.close()
    return man, out


def match_texts(ws, man):
    """[(key, text, problem or '')] for the package's texts in this game (problem: why it cannot be applied)."""
    out = []
    for it in man.get("texts", []):
        k, text = it["key"], it["text"]
        try:
            ws.texts.current(k)
            ws.texts.check(k, text)
            out.append((k, text, ""))
        except Exception as e:  # missing key, too long, characters the file cannot hold
            out.append((k, text, str(e)))
    return out


def prepare_import(ws, path, items, progress=None):
    """Fill `wem` for matched items: reuse the packaged WEM when it targets the very same original,
    otherwise re-encode the packaged audio for the local sound."""
    z, _ = read_mod(path)
    for im in items:
        if not im.sound:
            continue
        orig = ws.original_bytes(im.sound)
        wem = z.read(im.item["wem"])
        if im.how == "exact" and sha1_16(orig) == im.item.get("original_sha1") and ws.fits(im.sound, wem):
            im.wem, im.note = wem, "exact copy"
        else:
            x, sr = sf.read(io.BytesIO(z.read(im.item["audio"])), dtype="float32", always_2d=True)
            im.wem, info = ws.prepare(im.sound, x, sr, progress=progress)
            im.note = "re-encoded (%d Hz, q%.1f)" % (info["rate"], info["quality"])
    z.close()
    return items


def import_mod(ws, path, force=False, progress=None):
    """Match, prepare and apply a package. Returns ([ImportItem] (see .how/.note), [(key, text, problem)])."""
    man, items = match_mod(ws, path)
    prepare_import(ws, path, items, progress)
    source = {"source": "mod:" + path.replace("\\", "/").split("/")[-1]}
    changes = [(im.sound, im.wem, source) for im in items if im.wem]
    if changes:
        ws.apply(changes, force=force, progress=progress)
    texts = match_texts(ws, man)
    ok = {k: v for k, v, problem in texts if not problem}
    if ok:
        ws.texts.apply(ok, force=force, progress=progress, meta=source)
    return items, texts