"""Real-game write round trips for the mission-3 features (no game launch; reverts everything itself).

    python tools/e2e_features.py [--game DIR] [--lang ITA,INT,... | --lang all]

Runs the checks for each language (default: the game's language; `all` = every subtitle/menu language). A
language without its own voices (ESM, POL, POR, RUS, JPN, KOR) uses the INT voices, like the game and the editor.
For every check: the SHA-1 of each touched file is taken, the change is applied through Workspace (the code path
of the editor and the CLI), read back from the game files, reverted, and the files must be byte-identical again.
  1. in-memory sound with the most copies (every copy written; once, it is shared by all languages)
  2. localized in-memory speech line of the language's voices
  3. speech line of the language's voices with prefetch heads (WAD + every head copy), and one that a bank
     also stores whole (WAD + every whole copy)
  4. subtitle with several copies (RDialogueEvent), and a scripted subtitle (RSubtitle, Startup.upk)
  5. menu text
  6. sound + subtitle in the same package, reverted one at a time
  7. mod package made with the language's voices and imported into another voice language (INT; FRA for INT):
     matched by the game's line identity, so the same line (not the same WAD position) is replaced; a text key
     names its language, so the edited subtitle stays that language's subtitle (the other language's is untouched)
  8. 5-channel speech line replaced with effects (the Edit tab's chain: EQ, echo, PA channel layout): the
     decoded game file has the line's channels in the game's order, the surrounds carry the delayed voice
Uses temporary originals stores (the user's store and settings are not touched), deleted at the end when every
change was reverted (kept, with a warning, otherwise). Whatever happens (a failed check, an exception), every
change is reverted before the tool exits; a sound that is not in its original state is never written.
"""
import argparse
import hashlib
import os
import shutil
import sys
import tempfile

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

import numpy as np  # noqa: E402

from ackit import audio, banks, fx, settings, subtitles  # noqa: E402
from ackit.game import Game, default_game_dir  # noqa: E402
from ackit.modpack import export_mod, import_mod, match_mod  # noqa: E402
from ackit.wem import fill  # noqa: E402
from ackit.workspace import Workspace, MODIFIED, ORIGINAL  # noqa: E402

RESULTS = []
SPACES = []          # every Workspace created: reverted in main()'s finally block


def sha(p):
    h = hashlib.sha1()
    with open(p, "rb") as f:
        for c in iter(lambda: f.read(1 << 22), b""):
            h.update(c)
    return h.hexdigest()


def check(name, ok, detail=""):
    RESULTS.append((name, bool(ok)))
    print("%-62s %s %s" % (name, "PASS" if ok else "FAIL", detail))


def tone(seconds, freq):
    t = np.arange(int(seconds * 48000)) / 48000.0
    return (0.3 * np.sin(2 * np.pi * freq * t)).astype(np.float32)[:, None], 48000


def snapshot(paths):
    return {p: sha(p) for p in sorted(set(paths))}


def copies_of(ws, s):
    out = []
    for l in ws.locations(s):
        b = banks.open_file(ws.bank_path(l.file)).read(l.offset, l.size)
        out.append((l, b))
    return out


def sound_round_trip(ws, s, label):
    if ws.state(s) != ORIGINAL:
        check(label, False, "(%s is not original: %s; nothing written)" % (s.name, ws.state(s)))
        return
    wem, _ = ws.prepare(s, *tone(min(max(s.seconds, 0.3), 1.0), 700))
    files = ws.files(s)
    before = snapshot(files)
    ws.apply([(s, wem, {"source": "e2e"})])
    ok = ws.state(s) == MODIFIED and ws.current_bytes(s) == wem
    orig_size = ws.store.get(s.key)["orig_size"]
    for l, b in copies_of(ws, s):
        if l.kind == banks.MEMORY or l.size >= orig_size:     # a whole WEM in memory: fills its slot exactly
            ok &= b == fill(wem, l.size)
        else:                                  # prefetch head: the start of the WAD recording (zero-padded)
            full = wem + bytes(max(0, orig_size - len(wem)))
            ok &= b == (full + bytes(l.size))[:l.size]
    ws.revert([s])
    ok &= ws.state(s) == ORIGINAL and snapshot(files) == before and not ws.store.entries
    check(label, ok, "(%s, %d files, %d copies/heads)" % (s.name, len(files), len(ws.locations(s))))


