#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""《回声地窟》Stable Audio 3 small-sfx 音效后处理（SA3 补产批次，16-dungeon.md §9.2）。

批次 2 的 25 个 id 此前为 SYN 兜底（gen_sfx2.py，§9.2 降级条款）；本脚本产 SA3 候选版，
与在役 SYN 版做 A-B（响度/时长/包络对齐），质量明显更好才同名晋级 audio/sfx/（游戏层零改动）。
prompt 出自 §9.2 表（echo 距离档/tel 前摇的英文稿按中文说明补全）。

生成（GPU，flock 串行；模型 gated 走 hf-mirror，见 env.d/sfx.sh）：
  python3 gen_sfx_sa3.py --emit   # 打印生成命令，原样执行
生成 13 条 mono 源原始 wav 到 gen/staging/sfx_sa3/dungeon_<i>.wav（顺序 = 下表 idx）。
后处理：修整/去直流/8ms cos² 淡化/源真峰值 -4dBFS（同在役 SYN 档）→
mob_step_*/tel_* 派生预 pan 立体声（l=能量只在 L / r=只在 R / c=双声道同相复制，
工艺同 gen_sfx2.py pan_derive/stereo_dup）→
ogg libvorbis q4 + m4a aac 128k（highpass=f=30，前后垫 60/120ms 静音）→
解码回读 QC：逐声道真峰值 <= -3dBFS、|直流| <= 0.001；超限降增益 + 首轮重编加
18kHz 低通，最多 3 轮补偿。候选落 gen/staging/sfx_sa3/out/（不入库），A-B 后晋级。

