import argparse
import os

import soundfile as sf
import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import snapshot_download

from auffusion.auffusion_pipeline import Generator, denormalize_spectrogram


MODEL_ID = "dinhhung1508/SwiftAudio"
SAMPLE_RATE = 16000


def load_model(model_id=MODEL_ID, device=None):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.float16 if device == "cuda" else torch.float32
    model_path = model_id if os.path.isdir(model_id) else snapshot_download(model_id)

    pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
    pipeline = pipeline.to(device)
    pipeline.set_progress_bar_config(disable=True)

    vocoder = Generator.from_pretrained(model_path, subfolder="vocoder")
    vocoder = vocoder.to(device=device, dtype=dtype)
    return pipeline, vocoder, device, dtype


@torch.inference_mode()
def generate_audio(pipeline, vocoder, prompt, seed=42, device=None, dtype=None):
    device = device or pipeline.device.type
    dtype = dtype or pipeline.unet.dtype

    text_inputs = pipeline.tokenizer(
        prompt,
        padding="max_length",
        max_length=pipeline.tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    )
    prompt_embeds = pipeline.text_encoder(text_inputs.input_ids.to(device))[0]
    prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)

    generator = torch.Generator(device=device).manual_seed(seed)
    noise = torch.randn(
        (1, pipeline.unet.config.in_channels, 256 // 8, 1024 // 8),
        generator=generator,
        device=device,
        dtype=dtype,
    )

    timestep = torch.full(
        (1,),
        pipeline.scheduler.config.num_train_timesteps - 1,
        device=device,
        dtype=torch.int64,
    )
    model_pred = pipeline.unet(
        noise, timestep, prompt_embeds, return_dict=False
    )[0]

    alpha_t = 0.0047**0.5
    sigma_t = (1 - 0.0047) ** 0.5
    latents = (noise - sigma_t * model_pred) / alpha_t
    latents = latents / pipeline.vae.config.scaling_factor

    image = pipeline.vae.decode(latents, return_dict=False)[0]
    image = pipeline.image_processor.postprocess(
        image, output_type="pt", do_denormalize=[True]
    )[0]
    spectrogram = denormalize_spectrogram(image)
    return vocoder.inference(spectrogram, lengths=10 * SAMPLE_RATE)[0]


def parse_args():
    parser = argparse.ArgumentParser(description="Generate audio with SwiftAudio.")
    parser.add_argument("--prompt", required=True, help="Text description of the audio.")
    parser.add_argument("--output", default="output.wav", help="Output WAV path.")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--model", default=MODEL_ID, help="Hub model ID or local path.")
    parser.add_argument("--device", choices=["cuda", "cpu"], default=None)
    return parser.parse_args()


def main():
    args = parse_args()
    pipeline, vocoder, device, dtype = load_model(args.model, args.device)
    audio = generate_audio(
        pipeline, vocoder, args.prompt, args.seed, device=device, dtype=dtype
    )
    sf.write(args.output, audio, SAMPLE_RATE)
    print(f"Saved audio to {args.output}")


if __name__ == "__main__":
    main()