UnheardCity/ackit/subtitles.py

"""Subtitles: the RDialogueEvent objects inside the game's packages (docs/formats/subtitles.md).

Each dialogue object references its speech event (an export named like the speech line, e.g.
'Speech_CHURCH_B1_Ch2_POST_Explosion_BM_Cash_04') and holds the subtitle text of every language:
  LocalizedSubtitles: array of {Language: FString, Subtitles: array of {Text: FString, (float @0x0c)}}
in Rocksteady's compact tagged-property format (u16 type, u16 field offset, value; structs and arrays with
a byte size; struct elements end with a u16 0).

A dialogue object can be stored in several packages; a subtitle edit is written into every copy, in place:
the new text replaces the bytes of the old FString and is padded with NUL characters (the game stops at the
first NUL), so no size field, export or package layout changes. The text must therefore fit the original
string's bytes (cp1252 strings: one byte per character; UTF-16 strings: two).
"""
import csv
import gzip
import io
import os
import re
import struct
from collections import defaultdict, namedtuple


CLASS = "RDialogueEvent"
SUBTITLE_CLASS, LOOKUP_CLASS = "RSubtitle", "RSubtitleLookup"   # scripted radio/encryption subtitles
UNRESOLVED = "#lookup:%08X"     # speech of an RSubtitle until its lookup key is matched to a line name
T_STRUCT, T_STRING, T_OBJECT = 0x09, 0x0D, 0x10

# one subtitle string inside a package: pos = absolute uncompressed offset of the FString length field,
# size = bytes of the whole FString (length field included), enc 'a' (cp1252) or 'w' (UTF-16)
Field = namedtuple("Field", "lang index text pos size enc")
# speech: name of the referenced speech event; speaker: the object's speaker field (RDialogueEvent only)
Dialogue = namedtuple("Dialogue", "speech export fields speaker", defaults=("",))


class SubtitleError(ValueError):
    pass


# Windows-1252 as Windows itself maps it: the 5 bytes Python's cp1252 leaves undefined stand for the same code
# points (U+0081 ...). One line is UTF-8 bytes in a cp1252 FString in 4 languages (Speech_TAPE_J_Answer_1_00,
# docs/formats/subtitles.md); the RUS and KOR ones contain such bytes. Read and written back losslessly, as stored.
_CP1252_UNDEFINED = frozenset((0x81, 0x8D, 0x8F, 0x90, 0x9D))


def _decode_cp1252(raw):
    try:
        return raw.decode("cp1252")
    except UnicodeDecodeError:
        return "".join(chr(x) if x in _CP1252_UNDEFINED else bytes([x]).decode("cp1252") for x in raw)


def _encode_cp1252(text):
    """UnicodeEncodeError for a character that is not in Windows-1252."""
    try:
        return text.encode("cp1252")
    except UnicodeEncodeError:
        return b"".join(bytes([ord(c)]) if ord(c) in _CP1252_UNDEFINED else c.encode("cp1252") for c in text)


def read_fstring(b, o):
    """-> (text, end offset, enc)"""
    n = struct.unpack_from("<i", b, o)[0]
    if n >= 0:
        raw = bytes(b[o + 4:o + 4 + n])
        return _decode_cp1252(raw.split(b"\0", 1)[0]), o + 4 + n, "a"
    raw = bytes(b[o + 4:o + 4 - 2 * n])
    return raw.decode("utf-16-le", "replace").split("\0", 1)[0], o + 4 - 2 * n, "w"


