import tempfile
from math import ceil
from pathlib import Path

import gradio as gr
import spaces
import torch
import torch.nn.functional as F
import torchaudio
from flowhigh import FlowHighSR
from huggingface_hub import snapshot_download
from pyharp import ModelCard, build_endpoint
from pyharp.tags import Subcategory

SAMPLE_RATE = 48000
checkpoint_dir = snapshot_download(
    repo_id="ResembleAI/FlowHigh",
    revision="4281fe4119e5f3d209a9a893218fed85af2e5bbe",
    allow_patterns=["*.json", "*.pt"],
)
model = FlowHighSR.from_local(checkpoint_dir, device="cuda")

model_card = ModelCard(
    name="FLowHigh",
    description=(
        "Extend the bandwidth of low-resolution audio with flow matching. "
        "Repository: https://github.com/resemble-ai/flowhigh Paper: "
        "https://arxiv.org/abs/2501.04926"
    ),
    author="Jun-Hak Yun, Seung-Bin Kim, Seong-Whan Lee",
    tags=[Subcategory.BANDWIDTH_EXTENSION],
)


@spaces.GPU
@torch.inference_mode()
def process_fn(input_audio_path: str, output_sample_rate: str) -> str:
    if not input_audio_path:
        raise gr.Error("Choose an audio file.")
    audio, input_rate = torchaudio.load(input_audio_path)
    audio = audio.mean(dim=0, keepdim=True)
    if audio.shape[-1] == 0:
        raise gr.Error("The audio file is empty.")
    output_rate = int(output_sample_rate)
    output_length = ceil(audio.shape[-1] * output_rate / input_rate)
    native_length = ceil(audio.shape[-1] * SAMPLE_RATE / input_rate)

    # Upstream divides by the input peak, so silence bypasses inference.
    if audio.count_nonzero() == 0:
        result = torch.zeros(1, output_length)
    else:
        minimum_length = ceil(2048 * input_rate / SAMPLE_RATE)
        if audio.shape[-1] < minimum_length:
            audio = F.pad(audio, (0, minimum_length - audio.shape[-1]))
        result = model.generate(audio, input_rate, SAMPLE_RATE, timestep=1).cpu()
        result = result[..., :native_length]
        if output_rate != SAMPLE_RATE:
            result = torchaudio.functional.resample(result, SAMPLE_RATE, output_rate)
        result = result[..., :output_length]
        peak = result.abs().max()
        if peak > 1:
            result *= 0.99 / peak

    # mkdtemp, not a context manager: that would delete the directory on exit,
    # taking the written file with it before the caller can read it, which is
    # what the block-cache call here was working around. That also reached for
    # the output component, which exists only under __main__ and so is
    # undefined in the worker process that imports this module.
    directory = Path(tempfile.mkdtemp(prefix="flowhigh-"))
    output_path = directory / f"{directory.name}.wav"
    torchaudio.save(str(output_path), result, output_rate)

    return str(output_path)


if __name__ == "__main__":
    with gr.Blocks(title="FLowHigh", delete_cache=(600, 3600)) as demo:
        input_audio = (
            gr.Audio(type="filepath", label="Input Audio")
        .harp_required(True)
        .set_info("Low-bandwidth audio to extend. Up to 30 seconds.")
            .harp_required(True)
            .set_info("Audio up to 30 seconds. Stereo inputs are mixed to mono.")
        )
        output_rate = gr.Dropdown(
            choices=["24000", "48000", "96000"],
            value="48000",
            label="Output Sample Rate",
        ).set_info("The model runs at 48 kHz; 24 and 96 kHz outputs are resampled.")
        output_audio = gr.Audio(type="filepath", label="Enhanced Audio").set_info(
            "Mono WAV at the selected sample rate. Downloads are retained for one hour."
        )
        build_endpoint(
            model_card=model_card,
            input_components=[input_audio, output_rate],
            output_components=[output_audio],
            process_fn=process_fn,
        )


    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)