#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import wave
from array import array
from pathlib import Path

try:
    import soundfile as sf  # type: ignore
except Exception:  # pragma: no cover - optional runtime dependency
    sf = None

try:
    import torch  # type: ignore
except Exception:  # pragma: no cover - optional runtime dependency
    torch = None

try:
    import numpy as np  # type: ignore
except Exception:  # pragma: no cover - optional runtime dependency
    np = None


def main() -> int:
    parser = argparse.ArgumentParser(description="Local cinematic asset synthesizer")
    parser.add_argument("kind", choices=["voiceover", "music", "sfx"])
    parser.add_argument("--prompt-file", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--duration", type=float, default=8.0)
    parser.add_argument("--model", default="")
    parser.add_argument("--piper-binary", default=os.getenv("LOCAL_ASSET_PIPER_BINARY", "piper"))
    args = parser.parse_args()

    prompt = Path(args.prompt_file).read_text(encoding="utf-8").strip()
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)

    if args.kind == "voiceover":
        if synthesize_voiceover(prompt, output, args.piper_binary, args.model):
            return 0
        return 2

    if args.kind == "music":
        if synthesize_audiocraft("music", prompt, output, args.duration, args.model):
            return 0
        return 2

    if synthesize_audiocraft("sfx", prompt, output, args.duration, args.model):
        return 0
    return 2


def synthesize_voiceover(prompt: str, output: Path, piper_binary: str, model: str) -> bool:
    model_path = model.strip()
    if model_path and shutil.which(piper_binary):
        command = [piper_binary, "--model", model_path, "--output_file", str(output)]
        try:
            process = subprocess.run(
                command,
                input=prompt.encode("utf-8"),
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                check=False,
            )
            return process.returncode == 0 and output.is_file()
        except Exception:
            return False

    return False


def synthesize_audiocraft(kind: str, prompt: str, output: Path, duration: float, model_name: str) -> bool:
    model_path = Path(model_name).expanduser() if model_name else None
    if torch is None or model_path is None or not model_path.exists():
        return False
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
    try:
        if kind == "music":
            from audiocraft.models import MusicGen  # type: ignore

            model = MusicGen.get_pretrained(str(model_path.resolve()))
            model.set_generation_params(duration=max(1.0, float(duration)))
            audio = model.generate([prompt])[0]
            write_tensor_audio(audio, output, getattr(model, "sample_rate", 32000))
            return output.is_file()

        from audiocraft.models import AudioGen  # type: ignore

        model = AudioGen.get_pretrained(str(model_path.resolve()))
        model.set_generation_params(duration=max(0.5, float(duration)))
        audio = model.generate([prompt])[0]
        write_tensor_audio(audio, output, getattr(model, "sample_rate", 32000))
        return output.is_file()
    except Exception:
        return False

def write_tensor_audio(audio, output: Path, sample_rate: int) -> None:
    if np is None:
        raise RuntimeError("numpy is required to write model-generated audio")
    if hasattr(audio, "detach"):
        audio = audio.detach().cpu().numpy()
    elif torch is not None and isinstance(audio, torch.Tensor):
        audio = audio.cpu().numpy()
    audio = np.asarray(audio)
    if audio.ndim == 3:
        audio = audio[0]
    if audio.ndim == 2 and audio.shape[0] < audio.shape[1]:
        audio = audio.transpose(1, 0)
    if sf is not None:
        sf.write(str(output), audio, sample_rate)
        return
    write_wave_file(output, audio, sample_rate)


def write_wave_file(output: Path, waveform, sample_rate: int, channels: int | None = None) -> None:
    if hasattr(waveform, "ndim") and np is not None:
        waveform = np.asarray(waveform)
        if waveform.ndim == 1:
            waveform = waveform[:, None]
        waveform = np.clip(waveform, -1.0, 1.0)
        pcm = (waveform * 32767.0).astype(np.int16)
        channels = pcm.shape[1]
        frames = array("h", pcm.reshape(-1).tolist())
    else:
        frames = waveform
        if channels is None:
            channels = 1
    with wave.open(str(output), "wb") as wav:
        wav.setnchannels(channels)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(frames.tobytes())


if __name__ == "__main__":
    raise SystemExit(main())