def encode_fstring(text, size, enc):
    """The FString bytes (length field included) for `text`, exactly `size` bytes long (NUL padded).
    Raises SubtitleError when the text does not fit."""
    if "\0" in text:
        raise SubtitleError("the text must not contain NUL characters")
    if enc == "a":
        try:
            raw = _encode_cp1252(text)
        except UnicodeEncodeError:
            raise SubtitleError("this subtitle is stored as Western European (Windows-1252) text; '%s' cannot be "
                                "written"
                                % next(c for c in text if not _cp1252(c))) from None
        n = size - 4
        if len(raw) + 1 > n:
            raise SubtitleError("the text is %d characters long but this subtitle has room for %d"
                                % (len(raw), n - 1))
        return struct.pack("<i", n) + raw + bytes(n - len(raw))
    raw = text.encode("utf-16-le")
    n = (size - 4) // 2
    if len(raw) // 2 + 1 > n:
        raise SubtitleError("the text is %d characters long but this subtitle has room for %d" % (len(raw) // 2, n - 1))
    return struct.pack("<i", -n) + raw + bytes(2 * n - len(raw))


def _cp1252(c):
    try:
        _encode_cp1252(c)
        return True
    except UnicodeEncodeError:
        return False


def room(field):
    """Maximum number of characters for a subtitle string."""
    return (field.size - 4 - 1) if field.enc == "a" else (field.size - 4) // 2 - 1


def _parse_lines(b, o, end, base, lang):
    """Subtitles array content [o, end) -> [Field]."""
    count = struct.unpack_from("<I", b, o)[0]
    o += 4
    out = []
    for i in range(count):
        # element: props until u16 0; the text is the string at field offset 0
        while True:
            t, fo = struct.unpack_from("<HH", b, o)
            if t == 0:
                o += 2
                break
            if t == T_STRING:
                text, e, enc = read_fstring(b, o + 4)
                if fo == 0:
                    out.append(Field(lang, i, text, base + o + 4, e - (o + 4), enc))
                o = e
            elif t == 0x04:               # 4-byte value (the line's time)
                o += 8
            else:
                raise SubtitleError("unexpected property type %#x in a subtitle line" % t)
            if o > end:
                raise SubtitleError("subtitle line overruns its array")
    if o != end:
        raise SubtitleError("subtitle array size mismatch")
    return out


def parse_dialogue(b, base, names_idx):
    """Fields of one dialogue/subtitle export (bytes `b` starting at uncompressed offset `base`).
    names_idx: (name indices of the per-language array: LocalizedSubtitles or ArraySubtitles, of Subtitles)."""
    arrays, sub = names_idx
    i = -1
    for ls in arrays:
        i = b.find(struct.pack("<II", ls, 0))
        if i >= 4 and struct.unpack_from("<H", b, i - 4)[0] == T_STRUCT:
            break
    if i < 4 or struct.unpack_from("<H", b, i - 4)[0] != T_STRUCT:
        return []
    size = struct.unpack_from("<i", b, i + 8)[0]
    o, end = i + 16, i + 16 + size
    count = struct.unpack_from("<I", b, o)[0]
    o += 4
    fields = []
    for _ in range(count):
        lang = None
        while True:
            t, fo = struct.unpack_from("<HH", b, o)
            if t == 0:
                o += 2
                break
            if t == T_STRING:
                lang, o, _ = read_fstring(b, o + 4)
            elif t == T_STRUCT:
                nm, _nn, ssize, _ai = struct.unpack_from("<iiii", b, o + 4)
                c0 = o + 20
                if nm == sub:
                    fields += _parse_lines(b, c0, c0 + ssize, base, lang)
                o = c0 + ssize
            else:
                raise SubtitleError("unexpected property type %#x in a localized subtitle" % t)
    if o != end:
        raise SubtitleError("localized subtitle size mismatch")
    return fields


def _speaker(b):
    """The speaker FString of a dialogue object (property 0x0D at field offset 0x34), '' if none."""
    m = re.search(rb"\x0d\x00\x34\x00", b)
    try:
        return read_fstring(b, m.start() + 4)[0] if m else ""
    except (struct.error, UnicodeDecodeError, ValueError):
        return ""


def _speech_ref(b, exports, imports):
    """Name of the object referenced by the dialogue's first object property (its speech event)."""
    m = re.search(rb"\x10\x00\x2c\x00", b)
    if not m:
        return ""
    ref = struct.unpack_from("<i", b, m.start() + 4)[0]
    if ref > 0 and ref <= len(exports):
        return exports[ref - 1].name
    if ref < 0 and -ref <= len(imports):
        return imports[-ref - 1][1]
    return ""


class Collector:
    """`visit` hook for banks.scan: collects the dialogue objects of every package."""

    def __init__(self):
        self.dialogues = defaultdict(list)       # speech name -> [(file, Dialogue)]
        self.scripted = set()                    # speech names of RSubtitle objects (resolve)
        self.issues = []

    def __call__(self, rel, pkg, start, data):
        try:
            for d in scan_package(pkg, data, start):
                self.dialogues[d.speech].append((rel, d))
        except (SubtitleError, struct.error, ValueError, IndexError) as e:
            self.issues.append(("subtitles", "%s: %s" % (rel, e)))

    PREFIXES = ("", "RADIO_SCAN_", "RADIO_", "TV_", "CH_ZERO_", "GCPD_", "CHAR_", "SEQUENCE_")

    def resolve(self, line_names):
        """Give RSubtitle objects their speech line name: the name whose lookup_key equals the object's key,
        from `line_names` or built from the object's own name (Speech_ + prefix + name; exact hash match)."""
        by_key = {lookup_key(n): n for n in line_names}
        for sp in [s for s in self.dialogues if s.startswith("#lookup:")]:
            key = int(sp.split(":")[1], 16)
            name = by_key.get(key)
            if not name:
                exp = self.dialogues[sp][0][1].export
                name = next(("Speech_" + p + exp for p in self.PREFIXES if lookup_key("Speech_" + p + exp) == key), None)
            copies = self.dialogues.pop(sp)
            if not name:
                self.issues.append(("subtitles", "no speech line for subtitle %s" % copies[0][1].export))
                name = copies[0][1].export
            self.dialogues[name].extend((rel, d._replace(speech=name)) for rel, d in copies)
            self.scripted.add(name)


# --- catalog files: SUB_<LANG>.csv (texts) + SUB_locations.csv.gz (every copy) ------------------
LOCATIONS = "SUB_locations.csv.gz"
TEXT_FIELDS = ["speech", "line", "text", "room", "enc", "sound"]
class SubLine(namedtuple("SubLine", "lang speech line text room enc sound")):
    """A subtitle string. sound: the voice recording(s) its speech event plays in that language (sound keys,
    ';'-separated in the event's order: an event can play one of two versions of a line)."""
    __slots__ = ()

    @property
    def sounds(self):
        return tuple(self.sound.split(";")) if self.sound else ()
SubLoc = namedtuple("SubLoc", "lang speech line file pos size enc")


def key(lang, speech, line):
    return "SUB:%s:%s:%d" % (lang, speech, line)


def rows(collector, links):
    """({lang: [SubLine]}, [SubLoc]) from a Collector; links: {speech: {lang: sound key}}."""
    texts = defaultdict(list)
    locs = []
    for speech in sorted(collector.dialogues):
        copies = collector.dialogues[speech]
        rel0, d0 = copies[0]
        for f in d0.fields:
            ln = links.get(speech, {})
            snd = ln.get(f.lang) or ln.get("INT", "")
            texts[f.lang].append(SubLine(f.lang, speech, f.index, f.text, room(f), f.enc, snd))
        for rel, d in copies:
            for f in d.fields:
                locs.append(SubLoc(f.lang, speech, f.index, rel, f.pos, f.size, f.enc))
    return dict(texts), locs


def write_texts(path, lines):
    with open(path + ".tmp", "w", encoding="utf-8", newline="") as fh:
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(TEXT_FIELDS)
        for s in lines:
            w.writerow([s.speech, s.line, s.text, s.room, s.enc, s.sound])
    os.replace(path + ".tmp", path)


def read_texts(path, lang):
    with open(path, encoding="utf-8", newline="") as fh:
        return [SubLine(lang, r["speech"], int(r["line"]), r["text"], int(r["room"]), r["enc"], r["sound"])
                for r in csv.DictReader(fh)]


def write_locations(path, locs):
    raw = io.StringIO()
    w = csv.writer(raw, lineterminator="\n")
    w.writerow(["lang", "speech", "line", "file", "pos", "size", "enc"])
    for l in sorted(locs, key=lambda l: (l.lang, l.speech, l.line, l.file, l.pos)):
        w.writerow([l.lang, l.speech, l.line, l.file, "%X" % l.pos, l.size, l.enc])
    with open(path + ".tmp", "wb") as fh:
        with gzip.GzipFile(fileobj=fh, mode="wb", mtime=0) as gz:
            gz.write(raw.getvalue().encode("utf-8"))
    os.replace(path + ".tmp", path)


def read_locations(path, lang=None):
    """{SUB key: [SubLoc]} (only `lang` if given)."""
    out = defaultdict(list)
    with gzip.open(path, "rt", encoding="utf-8", newline="") as fh:
        for r in csv.DictReader(fh):
            if lang and r["lang"] != lang:
                continue
            l = SubLoc(r["lang"], r["speech"], int(r["line"]), r["file"], int(r["pos"], 16), int(r["size"]), r["enc"])
            out[key(l.lang, l.speech, l.line)].append(l)
    return dict(out)


def lookup_key(line_name):
    """Key of a speech line in the RSubtitleLookup table: FNV-1 of its name without 'Speech_'."""
    from .bnk import fnv1
    return fnv1(line_name[7:] if line_name.lower().startswith("speech_") else line_name)


def _lookup_table(b):
    """RSubtitleLookup export data -> {export index: key}: array of {int key @0, object @4}."""
    count = struct.unpack_from("<I", b, 24)[0]
    out, o = {}, 28
    for _ in range(count):
        if b[o:o + 4] != b"\x02\x00\x00\x00" or b[o + 8:o + 12] != b"\x10\x00\x04\x00":
            raise SubtitleError("unexpected subtitle lookup layout")
        key, ref = struct.unpack_from("<Ii", b, o + 4)[0], struct.unpack_from("<i", b, o + 12)[0]
        out[ref] = key
        o += 18
    return out


def scan_package(pkg, data=None, start=None, tables=None):
    """[Dialogue] of one package (empty if it has none): RDialogueEvent objects (speech = the referenced
    speech event) and RSubtitle objects (speech = UNRESOLVED % lookup key, see Collector.resolve)."""
    if data is None:
        start, data = pkg.read_all()
    names, imports, exports = tables or pkg.tables(data, start)
    by_class = {}
    for i, (c, o, _) in enumerate(imports):
        if c == "Class" and o in (CLASS, SUBTITLE_CLASS, LOOKUP_CLASS):
            by_class[-(i + 1)] = o
    arrays = [names.index(a) for a in ("LocalizedSubtitles", "ArraySubtitles") if a in names]
    if not by_class or not arrays or "Subtitles" not in names:
        return []
    idx = (arrays, names.index("Subtitles"))
    keys = {}
    for e in exports:
        if by_class.get(e.cls) == LOOKUP_CLASS and e.size > 0:
            keys.update(_lookup_table(data[e.offset - start:e.offset - start + e.size]))
    out = []
    for e in exports:
        kind = by_class.get(e.cls)
        if kind in (CLASS, SUBTITLE_CLASS) and e.size > 0:
            b = data[e.offset - start:e.offset - start + e.size]
            fields = parse_dialogue(b, e.offset, idx)
            if not fields:
                continue
            if kind == CLASS:
                speech = _speech_ref(b, exports, imports)
            elif e.index in keys:
                speech = UNRESOLVED % keys[e.index]
            else:
                speech = e.name                       # not in a lookup table: keep the object's name
            out.append(Dialogue(speech, e.name, fields, _speaker(b) if kind == CLASS else ""))
    return out