#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""《不眠频率》Stable Audio 3 small-sfx 音效后处理（SA3 补产批次，15-radio.md §5.2）。

生成（GPU，flock 串行；模型 gated 走 hf-mirror，见 env.d/sfx.sh）：
  python3 gen_sfx_sa3.py --emit   # 打印生成命令，原样执行
生成 12 条原始 wav 到 gen/staging/sfx_sa3/radio_<i>.wav（顺序 = 下表 idx）。
本脚本随后做后处理：修整/去直流/8ms cos² 淡化（loop 类只去直流不淡化）/
响度定档（源真峰值或 RMS 定档，escape amb 床 -26dBFS RMS 先例）→
ogg libvorbis q4 + m4a aac 128k（highpass=f=30，one-shot 前后垫 60/120ms 静音）→
解码回读 QC：逐声道真峰值 <= -3dBFS、|直流| <= 0.001；超限降增益 + 首轮重编
加 18kHz 低通，最多 3 轮补偿。loop 类（phone_ring）另验解码首/末样本 <= -60dBFS。
候选 ogg/m4a 落 gen/staging/sfx_sa3/out/（不入库），人工/A-B 复核后晋级 audio/sfx/。

prompt 均出自 15-radio.md §5.2（英文原稿）；定档理由：
  onair_sting/brake_screech/phone_ring -4（台标/危机拍/铃声，主层次）
  line_connect/line_drop -7（line_down 细节咔哒，subtle/soft）
  arg_muffled -8（C3 信息层线索，垫在人声下但可闻）
  tape_static -5（p9 直面道具声）
  amb_* RMS -26 / 峰值帽 -8（通话侧环境床，escape amb_cabinet/-26dBFS 先例）

用法：.venv-tts/bin/python web/radio/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, 模式 oneshot/loop/amb, 定档 dBFS)
ITEMS = [
    ("onair_sting", 3.0, "Short broadcast on-air light sting, warm chime with vintage radio click, 1 second", "oneshot", -4.0),
    ("phone_ring", 6.0, "Old studio telephone bell, two-strike bright metallic ring, clean", "loop", -4.0),
    ("line_connect", 3.0, "Telephone line connect click and short dial tone fade, subtle", "oneshot", -7.0),
    ("line_drop", 3.0, "Telephone hang-up click with short busy tone, soft", "oneshot", -7.0),
    ("amb_rain_city", 10.0, "Heavy rain on city street at night, distant traffic hiss, seamless loop", "amb", -26.0),
    ("amb_station", 10.0, "Empty train station hall at night, distant echo and rolling suitcase, seamless loop", "amb", -26.0),
    ("amb_kid_room", 10.0, "Quiet child bedroom at night, soft lamp buzz, muffled house silence, seamless loop", "amb", -26.0),
    ("arg_muffled", 5.0, "Muffled argument behind a door, dull thud of door slam, indistinct voices, 4 seconds", "oneshot", -8.0),
    ("amb_street_rain", 10.0, "Rain at an empty bus stop at night, distant wet road traffic, seamless loop", "amb", -26.0),
    ("amb_cab", 10.0, "Inside a moving truck cab at night, steady engine hum, rhythmic windshield wipers, rain outside, seamless loop", "amb", -26.0),
    ("brake_screech", 3.0, "Sudden brake screech on wet asphalt, brief tire friction, 2 seconds", "oneshot", -4.0),
    ("tape_static", 4.0, "Old cassette tape static and tape hiss with a click stop, 3 seconds", "oneshot", -5.0),
    # amb_studio：文档 §5.2 有、game.js 未接线（直播间底噪由 bgm/onair 承担）——产出备用
    ("amb_studio", 10.0, "Quiet radio studio room tone, console hum and soft air conditioning, seamless loop", "amb", -26.0),
]
RING_PERIOD = 5.2  # RING 循环周期 >=4s（设计 §5.2），烘进文件

