from collections.abc import AsyncGenerator
from pathlib import Path
import platform

import gradio as gr
import httpx
from httpx_sse import aconnect_sse
from openai import AsyncOpenAI

from speaches import kokoro_utils
from speaches.api_types import Voice
from speaches.config import Config, Task
from speaches.routers.speech import (
    MAX_SAMPLE_RATE,
    MIN_SAMPLE_RATE,
    SUPPORTED_RESPONSE_FORMATS,
)

TRANSCRIPTION_ENDPOINT = "/v1/audio/transcriptions"
TRANSLATION_ENDPOINT = "/v1/audio/translations"
TIMEOUT_SECONDS = 180
TIMEOUT = httpx.Timeout(timeout=TIMEOUT_SECONDS)
DEFAULT_TEXT = "A rainbow is an optical phenomenon caused by refraction, internal reflection and dispersion of light in water droplets resulting in a continuous spectrum of light appearing in the sky."  # noqa: E501

# NOTE: `gr.Request` seems to be passed in as the last positional (not keyword) argument


def base_url_from_gradio_req(request: gr.Request) -> str:
    # NOTE: `request.request.url` seems to always have a path of "/gradio_api/queue/join"
    assert request.request is not None
    return f"{request.request.url.scheme}://{request.request.url.netloc}"


def http_client_from_gradio_req(request: gr.Request, config: Config) -> httpx.AsyncClient:
    base_url = base_url_from_gradio_req(request)
    return httpx.AsyncClient(
        base_url=base_url,
        timeout=TIMEOUT,
        headers={"Authorization": f"Bearer {config.api_key}"} if config.api_key else None,
    )


def openai_client_from_gradio_req(request: gr.Request, config: Config) -> AsyncOpenAI:
    base_url = base_url_from_gradio_req(request)
    return AsyncOpenAI(base_url=f"{base_url}/v1", api_key=config.api_key if config.api_key else "cant-be-empty")


