import os
import base64
import gradio as gr
from pathlib import Path

from constants import STREAM_EVERY, APP_TMP_DIR, MODEL_ID, VAD_MODEL_ID
from ui import LED_DOT_OFF
from hf_dataset_utils import ALL_FILES

from stream_pipeline import (
    on_start_recording,
    on_stop_recording,
    shutdown_streamers,
    stream_step,
)
from offline_pipeline import (
    load_file_from_dataset,
    load_local_file,
    run_offline_pipeline,
)
from clean_up import purge_tmp_directory, cleanup_previous_run


_CSS_DIR = Path(__file__).resolve().parent / "assets"
# Embed the logo as a self-contained data URI so it never depends on an
# external CDN URL (the previous Mintlify link expired and broke the header).
_LOGO_B64 = base64.b64encode((_CSS_DIR / "logo.svg").read_bytes()).decode("ascii")
_LOGO_DATA_URI = f"data:image/svg+xml;base64,{_LOGO_B64}"

# Model identifiers for a subtle footer (short slugs straight from constants).
_MODEL_FOOTER = (
    '<div style="text-align:center; font-size:0.72rem; opacity:0.5; margin-top:0.75rem;">'
    f"Voice Focus: <code>{MODEL_ID}</code>"
    f" &nbsp;·&nbsp; VAD: <code>{VAD_MODEL_ID}</code>"
    "</div>"
)
with gr.Blocks() as demo:
    sample_stem = gr.State("")
    last_sample_stem = gr.State("")
    input_array = gr.State()
    streaming_sr = gr.State(None)
    current_sample_rate = gr.State(None)

    gr.HTML(
        '<a href="https://ai-coustics.com/" target="_blank">'
        f'<img src="{_LOGO_DATA_URI}" '
        'alt="ai-coustics Logo" width="400" style="display: block; margin: 0 auto;">'
        "</a>"
    )
    gr.Markdown(open("docs/intro.md", "r", encoding="utf-8").read())

    # ✅ Global controls (shared by both tabs)
    stt_model = gr.Radio(
        label="STT Model",
        choices=["Soniox STT-RT v3", "Deepgram Nova-3 RT"],
        value="Soniox STT-RT v3",
        interactive=True,
    )
    # Streamers are created on first start_recording to avoid Soniox "No audio received" timeout
    enhancement_level = gr.Slider(
        minimum=0,
        maximum=100,
        step=1,
        value=80,
        label="Enhancement level (%)",
        scale=2,
    )

    with gr.Tabs():
        with gr.Tab("Stream audio in real time") as stream_tab:
            with gr.Row(equal_height=False, elem_classes="stream-layout"):
                with gr.Column(scale=4, min_width=320):
                    with gr.Group(elem_classes="panel section-panel"):
                        gr.Markdown("### Input", elem_classes="title")
                        gr.Markdown(
                            open("docs/online.md", "r", encoding="utf-8").read(),
                            elem_classes="tab-description",
                        )

                        input_gain_db = gr.Slider(
                            minimum=0,
                            maximum=20,
                            step=0.5,
                            value=0,
                            label="Input gain (dB)",
                        )

                        audio_stream = gr.Audio(
                            sources=["microphone"],
                            streaming=True,
                            elem_id="audio_stream",
                        )

                with gr.Column(scale=6, min_width=420):
                    with gr.Group(elem_classes="status-panel output-panel"):
                        gr.Markdown("### Live Output", elem_classes="title")
                        with gr.Row(
                            equal_height=True, elem_classes="status-indicators"
                        ):
                            with gr.Group(elem_classes="status-card"):
                                gr.Markdown(
                                    "**System Status**",
                                    elem_classes="status-card__label",
                                )
                                system_status_led = gr.HTML(
                                    value=LED_DOT_OFF, show_label=False
                                )
                                system_status_text = gr.Markdown(
                                    value="Off",
                                    elem_classes="status-card__subtext",
                                )

                            with gr.Group(elem_classes="status-card"):
                                gr.Markdown(
                                    "**Voice Activity**",
                                    elem_classes="status-card__label",
                                )
                                vad_led = gr.HTML(value=LED_DOT_OFF, show_label=False)

                        with gr.Row(
                            equal_height=True,
                            elem_classes="transcript-row transcript-row--large",
                        ):
                            enhanced_text = gr.Textbox(
                                label="Enhanced Transcript",
                                lines=10,
                                autoscroll=False,
                            )
                            raw_text = gr.Textbox(
                                label="Raw Transcript",
                                lines=10,
                                autoscroll=False,
                            )

        with gr.Tab("Pick Example") as dataset_tab:
            with gr.Group(elem_classes="panel section-panel"):
                gr.Markdown("### Input", elem_classes="title")
                gr.Markdown(
                    open("docs/example_pick.md", "r", encoding="utf-8").read(),
                    elem_classes="tab-description",
                )

                dataset_dropdown = gr.Dropdown(
                    choices=ALL_FILES, value=ALL_FILES[0], label="Sample"
                )
                audio_file_from_dataset = gr.Audio(
                    label="Preview",
                    autoplay=False,
                    interactive=False,
                )

        with gr.Tab("Upload local file") as upload_tab:
            with gr.Group(elem_classes="panel section-panel"):
                gr.Markdown("### Input", elem_classes="title")
                gr.Markdown(
                    open("docs/local_file.md", "r", encoding="utf-8").read(),
                    elem_classes="tab-description",
                )

                audio_file_upload = gr.File(
                    file_types=[".wav", ".mp3", ".flac", ".m4a", ".ogg"],
                    file_count="single",
                    scale=3,
                )
                normalize = gr.Checkbox(label="Normalize audio", value=False)
                audio_preview = gr.Audio(
                    label="Preview",
                    autoplay=False,
                    interactive=False,
                )

    enhance_btn = gr.Button("Enhance with Quail Voice Focus", scale=2, visible=False)

    with gr.Group(elem_classes="panel results-card", visible=False) as results_card:
        result_title = gr.Markdown("", elem_classes="title")
        enhanced_audio = gr.Audio(
            type="numpy",
            interactive=False,
            buttons=["download"],
            autoplay=False,
        )

        with gr.Row(equal_height=True, elem_classes="results-row"):
            with gr.Column(scale=5, min_width=320, elem_classes="results-left"):
                noisy_image = gr.Image(
                    label="Input spectrogram", format="png", type="filepath"
                )
                enhanced_image = gr.Image(
                    label="Enhanced spectrogram", format="png", type="filepath"
                )

            with gr.Column(scale=5, min_width=320, elem_classes="results-right"):
                original_transcript = gr.Textbox(
                    label="Groundtruth transcription", lines=3, interactive=False
                )
                noisy_transcript = gr.Textbox(
                    label="Noisy transcription", lines=3, interactive=False
                )
                enhanced_transcript = gr.Textbox(
                    label="Transcription post enhancement", lines=3, interactive=False
                )

    gr.HTML(_MODEL_FOOTER)

    # ------------------------------------------------------
    # TAB CHANGES
    # ------------------------------------------------------
    stream_tab.select(
        lambda: (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(sources=["microphone"], streaming=True, interactive=True),
        ),
        outputs=[results_card, enhance_btn, audio_stream],
    )

    def _on_not_streaming_tab():
        shutdown_streamers()
        return (
            gr.update(streaming=False, interactive=False),
            gr.update(visible=True),
            LED_DOT_OFF,
            LED_DOT_OFF,
            "Off",
        )

    upload_tab.select(
        _on_not_streaming_tab,
        outputs=[
            audio_stream,
            enhance_btn,
            vad_led,
            system_status_led,
            system_status_text,
        ],
    ).then(
        load_local_file,
        inputs=[audio_file_upload, normalize],
        outputs=[input_array, sample_stem, audio_preview, current_sample_rate],
    )

    dataset_tab.select(
        _on_not_streaming_tab,
        outputs=[
            audio_stream,
            enhance_btn,
            vad_led,
            system_status_led,
            system_status_text,
        ],
    ).then(
        load_file_from_dataset,
        inputs=dataset_dropdown,
        outputs=[
            audio_file_from_dataset,
            input_array,
            sample_stem,
            current_sample_rate,
        ],
    )
    # ------------------------------------------------------
    # STREAMING EVENTS
    # ------------------------------------------------------
    audio_stream.stream(
        stream_step,
        inputs=[
            audio_stream,
            streaming_sr,
            stt_model,
            enhancement_level,
            input_gain_db,
        ],
        outputs=[
            streaming_sr,
            system_status_led,
            system_status_text,
            enhanced_text,
            raw_text,
            vad_led,
        ],
        stream_every=STREAM_EVERY,
        time_limit=60 * 2,
        concurrency_limit=1,
    )

    audio_stream.stop_recording(
        on_stop_recording,
        outputs=[vad_led, system_status_led, system_status_text, streaming_sr],
    ).then(
        shutdown_streamers,
    )

    audio_stream.start_recording(
        on_start_recording,
        outputs=[enhanced_text, raw_text, system_status_led, system_status_text],
    )

    # ------------------------------------------------------
    # OFFLINE EVENTS (DATASET + LOCAL FILE)
    # ------------------------------------------------------

    # Dataset dropdown selection triggers loading the audio file and hiding results until enhancement
    dataset_dropdown.change(
        lambda: gr.update(visible=False),
        outputs=results_card,
    ).then(
        load_file_from_dataset,
        inputs=dataset_dropdown,
        outputs=[
            audio_file_from_dataset,
            input_array,
            sample_stem,
            current_sample_rate,
        ],
    )

    # Uploading a local file triggers loading the audio file and hiding results until enhancement
    audio_file_upload.change(
        lambda: gr.update(visible=False),
        outputs=results_card,
    ).then(
        load_local_file,
        inputs=[audio_file_upload, normalize],
        outputs=[input_array, sample_stem, audio_preview, current_sample_rate],
    )

    normalize.change(
        load_local_file,
        inputs=[audio_file_upload, normalize],
        outputs=[input_array, sample_stem, audio_preview, current_sample_rate],
    )

    # Enhancement button: run pipeline with live transcript progress (dataset + local file modes).
    enhance_btn.click(cleanup_previous_run, inputs=[last_sample_stem]).then(
        lambda: gr.update(visible=True),
        outputs=results_card,
    ).then(
        run_offline_pipeline,
        inputs=[
            input_array,
            current_sample_rate,
            enhancement_level,
            stt_model,
            sample_stem,
        ],
        outputs=[
            enhanced_audio,
            enhanced_image,
            noisy_image,
            original_transcript,
            noisy_transcript,
            enhanced_transcript,
            last_sample_stem,
        ],
    ).failure(
        lambda: gr.Warning(
            "Enhancement failed. Please refresh page and make sure you have a stable connection."
        )
    )

os.makedirs(APP_TMP_DIR, exist_ok=True)
purge_tmp_directory(max_age_minutes=0, tmp_dir=APP_TMP_DIR)
demo.queue()
demo.launch(
    css=(_CSS_DIR / "styling.css").read_text(encoding="utf-8"),
    allowed_paths=[APP_TMP_DIR],
)