from __future__ import annotations

import json
import uuid
from pathlib import Path
from tempfile import gettempdir

import gradio as gr
import soundfile as sf

try:
    import spaces
except ImportError:
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                return args[0]

from pyharp import ModelCard, build_endpoint

from deepafx_st_runtime import style_transfer

MIN_AUDIO_SECONDS = 1
MAX_AUDIO_SECONDS = 30
OUTPUT_ROOT = Path(gettempdir()) / "deepafx_st_outputs"

model_card = ModelCard(
    name="DeepAFx-ST",
    description=(
        "Transfer the production style (EQ, compression) "
        "of a reference recording onto your audio."
    ),
    author="Adobe Research",
    tags=[
        "audio-effects",
        "style-transfer",
        "mixing",
        "mastering",
        "production",
    ],
)


def _validate_audio(path: str | None, label: str) -> str:
    if not path:
        raise gr.Error(f"Please upload {label}.")
    try:
        duration = sf.info(path).duration
    except Exception as exc:
        raise gr.Error(f"Could not read {label}: {exc}") from exc
    if duration < MIN_AUDIO_SECONDS:
        raise gr.Error(
            f"{label} must be at least {MIN_AUDIO_SECONDS} second long."
        )
    if duration > MAX_AUDIO_SECONDS:
        raise gr.Error(
            f"{label} must be no longer than {MAX_AUDIO_SECONDS} seconds. "
            f"Received {duration:.1f} seconds."
        )
    return path


@spaces.GPU(duration=120)
def process_fn(
    input_path: str | None,
    reference_path: str | None,
) -> tuple[str, str]:
    input_path = _validate_audio(input_path, "input audio")
    reference_path = _validate_audio(reference_path, "reference audio")

    try:
        audio, sample_rate, params = style_transfer(input_path, reference_path)
    except Exception as exc:
        raise gr.Error(f"DeepAFx-ST inference failed: {exc}") from exc

    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_dir.mkdir(parents=True, exist_ok=True)

    audio_path = output_dir / "styled_output.wav"
    sf.write(str(audio_path), audio.numpy(), sample_rate)

    params_path = output_dir / "dsp_parameters.json"
    params_path.write_text(
        json.dumps(
            {"model": "DeepAFx-ST", "variant": "autodiff", **params},
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )

    return str(audio_path), str(params_path)


with gr.Blocks(title="DeepAFx-ST") as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio")
        .harp_required(True)
        .set_info("Audio to process, 1 to 30 seconds."),
        gr.Audio(type="filepath", label="Reference Audio")
        .harp_required(True)
        .set_info("Target style — your input will be EQ'd and compressed to match this recording."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Styled Output")
        .set_info("Processed audio with the reference's production style."),
        gr.File(
            type="filepath",
            file_types=[".json"],
            label="DSP Parameters",
        ).set_info("Predicted EQ and compressor parameters."),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)