# 事件窗（raw 16.16s 定长 latent 床上截取目标事件；None=按 -50dBFS 活动窗自然修剪）。
# 窗中心取实测事件：arg_muffled 摔门峰 @5.00s、brake_screech 最强吱咬 @3.91s、
# tape_static 停带咔哒 @3.72s（seed 20260902 实测）。
WIN = {
    "arg_muffled": (1.0, 5.4),      # 争吵声垫 -> 摔门 -> 余渍，≈设计 4s
    "brake_screech": (3.6, 5.8),    # 最强急刹吱咬 2.2s（设计 2s 危机拍）
    "tape_static": (0.8, 4.15),     # 电流嘶声 -> 停带咔哒，≈设计 3s
}


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}/radio.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_oneshot(y, target_db, win=None):
    if win:
        a, b = int(SR * win[0]), int(SR * win[1])
        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]))      # 编码器软着陆
    pad_t = np.zeros((int(SR * 0.12), y.shape[1]))
    return np.concatenate([pad_h, y, pad_t])


def prep_loop(y, target_db, period_s):
    """铃声循环垫：修剪衰减尾 -> 循环边界 10ms 压零（不做整段淡化）-> 补静音到周期。"""
    y = trim_oneshot(y, floor_db=-50.0, head_ms=30, tail_ms=30, max_s=2.6)
    y = y - np.mean(y, axis=0, keepdims=True)
    nf = max(1, int(SR * 0.010))
    y[:nf] *= np.linspace(0, 1, nf)[:, None]            # 仅边界压零，防 loop 爆点
    y[-nf:] *= np.linspace(1, 0, nf)[:, None]
    y = y / (np.max(np.abs(y)) + 1e-9) * 10 ** (target_db / 20)
    need = int(SR * period_s)
    if len(y) < need:
        y = np.concatenate([y, np.zeros((need - len(y), y.shape[1]))])
    return y


def prep_amb(y, rms_db, peak_db):
    """环境床：只去直流，不淡化（游戏层 0.4s 重触发重叠掩接缝）；RMS 定档 + 峰值帽。"""
    y = y - np.mean(y, axis=0, keepdims=True)
    rms = 20 * np.log10(np.sqrt(np.mean(y ** 2)) + 1e-9)
    peak = 20 * np.log10(np.max(np.abs(y)) + 1e-9)
    s = min(10 ** ((rms_db - rms) / 20), 10 ** ((peak_db - peak) / 20))
    return y * s


def encode(name, sig, loop=False):
    base = sig.astype(np.float64)
    vol, lpf, dc_adj = 0.0, False, 0.0
    edge_n = 0  # loop 边界超限时逐轮加宽压零区
    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))
        if edge_n:
            ne = int(SR * 0.005 * edge_n)
            x[:ne] = 0.0
            x[-ne:] = 0.0
        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
            ok_loop = True
            if loop:
                edge = 10 ** (-60.0 / 20)
                ok_loop = (float(np.max(np.abs(y[0, :]))) <= edge
                           and float(np.max(np.abs(y[-1, :]))) <= edge)
            res[ext] = (tp, dcv, rms, y.shape[0] / sr)
            if not (ok_tp and ok_dc and ok_loop):
                if not ok_tp:
                    vol -= (tp + 3.0) + 0.25
                    lpf = True       # 首轮重编起加 18kHz 低通压 MDCT 过冲
                if not ok_dc:
                    dc_adj -= dcv
                if not ok_loop:
                    edge_n += 1      # 循环边界再压 5ms 零
                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':16s} | {'时长s':>6s} | {'ogg峰dBFS':>9s} | {'m4a峰dBFS':>9s} | {'RMS dBFS':>8s} | 直流 |")
    for i, (name, _d, _p, mode, tgt) in enumerate(ITEMS):
        raw = os.path.join(STAGING, f"radio_{i}.wav")
        y, sr = sf.read(raw, always_2d=True, dtype="float64")
        assert sr == SR, f"{name}: 采样率 {sr}"
        if y.shape[1] == 1:
            y = np.stack([y[:, 0], y[:, 0]], axis=1)
        if mode == "oneshot":
            sig = prep_oneshot(y, tgt, win=WIN.get(name))
        elif mode == "loop":
            sig = prep_loop(y, tgt, RING_PERIOD)
        else:
            sig = prep_amb(y, tgt, -8.0)
        res = encode(name, sig, loop=(mode == "loop"))
        tp_o, dc, rms, dur = res["ogg"]
        tp_m = res["m4a"][0]
        print(f"| {name:16s} | {dur:6.2f} | {tp_o:9.2f} | {tp_m:9.2f} | {rms:8.2f} | {dc:+.4f} |")
    print("radio SA3 候选完成 -> gen/staging/sfx_sa3/out/（晋级前先复核）")


if __name__ == "__main__":
    main()