def effects_round_trip(ws, voices, lang):
    """8. A replacement through the effects into a 5-channel line (L R C SL SR), read back and decoded."""
    label = "8 %s 5-channel line replaced with effects" % lang
    s = next((x for x in ws.catalog.sounds if x.archive == voices and x.group == "Speech" and x.channels == 5
              and 1.5 < x.seconds < 4 and ws.state(x) == ORIGINAL), None)
    if s is None:
        check(label, False, "no original 5-channel speech line found")
        return
    new, sr = tone(1.0, 400)
    new[int(0.6 * sr):] = 0                    # a pause, so the echo can be measured after the voice stops
    st = fx.defaults()
    st["eq"].update(on=True, low_cut=300, high_cut=3400)
    st["echo"].update(on=True, time=120, feedback=0.2, mix=0.3)
    st["channels"].update(on=True, front=-12, surround=-6, surround_delay=160)
    x = fx.render(new, sr, st, s.channels)
    wem, _ = ws.prepare(s, x, sr)
    files = ws.files(s)
    before = snapshot(files)
    ws.apply([(s, wem, {"source": "e2e", "effects": fx.summary(st)})])
    y, ysr = audio.decode_wem(ws.current_bytes(s))
    ok = ws.state(s) == MODIFIED and y.shape[1] == 5
    lag = -1.0
    if ok:
        def onset(c):                           # first sample above half of the channel's peak
            a = np.abs(y[:, c])
            return int(np.argmax(a > 0.5 * a.max())) / ysr
        lag = onset(3) - onset(2)               # SL vs C: the PA echo delay
        ok &= abs(lag - 0.16) < 0.02 and np.abs(y[:, 2]).max() > 2 * np.abs(y[:, 0]).max()
    ws.revert([s])
    ok &= ws.state(s) == ORIGINAL and snapshot(files) == before and not ws.store.entries
    check(label, ok, "(%s, surround delay %.3f s)" % (s.name, lag))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--game")
    ap.add_argument("--lang", help="comma-separated languages, or 'all' (default: the game's language)")
    a = ap.parse_args()
    game = Game(a.game or default_game_dir(settings.load()))
    td = tempfile.mkdtemp()
    try:
        first = True
        for lang in languages(game, a.lang, td):
            voices = lang if lang in game.languages else "INT"
            ws = Workspace(game, os.path.join(td, "store_" + lang), [voices, "SFX", "BANK"])
            SPACES.append(ws)
            print("\n== %s (voices %s): %d sounds" % (lang, voices, len(ws.catalog.sounds)))
            run_checks(game, lang, voices, td, ws, first)
            first = False
    finally:
        clean = True
        for w in SPACES:                      # never leave a change behind
            w.revert_all()
            w.texts.revert_all()
            if w.store.entries or w.texts.records():
                clean = False
                print("WARNING: could not revert everything; store: %s" % w.store.root)
        if clean:                             # the stores hold full copies of game files (GBs): drop them
            shutil.rmtree(td, ignore_errors=True)
    print("\n%d/%d PASS" % (sum(ok for _, ok in RESULTS), len(RESULTS)))
    return 0 if RESULTS and all(ok for _, ok in RESULTS) else 1


def languages(game, arg, td):
    if not arg:
        return [game.language() or "INT"]
    if arg.lower() != "all":
        return [x.strip().upper() for x in arg.split(",") if x.strip()]
    tx = Workspace(game, os.path.join(td, "store_list"), [game.languages[0]]).texts
    return sorted(set(game.languages) | set(tx.subtitle_languages()) | set(tx.menu_languages()))


