"""Step 7: join every note from its picked takes and encode it for the game (work/qoa, work/wav48).

- each take trimmed to its speech (top_db 40), 60 ms kept at both ends, 8 ms fades;
- pauses as the EN narrator makes them: PAUSE["s"] inside a paragraph, PAUSE["l"] after a line break,
  PAUSE["p"] between paragraphs (EN medians: phrases ~0.6 s, paragraphs ~0.9 s);
- 24 -> 48 kHz with soxr VHQ; loudness: active-speech RMS matched to the EN original, soft limiter
  at -1 dBFS; 100 ms of silence at both ends; frameless QOA.
The QOA encoder is pure Python: this step is CPU-bound (assemble.workers processes).
"""
import json
import os
import sys
from multiprocessing import Pool

import librosa
import numpy as np
import soundfile as sf

from .config import load
from .qoa import encode

PAUSE = {"s": 0.45, "l": 0.70, "p": 0.90}
CFG = None


def active_rms(x):
    v = x[::5]
    v = v[np.abs(v) > 200 / 32768]
    return float(np.sqrt(np.mean(v * v))) if len(v) else 1e-3


def trim(y, sr):
    iv = librosa.effects.split(y, top_db=40, frame_length=1024, hop_length=240)
    if not len(iv):
        return y
    y = y[max(0, iv[0][0] - int(0.06 * sr)): min(len(y), iv[-1][1] + int(0.06 * sr))].copy()
    f = int(0.008 * sr)
    y[:f] *= np.linspace(0, 1, f)
    y[-f:] *= np.linspace(1, 0, f)
    return y


def job(args):
    n, segs, pick = args
    cfg = CFG or load()
    meta = cfg.w("qoa", f"{n['k']}.json")
    if meta.exists():
        return
    parts = []
    for s in segs:
        y, sr = sf.read(cfg.w("takes", f"{s['take']}_s{pick[s['id']]}.wav"), dtype="float32")
        if parts:
            parts.append(np.zeros(int(PAUSE.get(s["join"], 0.45) * sr), np.float32))
        parts.append(trim(y, sr))
    y = librosa.resample(np.concatenate(parts), orig_sr=24000, target_sr=48000, res_type="soxr_vhq")
    en, _ = sf.read(cfg.w("en24", f"{n['k']}.wav"), dtype="float32")
    gain = active_rms(en) / active_rms(y)
    y = y * gain
    lim = 10 ** (-1 / 20)
    knee = 0.7 * lim
    over = np.abs(y) > knee
    y[over] = np.sign(y[over]) * (knee + (lim - knee) * np.tanh((np.abs(y[over]) - knee) / (lim - knee)))
    pad = np.zeros(4800, np.float32)
    y = np.concatenate([pad, y, pad])
    sf.write(cfg.w("wav48", f"{n['k']}.wav"), y, 48000, subtype="PCM_16")
    pcm = np.clip(np.round(y * 32767), -32768, 32767).astype(np.int32).tolist()
    cfg.w("qoa", f"{n['sound']}.qoa").write_bytes(encode(pcm))
    json.dump(dict(k=n["k"], name=n["name"], vo=n["vo"], sound=n["sound"], frames=len(pcm),
                   seconds=len(pcm) / 48000, gain=round(gain, 3), segments=len(segs)), open(meta, "w"))


def main():
    cfg = load()
    notes = json.load(open(cfg.w("notes.json"), encoding="utf-8"))
    segs = json.load(open(cfg.w("segs.json"), encoding="utf-8"))
    pick = json.load(open(cfg.w("pick.json")))
    by = {}
    for s in segs:
        by.setdefault(s["k"], []).append(s)
    notes = [n for n in notes if n["k"] in by]  # segments may cover a subset (ENSVOICE_NOTES)
    if "--force" in sys.argv:  # re-encode every note (after a new pick)
        for f in cfg.w("qoa", "x").parent.glob("n*.json"):
            f.unlink()
    else:  # re-encode the notes whose picked takes changed
        for n in notes:
            m = cfg.w("qoa", f"{n['k']}.json")
            if m.exists() and json.load(open(m)).get("takes") != [f"{s['take']}_s{pick.get(s['id'])}" for s in by[n["k"]]]:
                m.unlink()
    jobs = [(n, by[n["k"]], pick) for n in notes if not cfg.w("qoa", f"{n['k']}.json").exists()]
    print("notes to encode", len(jobs), flush=True)
    with Pool(int(cfg.get("assemble", "workers", max(1, (os.cpu_count() or 2) - 1)))) as p:
        for i, _ in enumerate(p.imap_unordered(job, jobs, chunksize=1)):
            if i % 25 == 0:
                print(f"{i + 1}/{len(jobs)}", flush=True)
    for n in notes:  # remember which takes each note was built from
        m = cfg.w("qoa", f"{n['k']}.json")
        d = json.load(open(m))
        d["takes"] = [f"{s['take']}_s{pick[s['id']]}" for s in by[n["k"]]]
        json.dump(d, open(m, "w"))
    total = sum(json.load(open(cfg.w("qoa", f"{n['k']}.json")))["seconds"] for n in notes)
    print("ASSEMBLED", len(notes), "notes,", round(total / 60, 1), "minutes")


if __name__ == "__main__":
    main()