"""PoC (работа A, этап A0): text-to-audio через AudioLDM2.

Главный кандидат плана — Stable Audio Open Small — gated (нужен HF-токен),
поэтому PoC делаем на фолбэке AudioLDM2 (cvssp/audioldm2, CC-BY-NC-SA —
в ОСНОВНОЙ пайплайн с некоммерческими весами не пойдёт, см. docs/plan.md).

Запуск:
  gen_audioldm2.py                      # все клипы из LIST с seed по умолчанию
  gen_audioldm2.py <id> [id[:seed]...]  # один/несколько клипов
Клип = запись LIST (промпт, длительность, шаги); выход — review/<id>_<seed>.wav.

Детерминизм: seed + шаги фиксированы, промпт в манифесте; детерминизм
по-устройству (fp16 GPU ≠ fp32 CPU). GPU (RTX 3090) при наличии: fp16.
"""
import sys
from pathlib import Path

import numpy as np
import torch

HERE = Path(__file__).resolve().parent
REVIEW = HERE / "review"

NEG_NOISE = (
    "music, melody, rhythm, drums, speech, voices, sudden loud noises, "
    "harsh noise, distortion"
)

# Клипы PoC: амбиенты (длинные, лупуются пост-обработкой) и SFX (короткие,
# без лупа). Эталоны процедурных — tools/audio/gen.mjs и data/sfxSpecs.ts.
LIST = {
    # --- амбиенты (10 с, луп) ---
    "amb_ponds": {
        "positive": (
            "field recording, calm misty pond at dusk, gentle water lapping, "
            "soft wind through reeds, occasional distant croaking, eerie "
            "quiet melancholic atmosphere, seamless ambience"
        ),
        "negative": NEG_NOISE,
        "duration": 10.0,
    },
    "amb_zvenets": {
        "positive": (
            "field recording, quiet rural settlement at dusk, low distant "
            "metallic bell toll, soft wind, faint rustling, somber "
            "melancholic atmosphere, seamless ambience"
        ),
        "negative": NEG_NOISE,
        "duration": 10.0,
    },
    "amb_meadows": {
        "positive": (
            "field recording, desolate burned field, cold dry wind blowing "
            "over dead grass, occasional faint debris rattling, bleak "
            "post-fire emptiness, seamless ambience"
        ),
        "negative": NEG_NOISE,
        "duration": 10.0,
    },
    "amb_ponds_night": {
        "positive": (
            "field recording, pond at night, deeper water lapping, night "
            "insects, distant eerie echoes, dark oppressive quiet, "
            "seamless ambience"
        ),
        "negative": NEG_NOISE,
        "duration": 10.0,
    },
    "amb_tower_hum": {
        "positive": (
            "deep metallic resonance hum, giant bell tower vibrating, low "
            "sustained drone with slow shimmer, ominous and ancient"
        ),
        "negative": NEG_NOISE,
        "duration": 10.0,
    },
    # --- SFX (короткие, без лупа) ---
    "sfx_step": {
        "positive": (
            "single soft footstep on dry crumbly ash and dead grass, short "
            "subtle crunch, close up, dry foley"
        ),
        "negative": NEG_NOISE,
        "duration": 2.5,
    },
    "sfx_click": {
        "positive": (
            "single short UI click, soft wooden tap, clean and dry, "
            "interface sound"
        ),
        "negative": NEG_NOISE,
        "duration": 1.5,
    },
    "sfx_bell": {
        "positive": (
            "single bronze bell strike, resonant metallic ring with long "
            "natural decay, solemn, temple bell, one hit only"
        ),
        "negative": NEG_NOISE,
        "duration": 5.0,
    },
    "sfx_whoosh": {
        "positive": (
            "single quick air whoosh, fast swing passing by, short swoosh "
            "transition"
        ),
        "negative": NEG_NOISE,
        "duration": 2.0,
    },
    "sfx_hit": {
        "positive": (
            "single dull impact hit, heavy blunt thud with brief echo, "
            "combat punch sound, one hit only"
        ),
        "negative": NEG_NOISE,
        "duration": 2.0,
    },
    "sfx_pickup": {
        "positive": (
            "single gentle collect chime, soft bright pling, short "
            "delicate pickup sound, one note only"
        ),
        "negative": NEG_NOISE,
        "duration": 2.0,
    },
    "sfx_drop": {
        "positive": (
            "single water drop into calm pond, soft plop with tiny ripple, "
            "close up, one drop only"
        ),
        "negative": NEG_NOISE,
        "duration": 2.0,
    },
}

DEFAULT_SEED = 1


def load_pipeline():
    from diffusers import AudioLDM2Pipeline

    use_gpu = torch.cuda.is_available()
    dtype = torch.float16 if use_gpu else torch.float32
    pipe = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2", torch_dtype=dtype)
    pipe.to("cuda" if use_gpu else "cpu")
    if not use_gpu:
        torch.set_num_threads(12)
    return pipe, use_gpu


def main() -> None:
    # Аргументы — id[:seed]; без аргументов — весь LIST с DEFAULT_SEED.
    jobs: list[tuple[str, int]] = []
    for arg in sys.argv[1:]:
        ident, _, seed = arg.partition(":")
        if ident not in LIST:
            print(f"неизвестный клип {ident!r}; доступно: {', '.join(LIST)}")
            sys.exit(1)
        jobs.append((ident, int(seed) if seed else DEFAULT_SEED))
    if not jobs:
        jobs = [(name, DEFAULT_SEED) for name in LIST]

    import diffusers  # noqa: F401  (после разбора аргументов — чтобы --help был быстрым)

    pipe, use_gpu = load_pipeline()
    dev = "cuda fp16" if use_gpu else "cpu fp32"
    REVIEW.mkdir(parents=True, exist_ok=True)

    for ident, seed in jobs:
        spec = LIST[ident]
        torch.manual_seed(seed)
        np.random.seed(seed)
        print(f"генерация {ident} (seed={seed}, {spec['duration']} с, {dev})...")
        audio = pipe(
            spec["positive"],
            negative_prompt=spec["negative"],
            num_inference_steps=100,
            guidance_scale=3.0,
            audio_length_in_s=spec["duration"],
        ).audios[0]
        audio = np.clip(audio, -1.0, 1.0)
        pcm = (audio * 32767.0).astype(np.int16)
        import wave

        out = REVIEW / f"{ident}_{seed}.wav"
        with wave.open(str(out), "wb") as w:
            w.setnchannels(1)
            w.setsampwidth(2)
            w.setframerate(16000)
            w.writeframes(pcm.tobytes())
        print(f"готово: {out.relative_to(HERE)} ({pcm.shape[0] / 16000:.1f} с)")


if __name__ == "__main__":
    main()