from __future__ import annotations

import subprocess
import sys
import tempfile
import threading
from pathlib import Path

MODEL_REPO = "stabilityai/stable-audio-3-small-music"
MODEL_REVISION = "0fef1392cd842149a2b6d445e181c97608faac06"
STABLE_AUDIO_TOOLS_REVISION = "3241adba4fc2a85cf5b29d9eb68d42f40a28e820"
OUTPUT_ROOT = Path(tempfile.gettempdir()) / "stable_audio_3_outputs"

_MODEL = None
_MODEL_CONFIG = None
_MODEL_LOCK = threading.Lock()


def _ensure_stable_audio_tools() -> None:
    try:
        import stable_audio_tools  # noqa: F401

        return
    except ImportError:
        pass

    subprocess.check_call(
        [
            sys.executable,
            "-m",
            "pip",
            "install",
            "--quiet",
            "--no-deps",
            (
                "git+https://github.com/Stability-AI/stable-audio-tools.git@"
                f"{STABLE_AUDIO_TOOLS_REVISION}"
            ),
        ]
    )


_ensure_stable_audio_tools()


def _load_model():
    global _MODEL, _MODEL_CONFIG
    if _MODEL is not None:
        return _MODEL, _MODEL_CONFIG

    with _MODEL_LOCK:
        if _MODEL is not None:
            return _MODEL, _MODEL_CONFIG

        import torch
        from stable_audio_tools.models import pretrained

        original_download = pretrained.hf_hub_download

        def pinned_download(repo_id, *args, **kwargs):
            if repo_id == MODEL_REPO:
                kwargs.setdefault("revision", MODEL_REVISION)
            return original_download(repo_id, *args, **kwargs)

        pretrained.hf_hub_download = pinned_download
        try:
            model, config = pretrained.get_pretrained_model(MODEL_REPO)
        finally:
            pretrained.hf_hub_download = original_download

        model = model.to("cuda").to(torch.float16)
        model.eval().requires_grad_(False)
        _MODEL = model
        _MODEL_CONFIG = config
        return _MODEL, _MODEL_CONFIG


def _load_audio(path: str):
    import torch
    import torchaudio

    audio, sample_rate = torchaudio.load(path)
    if audio.shape[0] > 2:
        audio = audio[:2]
    return int(sample_rate), audio.to(torch.float32)


def _save_audio(output, sample_rate: int) -> str:
    import soundfile as sf
    import torch

    output = output.permute(1, 0, 2).reshape(output.shape[1], -1)
    output = output.to(torch.float32)
    peak = output.abs().max().clamp(min=1e-9)
    output = output.div(peak).clamp(-1, 1).cpu().numpy().T

    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
    output_dir = Path(tempfile.mkdtemp(dir=OUTPUT_ROOT))
    output_path = output_dir / "stable_audio_3_edit.wav"
    sf.write(output_path, output, sample_rate, subtype="PCM_16")
    return str(output_path)


def edit_audio(
    audio_path: str,
    prompt: str,
    mode: str,
    edit_start: float,
    edit_end: float,
    continuation_length: float,
    strength: float,
    seed: int,
) -> str:
    import torch
    import torchaudio

    model, config = _load_model()
    from stable_audio_tools.inference.generation import (
        generate_diffusion_cond_inpaint,
    )
    sample_rate = int(config["sample_rate"])
    sample_size = int(config["sample_size"])
    source_rate, source = _load_audio(audio_path)
    source_duration = source.shape[-1] / source_rate

    if source_rate != sample_rate:
        source = torchaudio.functional.resample(source, source_rate, sample_rate)
    model_dtype = next(model.parameters()).dtype
    source_tuple = (sample_rate, source.to(model_dtype))

    if mode == "Continue":
        output_duration = source_duration + continuation_length
    else:
        output_duration = source_duration

    conditioning = [{"prompt": prompt, "seconds_total": output_duration}]
    kwargs = {
        "steps": 8,
        "cfg_scale": 1.0,
        "conditioning": conditioning,
        "sample_size": sample_size,
        "sampler_type": "pingpong",
        "seed": int(seed) if seed > 0 else -1,
        "device": "cuda",
        "sigma_max": 1.0,
        "apg_scale": 1.0,
        "duration_padding_sec": 6.0,
    }

    if mode == "Restyle":
        kwargs["init_audio"] = source_tuple
        kwargs["init_noise_level"] = float(strength)
    elif mode == "Inpaint":
        kwargs["inpaint_audio"] = source_tuple
        kwargs["inpaint_mask_start_seconds"] = float(edit_start)
        kwargs["inpaint_mask_end_seconds"] = float(edit_end)
    elif mode == "Continue":
        kwargs["inpaint_audio"] = source_tuple
        kwargs["inpaint_mask_start_seconds"] = float(source_duration)
        kwargs["inpaint_mask_end_seconds"] = float(output_duration)
    else:
        raise ValueError(f"Unknown edit mode: {mode}")

    with torch.inference_mode():
        output = generate_diffusion_cond_inpaint(model, **kwargs)

    output = output[..., : int(output_duration * sample_rate)]
    return _save_audio(output, sample_rate)