def run_checks(game, lang, voices, td, ws, first):
    bank = [s for s in ws.catalog.sounds if s.archive == "BANK" and s.size >= 2500]
    if first:
        sound_round_trip(ws, max((s for s in bank if not s.language), key=lambda s: len(ws.locations(s))),
                         "1 in-memory sound, every copy")
    loc = [s for s in bank if s.language == voices]
    sound_round_trip(ws, max(loc, key=lambda s: len(ws.locations(s))),
                     "2 %s localized in-memory speech line (%s)" % (lang, voices))
    heads = [ws.catalog.by_key[k] for k in ws.catalog.heads if k in ws.catalog.by_key
             and ws.catalog.by_key[k].archive == voices]
    sound_round_trip(ws, max(heads, key=lambda s: len(ws.locations(s))),
                     "3a %s speech line with prefetch heads" % lang)
    whole = [s for s in heads if any(l.size >= s.size for l in ws.locations(s))]
    if whole:
        sound_round_trip(ws, max(whole, key=lambda s: len(ws.locations(s))),
                         "3b %s speech line also stored whole in a bank" % lang)
    else:
        check("3b %s speech line also stored whole in a bank" % lang, False, "none found")

    tx = ws.texts
    subs = tx.subtitles(lang)

    def text_round_trip(key, new, label):
        files = tx.files(key)
        before = snapshot(files)
        tx.apply({key: new})
        ok = tx.current(key) == new and tx.state(key) == MODIFIED
        for l in tx.subtitle_locations(key) if key.startswith("SUB:") else []:
            b = banks.open_file(ws.bank_path(l.file)).read(l.pos, l.size)
            ok &= subtitles.read_fstring(b, 0)[0] == new
        tx.revert([key])
        ok &= snapshot(files) == before and not tx.records()
        check(label, ok, "(%s, %d files)" % (key[:60], len(files)))

    multi = max((s for s in subs if s.room >= 8), key=lambda s: len(tx.subtitle_locations(subtitles.key(lang, s.speech, s.line))))
    text_round_trip(subtitles.key(lang, multi.speech, multi.line), "E2E test.", "4a %s subtitle with several copies" % lang)
    rsub = next((s for s in subs if s.room >= 10 and
                 {l.file for l in tx.subtitle_locations(subtitles.key(lang, s.speech, s.line))} == {"Startup.upk"}), None)
    if rsub:
        text_round_trip(subtitles.key(lang, rsub.speech, rsub.line), "E2E radio.", "4b %s scripted subtitle (RSubtitle)" % lang)
    else:
        check("4b %s scripted subtitle (RSubtitle)" % lang, False, "no scripted subtitle found in Startup.upk")
    e = next(x for x in tx.menu(lang) if x.file.upper().startswith("FRONTEND"))
    text_round_trip("MENU:" + e.key, "E2E " + e.value, "5 %s menu text" % lang)

    # 6. sound + subtitle in one package
    sub_files = {}
    for s in subs:
        if s.room < len("Shared package."):
            continue
        for l in tx.subtitle_locations(subtitles.key(lang, s.speech, s.line)):
            sub_files.setdefault(l.file, subtitles.key(lang, s.speech, s.line))
    pair = next(((s, sub_files[l.file]) for s in bank if not s.language for l in ws.locations(s)
                 if l.file in sub_files and len(ws.locations(s)) <= 3), None)
    if pair:
        s, key = pair
        files = set(ws.files(s)) | set(tx.files(key))
        before = snapshot(files)
        wem, _ = ws.prepare(s, *tone(0.3, 900))
        ws.apply([(s, wem, {})])
        tx.apply({key: "Shared package."})
        ok = ws.state(s) == MODIFIED and tx.current(key) == "Shared package."
        ws.revert([s])
        ok &= ws.state(s) == ORIGINAL and tx.current(key) == "Shared package."
        tx.revert([key])
        ok &= snapshot(files) == before
        check("6 %s sound + subtitle in one package, reverted separately" % lang, ok, "(%s + %s)" % (s.name, key[:40]))
    else:
        check("6 %s sound + subtitle in one package" % lang, False, "no such package found")

    effects_round_trip(ws, voices, lang)

    # 7. cross-language mod: the same game line, not the same WAD position
    from ackit.catalog import read_csv
    from ackit.paths import catalog_dir
    other = "FRA" if voices == "INT" else "INT"
    other_target = {s.name: s.target for s in read_csv(os.path.join(catalog_dir(), other + ".csv"))}
    other_subs = {(s.speech, s.line) for s in tx.subtitles(other)}
    by_sound = {}
    for s in subs:
        for k in s.sounds:
            by_sound.setdefault(k, s)
    line = next((s for s in ws.catalog.sounds if s.archive == voices and ":" in s.target
                 and s.key in by_sound and by_sound[s.key].room >= 9
                 and (by_sound[s.key].speech, by_sound[s.key].line) in other_subs   # the line exists there too
                 and other_target.get(s.name, s.target) != s.target), None)      # at another position there
    if line is None:
        check("7 mod %s -> %s" % (lang, other), False, "no line found at another position in %s" % other)
        return
    sub = by_sound[line.key]
    sub_key = subtitles.key(lang, sub.speech, sub.line)
    if ws.state(line) != ORIGINAL or tx.state(sub_key) != ORIGINAL:
        check("7 mod %s -> %s" % (lang, other), False, "(%s is not original; nothing written)" % line.name)
        return
    wem, _ = ws.prepare(line, *tone(1.0, 500))
    ws.apply([(line, wem, {})])
    tx.apply({sub_key: "Mod text."})
    pkg = os.path.join(td, "e2e_%s.acmod" % lang)
    export_mod(ws, pkg, ws.modified_keys(), "e2e", texts=tx.modified())
    ws.revert_all()
    tx.revert_all()
    wo = Workspace(game, os.path.join(td, "store_%s_%s" % (lang, other)), [other, "SFX", "BANK"])
    SPACES.append(wo)
    _, items = match_mod(wo, pkg)
    target = items[0].sound
    same_line = target is not None and target.name == line.name
    moved = target is not None and target.target != line.target         # positions differ between languages
    other_sub = subtitles.key(other, sub.speech, sub.line)
    files = set(wo.files(target)) if target else set()
    files |= set(wo.texts.files(other_sub)) | set(wo.texts.files(sub_key))
    before = snapshot(files)
    other_text = wo.texts.current(other_sub)
    res, texts = import_mod(wo, pkg)
    ok = same_line and wo.state(target) == MODIFIED
    ok &= [k for k, _, p in texts if not p] == [sub_key] and wo.texts.current(sub_key) == "Mod text."
    ok &= wo.texts.current(other_sub) == other_text
    wo.revert_all()
    wo.texts.revert_all()
    ok &= snapshot(files) == before
    check("7 mod %s -> %s: same game line%s" % (lang, other, " (another WAD position)" if moved else ""), ok,
          "(%s: %s:%s -> %s:%s)" % (line.name, voices, line.target, other, target.target if target else "-"))


if __name__ == "__main__":
    sys.exit(main())