#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""《深夜茶馆》Stable Audio 3 small-sfx 音效（20-teahouse.md §8.2 SA3 条目，5 件）。

生成（GPU，flock 串行；模型已入缓存，HF_HUB_OFFLINE 可跑，见 env.d/sfx.sh）：
  cd tools/stable-audio-3 && flock /home/xsl/SoundGame/.gpu.lock \
    .venv/bin/stable-audio --model small-sfx --steps 8 --seed 20 \
    -p "<prompts>" --duration <secs> -o <staging>/teahouse.wav
本脚本随后做后处理：自然修剪/去直流/8ms cos² 淡化/响度定档 → ogg q3 + m4a 128k
→ 解码回读 QC：真峰值 ≤ -3dBFS、|直流| ≤ 0.001；超限降增益 + 重编，最多 3 轮。

prompt：20-teahouse.md §8.2 条目英文化。定档（§3.3 频带分离 + 操作音不抢人声）：
  door_bell -8（铜铃两响，进门标记）、steps_in/out -12（脚步远低近高，床级）
  water_pour -8（注水 2s 主操作音）、pour_cup -8（上茶落盏）

用法：.venv-tts/bin/python web/teahouse/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")
DST = os.path.join(HERE, "..", "audio", "sfx")
SA3 = "/home/xsl/SoundGame/tools/stable-audio-3"
FFMPEG = "/home/xsl/SoundGame/tools/bin/ffmpeg"
SEED = 20
os.makedirs(STAGING, exist_ok=True)
os.makedirs(DST, exist_ok=True)

# (id, gen 秒, prompt, 定档 dBFS, trim 上限 s)
ITEMS = [
    ("door_bell", 4.0, "Small brass shop door bell rings twice, warm clear chime, quiet night tea house, close", -8.0, 1.2),
    ("steps_in", 6.0, "Soft cloth shoes walking on wooden floor, slow approach from far to near, gentle steps, quiet room, close", -12.0, 3.5),
    ("steps_out", 6.0, "Soft cloth shoes walking on wooden floor, slow walk away from near to far, gentle steps fading, quiet room, close", -12.0, 3.5),
    ("water_pour", 4.0, "Water pouring into a metal kettle, steady pour, two seconds, close, clean recording", -8.0, 2.6),
    ("pour_cup", 4.0, "Hot tea poured into a small ceramic cup, gentle pour ending with cup placed on saucer, close", -8.0, 2.6),
]

WIN = {}


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}/teahouse.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, max_s=None):
    if win:
        a, b = int(SR * win[0]), int(SR * win[1])
        y = y[a:b]
        nf2 = max(1, int(SR * 0.050))
        y[-nf2:] *= (0.5 * (1 + np.cos(np.pi * np.arange(nf2) / nf2)))[:, None]
    else:
        y = trim_oneshot(y, max_s=max_s)
    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]
    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], axis=0)


def encode(wav_path, base):
    for args, ext in [(["-c:a", "libvorbis", "-q:a", "3"], ".ogg"),
                      (["-c:a", "aac", "-b:a", "128k"], ".m4a")]:
        subprocess.run([FFMPEG, "-y", "-v", "error", "-i", wav_path,
                        "-ar", "44100", *args, base + ext], check=True)


def true_peak_db(path):
    r = subprocess.run([FFMPEG, "-i", path, "-af", "ebur128=peak=true", "-f", "null", "-"],
                       capture_output=True, text=True)
    peaks = [float(line.split("Peak:")[1].strip().split()[0])
             for line in r.stderr.splitlines() if "Peak:" in line and "dBFS" in line]
    return max(peaks) if peaks else float("nan")


def check_tp(base):
    return max(true_peak_db(base + ".ogg"), true_peak_db(base + ".m4a"))


def main():
    # stable-audio CLI 多 prompt 输出 <out>_<idx>.wav（每 prompt 一支）
    if not os.path.exists(os.path.join(STAGING, "teahouse_0.wav")):
        emit()
        sys.exit("请先执行 --emit 打印的命令生成 raw wav 后重跑本脚本")
    report = []
    problems = []
    for i, (fid, _d, _p, db, max_s) in enumerate(ITEMS):
        y, sr = sf.read(os.path.join(STAGING, f"teahouse_{i}.wav"))
        assert sr == SR, f"sr {sr}" if False else None
        wav_path = os.path.join(STAGING, fid + ".wav")
        base = os.path.join(DST, fid)
        out = prep_oneshot(y, db, win=WIN.get(fid), max_s=max_s)
        sf.write(wav_path, out, SR)
        encode(wav_path, base)
        tp = check_tp(base)
        for _round in range(3):
            if tp <= -3.0:
                break
            gain = 10 ** ((-3.0 - tp) / 20) * 0.98
            out = out * gain
            sf.write(wav_path, out, SR)
            encode(wav_path, base)
            tp = check_tp(base)
        dc_val = float(np.max(np.abs(np.mean(out, axis=0))))
        report.append((fid, len(out) / SR, tp, dc_val))
        if tp > -3.0:
            problems.append(f"{fid}: 真峰值 {tp:.2f} 超 -3dBFS")
        if dc_val > 0.001:
            problems.append(f"{fid}: 直流 {dc_val:.5f}")
    print(f"{'id':12s} {'dur(s)':>7s} {'TP(dBFS)':>9s} {'|dc|':>8s}")
    for fid, dur, tp, dcv in report:
        print(f"{fid:12s} {dur:7.2f} {tp:9.2f} {dcv:8.5f}")
    if problems:
        print("PROBLEMS:")
        for p in problems:
            print(" ", p)
        return 1
    print(f"ALL OK: {len(ITEMS)} SA3 -> {DST}")
    return 0


if __name__ == "__main__":
    if "--emit" in sys.argv:
        emit()
        sys.exit(0)
    sys.exit(main())