用法：.venv-tts/bin/python web/dungeon/gen/gen_sfx_sa3.py [--emit]
"""
import os
import subprocess
import sys

import numpy as np
import soundfile as sf

SR = 44100
HERE = os.path.dirname(os.path.abspath(__file__))
STAGING = os.path.join(HERE, "staging", "sfx_sa3")
OUT = os.path.join(STAGING, "out")
DST = os.path.join(HERE, "..", "audio", "sfx")
SA3 = "/home/xsl/SoundGame/tools/stable-audio-3"
FFMPEG = "/home/xsl/SoundGame/tools/bin/ffmpeg"
SEED = 20260902

# (id, gen 秒, prompt) —— §9.2 标注 SA3 的 13 个 mono 源；定档全部 -4dBFS（同在役 SYN）
ITEMS = [
    ("echo_1", 2.5, "close stone wall echo slap-back, bright, short, dry"),
    ("echo_2", 3.0, "mid distance cave echo, moderate tail"),
    ("echo_3", 3.5, "far cave echo, darker longer"),
    ("echo_4", 4.5, "very far cavern echo, long dark tail, huge space"),
    ("stairs_down", 4.0, "descending old metal stairs, deep cavern reverb"),
    ("atk_hit", 2.0, "blunt impact on hard insect chitin shell, dull thud hit"),
    ("enemy_die", 2.5, "monster body collapsing on stone floor, heavy fall thud and debris"),
    ("mob_step_small", 2.5, "small beast skittering steps, dry"),
    ("mob_step_bug", 2.5, "insect chitin crawling clicks"),
    ("mob_step_heavy", 2.5, "heavy armored thudding steps"),
    ("tel_small", 2.5, "small beast pounce hiss, short aggressive screech"),
    ("tel_bug", 2.5, "insect legs chitin rubbing scrape, creepy skitter"),
    ("tel_heavy", 2.5, "huge monster stomp and low growl, heavy rumble"),
]
# pan 派生表（与在役文件集一致：mob_step 有 _c，tel 无 _c）
DERIVE = {}
for _n in ["mob_step_small", "mob_step_bug", "mob_step_heavy"]:
    DERIVE[_n] = ["_l", "_c", "_r"]
for _n in ["tel_small", "tel_bug", "tel_heavy"]:
    DERIVE[_n] = ["_l", "_r"]

# 事件窗长度（秒）：raw 10.59s 定长 latent 床上按自响起点截取；None=按 -50dBFS 活动窗
# 自然修剪。echo 档长 0.65<1.25<2.0<2.8 保持与在役 SYN 档（0.58/1.08/1.68/2.38）同序递增
# ——回声尾长=墙距是本作核心信息，A-B 不满足递增即不晋级。
WIN = {
    "echo_1": 0.65, "echo_2": 1.25, "echo_3": 2.0, "echo_4": 2.8,
    "stairs_down": None, "atk_hit": None, "enemy_die": 1.7,
    "mob_step_small": 0.45, "mob_step_bug": 0.55, "mob_step_heavy": 0.65,
    "tel_small": 1.0, "tel_bug": 1.1, "tel_heavy": 1.3,
}
# 响度对齐：候选按在役 SYN 版（同名 ogg）RMS 定档（-4dBFS 峰值帽不变），保游戏混音参数
SYN_DIR = os.path.join(HERE, "..", "audio", "sfx")


def emit():
    durs = " ".join(str(d) for _, d, *_ in ITEMS)
    ps = " ".join(f'"{p}"' for _, _, p, *_ in ITEMS)
    print(f"cd {SA3} && flock /home/xsl/SoundGame/.gpu.lock .venv/bin/stable-audio "
          f"--model small-sfx --steps 8 --seed {SEED} "
          f"-p {ps} --duration {durs} -o {STAGING}/dungeon.wav")


def trim_oneshot(y, floor_db=-50.0, head_ms=50, tail_ms=140, max_s=None):
    env = np.max(np.abs(y), axis=1)
    act = np.where(env > 10 ** (floor_db / 20))[0]
    a = max(0, act[0] - int(SR * head_ms / 1000))
    b = min(len(y), act[-1] + int(SR * tail_ms / 1000))
    if max_s:
        b = min(b, a + int(SR * max_s))
    return y[a:b]


def prep(y, target_db=-4.0, win_len=None, syn_rms_db=None):
    if y.ndim == 1:
        y = y[:, None]
    if win_len is not None:
        # 自响起点：10ms 窗包络首个超过 (峰值-20dB) 的位置，提前 20ms 起
        w = int(SR * 0.01)
        nfr = len(y) // w
        eb = 20 * np.log10(np.max(np.abs(y[:nfr * w, 0]).reshape(nfr, w), axis=1) + 1e-9)
        pk = 20 * np.log10(np.max(np.abs(y[:, 0])) + 1e-9)
        on = np.where(eb > pk - 20)[0][0] * w
        a = max(0, on - int(SR * 0.02))
        b = min(len(y), a + int(SR * win_len))
        y = y[a:b]
        nf2 = max(1, int(SR * 0.050))                    # 窗切在持续声中，切点 50ms 缓出
        y[-nf2:] *= (0.5 * (1 + np.cos(np.pi * np.arange(nf2) / nf2)))[:, None]
    else:
        y = trim_oneshot(y)
    y = y - np.mean(y, axis=0, keepdims=True)          # 去直流
    nf = max(1, int(SR * 0.008))
    fade = 0.5 * (1 - np.cos(np.pi * np.arange(nf) / nf))
    y[:nf] *= fade[:, None]                             # 8ms cos² 淡入淡出
    y[-nf:] *= fade[::-1][:, None]
    y = y / (np.max(np.abs(y)) + 1e-9) * 10 ** (target_db / 20)
    pad_h = np.zeros((int(SR * 0.06), y.shape[1]))      # 编码器软着陆（同 gen_sfx2）
    pad_t = np.zeros((int(SR * 0.12), y.shape[1]))
    y = np.concatenate([pad_h, y, pad_t])
    if syn_rms_db is not None:                          # A-B 响度对齐（峰值帽 -4 不变）。
        rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9)  # 垫后计算：与解码回读同口径
        s = min(10 ** ((syn_rms_db - rms) / 20), 1.0)
        y = y * s
    return y


def _syn_metrics(name):
    """在役 SYN 版 ogg 解码回读（时长, RMS）——响度对齐基准。"""
    p = os.path.join(SYN_DIR, name + ".ogg")
    subprocess.run([FFMPEG, "-y", "-v", "error", "-i", p, "-c:a", "pcm_f32le", "/tmp/_syn_qc.wav"], check=True)
    y, sr = sf.read("/tmp/_syn_qc.wav", always_2d=True)
    return y.shape[0] / sr, 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9)


def derive(name, mono):
    """mono -> 预 pan 立体声（能量只在 L 或 R；c = 双声道同相复制），工艺同 gen_sfx2。"""
    outs = []
    for suf in DERIVE.get(name, []):
        if suf == "_l":
            outs.append((name + suf, np.stack([mono[:, 0], np.zeros(len(mono))], axis=1)))
        elif suf == "_r":
            outs.append((name + suf, np.stack([np.zeros(len(mono)), mono[:, 0]], axis=1)))
        elif suf == "_c":
            outs.append((name + suf, np.stack([mono[:, 0], mono[:, 0]], axis=1)))
    return outs


def encode(name, sig):
    base = sig.astype(np.float64)
    vol, lpf, dc_adj = 0.0, False, 0.0
    tmp = os.path.join(OUT, "_enc.wav")
    chk = os.path.join(OUT, "_qc.wav")
    res = {}
    for attempt in range(4):  # 首发 + 最多 3 轮补偿（首轮重编加 18kHz 低通）
        x = base - dc_adj
        if vol < 0:
            x = x * (10 ** (vol / 20))
        sf.write(tmp, x.astype(np.float32), SR)
        af = "highpass=f=30,lowpass=f=18000" if lpf else "highpass=f=30"
        for ext, cmd in [("ogg", ["-c:a", "libvorbis", "-q:a", "4"]),
                         ("m4a", ["-c:a", "aac", "-b:a", "128k"])]:
            out = os.path.join(OUT, f"{name}.{ext}")
            subprocess.run([FFMPEG, "-y", "-v", "error", "-i", tmp, "-af", af,
                            "-ar", "44100", *cmd, out], check=True)
            subprocess.run([FFMPEG, "-y", "-v", "error", "-i", out, "-c:a", "pcm_f32le", chk], check=True)
            y, sr = sf.read(chk, always_2d=True)
            tp = 20 * np.log10(np.max(np.abs(y)) + 1e-9)
            dcv = float(np.mean(y.mean(axis=1)))
            rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9)
            ok_tp, ok_dc = tp <= -3.0, abs(dcv) <= 0.001
            res[ext] = (tp, dcv, rms, y.shape[0] / sr)
            if not (ok_tp and ok_dc):
                if not ok_tp:
                    vol -= (tp + 3.0) + 0.25
                    lpf = True       # 首轮重编起加 18kHz 低通压 MDCT 过冲
                if not ok_dc:
                    dc_adj -= dcv
                break
        else:
            break
    else:
        raise AssertionError(f"{name}: 补偿 3 轮后仍超限 {res}")
    return res


def main():
    os.makedirs(STAGING, exist_ok=True)
    os.makedirs(OUT, exist_ok=True)
    if "--emit" in sys.argv:
        emit()
        return
    print(f"| {'id':18s} | {'时长s':>6s} | {'ogg峰dBFS':>9s} | {'m4a峰dBFS':>9s} | {'RMS dBFS':>8s} | {'SYN时长':>7s} | {'SYN RMS':>8s} | 直流 | 声道 |")
    for i, (name, _d, _p) in enumerate(ITEMS):
        raw = os.path.join(STAGING, f"dungeon_{i}.wav")
        y, sr = sf.read(raw, always_2d=True, dtype="float64")
        assert sr == SR, f"{name}: 采样率 {sr}"
        mono = y[:, :1]
        syn_dur, syn_rms_db = _syn_metrics(name)
        # 响度对齐只做一次（按 mono 基准），派生 l/r/c 直接复制基准样本——
        # 与在役 SYN 派生工艺一致：l/r 文件 RMS = 基准 -3dB（能量单声道承载），逐声道振幅不变
        base_sig = prep(mono, win_len=WIN.get(name), syn_rms_db=syn_rms_db)
        for eid, sig in [(name, base_sig)] + derive(name, base_sig):
            res = encode(eid, sig)
            tp_o, dc, rms, dur = res["ogg"]
            tp_m = res["m4a"][0]
            ch = 1 if eid == name else 2
            sd = f"{syn_dur:6.2f}" if eid == name else "   -"
            print(f"| {eid:18s} | {dur:6.2f} | {tp_o:9.2f} | {tp_m:9.2f} | {rms:8.2f} | {sd} | {syn_rms_db:7.2f} | {dc:+.4f} | {ch:d} |")
    print("dungeon SA3 候选完成 -> gen/staging/sfx_sa3/out/（A-B 复核后晋级）")


if __name__ == "__main__":
    main()