from __future__ import annotations

import tempfile
import uuid
from pathlib import Path

import spaces
import gradio as gr
import torch
import torchaudio
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint
from pyharp.tags import Subcategory

from data.tokenizer import AudioTokenizer, TextTokenizer
from inference_speech_editing_scale import inference_one_sample
from models.voicecraft import VoiceCraft


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CODEC_RATE = 50
OUTPUT_DIR = Path(tempfile.gettempdir()) / "voicecraft_outputs"
OUTPUT_DIR.mkdir(exist_ok=True)

model = VoiceCraft.from_pretrained("pyp1/VoiceCraft_830M_TTSEnhanced").to(DEVICE).eval()
codec_path = hf_hub_download("pyp1/VoiceCraft", "encodec_4cb2048_giga.th")
audio_tokenizer = AudioTokenizer(device=DEVICE, signature=codec_path)
text_tokenizer = TextTokenizer(backend="espeak")


@spaces.GPU(duration=60)
def edit_speech(audio_path: str, target_text: str, start_time: float, end_time: float) -> str:
    if not audio_path:
        raise gr.Error("Upload a speech recording.")
    target_text = target_text.strip()
    if not target_text:
        raise gr.Error("Enter the full transcript with the replacement text.")

    info = torchaudio.info(audio_path)
    duration = info.num_frames / info.sample_rate
    if not 0 <= start_time < end_time <= duration:
        raise gr.Error(f"Select an edit interval within the {duration:.2f} second recording.")

    start = max(start_time - 0.08, 1 / CODEC_RATE)
    end = min(end_time + 0.08, duration)
    mask = torch.tensor([[round(start * CODEC_RATE), round(end * CODEC_RATE)]])
    decode_config = {
        "top_k": 0,
        "top_p": 0.9,
        "temperature": 1.0,
        "stop_repetition": 3,
        "kvcache": 1,
        "codec_sr": CODEC_RATE,
        "silence_tokens": [1388, 1898, 131],
    }
    generated = inference_one_sample(
        model,
        model.args,
        model.args.phn2num,
        text_tokenizer,
        audio_tokenizer,
        audio_path,
        target_text,
        mask,
        DEVICE,
        decode_config,
    )
    output_path = OUTPUT_DIR / f"{uuid.uuid4().hex}.wav"
    torchaudio.save(str(output_path), generated[0].cpu(), 16000)
    return str(output_path)


if __name__ == "__main__":
    with gr.Blocks(title="VoiceCraft") as demo:
        build_endpoint(
            model_card=ModelCard(
                name="VoiceCraft",
                description=(
                    "Edit English speech while retaining the surrounding recording. "
                    "Repository: https://github.com/jasonppy/VoiceCraft Paper: "
                    "https://arxiv.org/abs/2403.16973"
                ),
                author="Puyuan Peng, Po-Yao Huang, Abdelrahman Mohamed, David Harwath",
                tags=[Subcategory.EDITING, 'speech-editing'],
            ),
            input_components=[
                gr.Audio(type="filepath", label="Speech Recording")
                .harp_required(True)
                .set_info("English speech to edit."),
                gr.Textbox(label="Target Transcript", lines=4)
                .harp_required(True)
                .set_info("Full transcript of the recording, with the replacement words in place."),
                gr.Number(label="Edit Start (seconds)", value=0.5, precision=2)
                .set_info("Start of the words to replace."),
                gr.Number(label="Edit End (seconds)", value=1.5, precision=2)
                .set_info("End of the words to replace."),
            ],
            output_components=[
                gr.Audio(type="filepath", label="Edited Speech")
                .set_info("16 kHz WAV recording."),
            ],
            process_fn=edit_speech,
        )


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