"""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 (Latin-1/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

from .upk import Package

CLASS = "RDialogueEvent"
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")
Dialogue = namedtuple("Dialogue", "speech export fields")   # speech: name of the referenced speech event


class SubtitleError(ValueError):
    pass


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 raw.split(b"\0", 1)[0].decode("cp1252", "replace"), 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 = text.encode("cp1252")
        except UnicodeEncodeError:
            raise SubtitleError("this subtitle is stored as Latin-1 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:
        c.encode("cp1252")
        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 RDialogueEvent export (bytes `b` starting at uncompressed offset `base`)."""
    ls, sub = names_idx
    i = b.find(struct.pack("<II", ls, 0))
    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 _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.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)))


# --- 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"]
SubLine = namedtuple("SubLine", "lang speech line text room enc sound")      # sound: linked WAD sound key
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 scan_package(pkg, data=None, start=None, tables=None):
    """[Dialogue] of one package (empty if it has none)."""
    if data is None:
        start, data = pkg.read_all()
    names, imports, exports = tables or pkg.tables(data, start)
    cls = {-(i + 1) for i, (c, o, _) in enumerate(imports) if c == "Class" and o == CLASS}
    if not cls or "LocalizedSubtitles" not in names or "Subtitles" not in names:
        return []
    idx = (names.index("LocalizedSubtitles"), names.index("Subtitles"))
    out = []
    for e in exports:
        if e.cls in cls and e.size > 0:
            b = data[e.offset - start:e.offset - start + e.size]
            fields = parse_dialogue(b, e.offset, idx)
            if fields:
                out.append(Dialogue(_speech_ref(b, exports, imports), e.name, fields))
    return out