import spaces

# isort: split

from collections.abc import Iterator
from itertools import batched

import gradio as gr
import torch
from silero_vad import load_silero_vad

from src.environment import load_hf_token
from src.shared.transcription import transcribe
from src.space.runtime import SpaceConfig, load_space_components, read_recording


def create_app() -> gr.Blocks:
    # ---------------------------------------------------------
    # Place model weights on CUDA during startup for ZeroGPU sharing.
    # ---------------------------------------------------------
    config = SpaceConfig()
    components = load_space_components(load_hf_token())
    device = torch.device("cuda")
    components.model.to(device)
    vad_model = load_silero_vad(onnx=True)

    @spaces.GPU(duration=30)
    def generate(audio: list[bytes]) -> list[str]:
        results = transcribe(components.model, components.feature_extractor, components.tokenizer, config, device, audio)
        return [result.text for result in results]

    def transcribe_recording(audio_path: str | None) -> Iterator[str]:
        segments = read_recording(audio_path, vad_model)
        if not segments:
            yield "音声内に発話が見つかりませんでした。"
            return
        texts: list[str] = []
        for batch in batched(segments, 4):
            texts.extend(generate(list(batch)))
            yield "".join(texts).replace("\r", "").replace("\n", "")

    # ---------------------------------------------------------
    # Provide one audio input and one transcription action.
    # ---------------------------------------------------------
    with gr.Blocks(title="LambdaAudio v6", delete_cache=(3600, 3600)) as app:
        gr.Markdown("# LambdaAudio v6\n日本語の音声をアップロードするか、マイクで録音してください。\n\n[モデルの詳細](https://huggingface.co/KeisukeMiyamoto/lambda-audio-v6)")
        audio = gr.Audio(sources=["microphone", "upload"], type="filepath", format="wav", label="音声", editable=False, elem_id="recording-input")
        button = gr.Button("文字起こし", variant="primary", elem_id="transcribe-button")
        submit = gr.Button(visible="hidden", elem_id="transcribe-submit")
        output = gr.Textbox(label="文字起こし結果", lines=5, interactive=False, buttons=["copy"])
        submit.click(transcribe_recording, inputs=audio, outputs=output, api_name="transcribe", concurrency_limit=1)

        # ---------------------------------------------------------
        # Stop the browser recorder before submitting its uploaded audio.
        # The stop_recording event fires after Gradio finishes uploading.
        # ---------------------------------------------------------
        button.click(
            fn=None,
            queue=False,
            js="""() => {
                const audio = document.getElementById('recording-input');
                if (audio.dataset.submitAfterRecording) return;
                const stop = [...audio.querySelectorAll('.stop-button, .stop-button-paused')]
                    .find(button => button.checkVisibility());
                if (stop) {
                    audio.dataset.submitAfterRecording = 'true';
                    stop.click();
                    return;
                }
                document.getElementById('transcribe-submit').click();
            }""",
        )
        audio.stop_recording(
            fn=None,
            queue=False,
            js="""() => {
                const audio = document.getElementById('recording-input');
                if (!audio.dataset.submitAfterRecording) return;
                delete audio.dataset.submitAfterRecording;
                document.getElementById('transcribe-submit').click();
            }""",
        )
    return app.queue(max_size=10)


if __name__ == "__main__":
    create_app().launch(max_file_size="200mb")