def create_gradio_demo(config: Config) -> gr.Blocks:  # noqa: C901, PLR0915
    async def whisper_handler(
        file_path: str, model: str, task: Task, temperature: float, stream: bool, request: gr.Request
    ) -> AsyncGenerator[str, None]:
        http_client = http_client_from_gradio_req(request, config)
        if task == Task.TRANSCRIBE:
            endpoint = TRANSCRIPTION_ENDPOINT
        elif task == Task.TRANSLATE:
            endpoint = TRANSLATION_ENDPOINT

        if stream:
            previous_transcription = ""
            async for transcription in streaming_audio_task(http_client, file_path, endpoint, temperature, model):
                previous_transcription += transcription
                yield previous_transcription
        else:
            yield await audio_task(http_client, file_path, endpoint, temperature, model)

    async def audio_task(
        http_client: httpx.AsyncClient, file_path: str, endpoint: str, temperature: float, model: str
    ) -> str:
        with Path(file_path).open("rb") as file:  # noqa: ASYNC230
            response = await http_client.post(
                endpoint,
                files={"file": file},
                data={
                    "model": model,
                    "response_format": "text",
                    "temperature": temperature,
                },
            )

        response.raise_for_status()
        return response.text

    async def streaming_audio_task(
        http_client: httpx.AsyncClient, file_path: str, endpoint: str, temperature: float, model: str
    ) -> AsyncGenerator[str, None]:
        with Path(file_path).open("rb") as file:  # noqa: ASYNC230
            kwargs = {
                "files": {"file": file},
                "data": {
                    "response_format": "text",
                    "temperature": temperature,
                    "model": model,
                    "stream": True,
                },
            }
            async with aconnect_sse(http_client, "POST", endpoint, **kwargs) as event_source:
                async for event in event_source.aiter_sse():
                    yield event.data

    async def update_whisper_model_dropdown(request: gr.Request) -> gr.Dropdown:
        openai_client = openai_client_from_gradio_req(request, config)
        models = (await openai_client.models.list()).data
        model_names: list[str] = [model.id for model in models]
        assert config.whisper.model in model_names
        recommended_models = {model for model in model_names if model.startswith("Systran")}
        other_models = [model for model in model_names if model not in recommended_models]
        model_names = list(recommended_models) + other_models
        return gr.Dropdown(
            choices=model_names,
            label="Model",
            value=config.whisper.model,
        )

    async def update_voices_and_language_dropdown(model_id: str | None, request: gr.Request) -> dict:
        params = httpx.QueryParams({"model_id": model_id})
        http_client = http_client_from_gradio_req(request, config)
        res = (await http_client.get("/v1/audio/speech/voices", params=params)).raise_for_status()
        voice_ids = [Voice.model_validate(x).voice_id for x in res.json()]
        return {
            voice_dropdown: gr.update(choices=voice_ids, value=voice_ids[0]),
            language_dropdown: gr.update(visible=model_id == "hexgrad/Kokoro-82M"),
        }

    async def handle_audio_speech(
        text: str,
        model: str,
        voice: str,
        language: str | None,
        response_format: str,
        speed: float,
        sample_rate: int | None,
        request: gr.Request,
    ) -> Path:
        openai_client = openai_client_from_gradio_req(request, config)
        res = await openai_client.audio.speech.create(
            input=text,
            model=model,
            voice=voice,  # pyright: ignore[reportArgumentType]
            response_format=response_format,  # pyright: ignore[reportArgumentType]
            speed=speed,
            extra_body={"language": language, "sample_rate": sample_rate},
        )
        audio_bytes = res.response.read()
        file_path = Path(f"audio.{response_format}")
        with file_path.open("wb") as file:  # noqa: ASYNC230
            file.write(audio_bytes)
        return file_path

    with gr.Blocks(title="Speaches Playground") as demo:
        gr.Markdown("# Speaches Playground")
        gr.Markdown(
            "### Consider supporting the project by starring the [speaches-ai/speaches repository on GitHub](https://github.com/speaches-ai/speaches)."
        )
        gr.Markdown("### Documentation Website: https://speaches-ai.github.io/speaches")
        gr.Markdown(
            "### For additional details regarding the parameters, see the [API Documentation](https://speaches-ai.github.io/speaches/api)"
        )

        with gr.Tab(label="Speech-to-Text"):
            audio = gr.Audio(type="filepath")
            whisper_model_dropdown = gr.Dropdown(
                choices=[config.whisper.model],
                label="Model",
                value=config.whisper.model,
            )
            task_dropdown = gr.Dropdown(
                choices=[task.value for task in Task],
                label="Task",
                value=Task.TRANSCRIBE,
            )
            temperature_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, label="Temperature", value=0.0)
            stream_checkbox = gr.Checkbox(label="Stream", value=True)
            button = gr.Button("Generate")

            output = gr.Textbox()

            # NOTE: the inputs order must match the `whisper_handler` signature
            button.click(
                whisper_handler,
                [audio, whisper_model_dropdown, task_dropdown, temperature_slider, stream_checkbox],
                output,
            )

        with gr.Tab(label="Text-to-Speech"):
            model_dropdown_choices = ["hexgrad/Kokoro-82M", "rhasspy/piper-voices"]
            if platform.machine() != "x86_64":
                model_dropdown_choices.remove("rhasspy/piper-voices")
                gr.Textbox("Speech generation using `rhasspy/piper-voices` model is only supported on x86_64 machines.")

            text = gr.Textbox(
                label="Input Text",
                value=DEFAULT_TEXT,
            )
            stt_model_dropdown = gr.Dropdown(
                choices=model_dropdown_choices,
                label="Model",
                value="hexgrad/Kokoro-82M",
            )
            voice_dropdown = gr.Dropdown(
                choices=["af"],
                label="Voice",
                value="af",
            )
            language_dropdown = gr.Dropdown(
                choices=kokoro_utils.LANGUAGES, label="Language", value="en-us", visible=True
            )
            stt_model_dropdown.change(
                update_voices_and_language_dropdown,
                inputs=[stt_model_dropdown],
                outputs=[voice_dropdown, language_dropdown],
            )
            response_fromat_dropdown = gr.Dropdown(
                choices=SUPPORTED_RESPONSE_FORMATS,
                label="Response Format",
                value="wav",
            )
            speed_slider = gr.Slider(minimum=0.25, maximum=4.0, step=0.05, label="Speed", value=1.0)
            sample_rate_slider = gr.Number(
                minimum=MIN_SAMPLE_RATE,
                maximum=MAX_SAMPLE_RATE,
                label="Desired Sample Rate",
                info="""
Setting this will resample the generated audio to the desired sample rate.
You may want to set this if you are going to use 'rhasspy/piper-voices' with voices of different qualities but want to keep the same sample rate.
Default: None (No resampling)
""",  # noqa: E501
                value=lambda: None,
            )
            button = gr.Button("Generate Speech")
            output = gr.Audio(type="filepath")
            button.click(
                handle_audio_speech,
                [
                    text,
                    stt_model_dropdown,
                    voice_dropdown,
                    language_dropdown,
                    response_fromat_dropdown,
                    speed_slider,
                    sample_rate_slider,
                ],
                output,
            )

        demo.load(update_whisper_model_dropdown, inputs=None, outputs=whisper_model_dropdown)
        demo.load(
            update_voices_and_language_dropdown,
            inputs=[stt_model_dropdown],
            outputs=[voice_dropdown, language_dropdown],
        )
    return demo