from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import gradio as gr

from pyharp import *


import contextlib
import tempfile
from pathlib import Path

import soundfile as sf
import torch
from huggingface_hub import hf_hub_download

model_card = ModelCard(
    name="VoiceFixer",
    description=(
        "VoiceFixer aims to restore human speech regardless how serious its "
        "degraded. It can handle noise, reverberation, low resolution "
        "(2kHz~44.1kHz) and clipping (0.1-1.0 threshold) effect within one "
        "model. The model is fixed-rate and mono, so the result is resampled "
        "back to the input's rate and copied across its channels: the output "
        "has the format it came in with, but no stereo image. Repository: "
        "https://github.com/haoheliu/voicefixer Paper: "
        "https://arxiv.org/abs/2109.13731"
    ),
    author="Haohe Liu, Qiuqiang Kong, Qiao Tian, Yan Zhao, DeLiang Wang, Chuanzeng Huang, Yuxuan Wang",
    tags=[Subcategory.RESTORATION, SampleRate(44100), Channels(1), 'speech', 'domain:speech'],
)


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is fetched and built there alone. voicefixer downloads its checkpoints from
# Zenodo when it is imported, at about 1 MB/s, so the same files, mirrored on the Hub, are linked
# where it looks before downloading, ahead of the import. It picks CUDA itself where there is one.
if __name__ != "__main__":
    _VOICEFIXER_CACHE = Path.home() / ".cache" / "voicefixer"
    for _name, _path in (
        ("vf.ckpt", "analysis_module/checkpoints/vf.ckpt"),
        ("model.ckpt-1490000_trimed.pt", "synthesis_module/44100/model.ckpt-1490000_trimed.pt"),
    ):
        _link = _VOICEFIXER_CACHE / _path
        if not _link.exists():
            _link.parent.mkdir(parents=True, exist_ok=True)
            with contextlib.suppress(FileExistsError):
                _link.symlink_to(hf_hub_download(
                    "teamup-tech/VoiceFixer-weights", _name,
                    revision="11ca799c59df55bbbc438f223bf3b55a64ebced6",
                ))

    from voicefixer import VoiceFixer

    voicefixer_model = VoiceFixer()


def match_input_format(output_path: str, input_path: str) -> str:
    """
    Rewrite an output file at the sample rate and channel count it came in with.

    The models here are fixed-rate and mono, and hand the writing to vendored
    upstream code, so the format is corrected afterwards rather than threaded
    through inference. Duplicated channels are copies, not a recovered stereo
    image; the model card says so.

    Args:
        output_path (str): File the model wrote. Overwritten in place.
        input_path (str): File the request supplied.

    Returns:
        output_path (str): The same path, for use in a return statement.
    """

    import librosa
    import numpy as np
    import soundfile

    want = soundfile.info(input_path)
    have = soundfile.info(output_path)

    if have.samplerate == want.samplerate and have.channels == want.channels:
        return output_path

    data, rate = soundfile.read(output_path, always_2d=True)

    if rate != want.samplerate:
        data = librosa.resample(
            data.T, orig_sr=rate, target_sr=want.samplerate
        ).T

    if data.shape[1] == 1 and want.channels > 1:
        data = np.tile(data, (1, want.channels))
    elif data.shape[1] > 1 and want.channels == 1:
        data = data.mean(axis=1, keepdims=True)

    soundfile.write(output_path, data, want.samplerate)

    return output_path


# The recording is restored in 30-second segments, so the cost grows with its length
# PyHARP stops a job that runs past its timeout_s, 15 minutes by default, so a request whose
# reservation would exceed that is refused rather than started. A shorter one is left to
# ZeroGPU, which admits it when the caller's own allowance covers it.
MAX_GPU_SECONDS = 900


def gpu_seconds(input_audio, output_file, mode):
    return 15 + 0.1 * sf.info(input_audio).duration


@spaces.GPU(duration=gpu_seconds)
def restore(input_audio, output_file, mode):
    voicefixer_model.restore(input_audio, output_file, cuda=torch.cuda.is_available(),
                             mode=int(mode))


def process_fn(input_audio, mode):
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not input_audio:
        raise gr.Error("Provide the speech to restore.")
    try:
        duration = sf.info(input_audio).duration
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration == 0:
        raise gr.Error("The audio file is empty.")
    if gpu_seconds(input_audio, None, mode) > MAX_GPU_SECONDS:
        raise gr.Error(
            f"The recording is {duration:.0f} s long, more than one request can restore in the "
            f"{MAX_GPU_SECONDS // 60} minutes it may run. Trim it."
        )

    output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    restore(input_audio, output_file, mode)

    # The format is matched outside the GPU call, which it does not need
    output_file = match_input_format(output_file, input_audio)

    # A lossless input keeps its container. A lossy one comes back as WAV, since encoding the
    # result lossily again would expose the artifacts of the first encode and add a second set.
    input_ext = Path(input_audio).suffix.lower()
    if input_ext in {".flac"}:
        data, rate = sf.read(output_file, always_2d=True)
        lossless_file = str(Path(output_file).with_suffix(input_ext))
        sf.write(lossless_file, data, rate)
        output_file = lossless_file
    return output_file


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info("Upload an audio file to be processed by VoiceFixer."),
            gr.Dropdown(choices=["0", "1", "2"], value="0", label="Processing Mode", info="Select the VoiceFixer processing mode:\n0: Original Model (suggested by default)\n1: Add preprocessing module (remove higher frequency)\n2: Train mode (might work sometimes on seriously degraded real speech)"),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Fixed Audio").set_info("The enhanced audio output from VoiceFixer."),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

    demo.queue().launch(show_error=True, pwa=True, ssr_mode=False)