"""compare-codec — upload audio, pick codecs, hear reconstructions side by side."""

from __future__ import annotations

import time
from pathlib import Path

import gradio as gr
import numpy as np

from compare_codec import CodecConfig, get_all
from compare_codec.spectrogram import generate as make_spectrogram

MAX_DURATION_S = 30.0


def _encode_decode_one(
    audio_path: Path, codec_name: str, cfg: CodecConfig
) -> tuple[np.ndarray, int, float]:
    """Run one codec config and return (audio_array, sample_rate, elapsed_seconds)."""
    codec = get_all()[codec_name]
    sr = cfg.params.get("sample_rate", codec.sample_rate)

    t0 = time.perf_counter()
    audio_out = codec.encode_decode(audio_path, cfg)
    elapsed = time.perf_counter() - t0

    max_samples = int(MAX_DURATION_S * sr)
    if len(audio_out) > max_samples:
        audio_out = audio_out[:max_samples]

    return audio_out, sr, elapsed


def build_ui() -> gr.Blocks:
    codecs = get_all()

    with gr.Blocks(title="compare-codec") as demo:
        gr.Markdown("# compare-codec")

        audio_in = gr.Audio(
            sources=["upload", "microphone"],
            type="filepath",
            label="Input audio",
        )
        run_btn = gr.Button("Reconstruct", variant="primary")

        tab_components: dict[str, dict] = {}

        with gr.Tabs():
            for codec_name, codec in codecs.items():
                configs = codec.configs()
                config_labels = [c.name for c in configs]

                with gr.Tab(label=codec_name) as tab:
                    config_dd = gr.Dropdown(
                        choices=config_labels,
                        value=config_labels[0],
                        label="Configuration",
                    )
                    audio_out = gr.Audio(
                        label="Reconstructed",
                        type="numpy",
                        interactive=False,
                    )
                    stats_md = gr.Markdown(value="*Upload audio to compare.*")
                    spec_img = gr.Image(
                        label="Spectrogram",
                        type="filepath",
                        interactive=False,
                    )

                tab_components[codec_name] = {
                    "tab": tab,
                    "config_dd": config_dd,
                    "audio_out": audio_out,
                    "stats_md": stats_md,
                    "spec_img": spec_img,
                    "configs": configs,
                    "config_labels": config_labels,
                }

        active_tab = gr.State(value=list(codecs.keys())[0])
        ordered_names = list(codecs.keys())

        all_outputs = []
        for name in ordered_names:
            c = tab_components[name]
            all_outputs.extend([c["audio_out"], c["stats_md"], c["spec_img"]])

        def process_active(audio_path: str | None, current_tab: str, *dropdown_values):
            """Reconstruct only the active tab's codec."""
            if audio_path is None:
                return [gr.update()] * len(all_outputs)

            dd_map = dict(zip(ordered_names, dropdown_values))
            comps = tab_components[current_tab]
            cfg_label = dd_map[current_tab]
            cfg = next(c for c in comps["configs"] if c.name == cfg_label)

            audio_out, sr, elapsed = _encode_decode_one(
                Path(audio_path), current_tab, cfg
            )
            spec_path = make_spectrogram(audio_out, sr)

            stats_text = (
                f"**{elapsed:.2f}s**  |  "
                f"{sr / 1000:.0f} kHz  |  "
                f"{cfg_label}"
            )

            flat = []
            for name in ordered_names:
                if name == current_tab:
                    flat.extend(
                        [
                            gr.update(value=(sr, audio_out)),
                            gr.update(value=stats_text),
                            gr.update(value=str(spec_path)),
                        ]
                    )
                else:
                    flat.extend([gr.update(), gr.update(), gr.update()])
            return flat

        all_dropdowns = [tab_components[n]["config_dd"] for n in ordered_names]

        btn_event = run_btn.click(
            fn=process_active,
            inputs=[audio_in, active_tab] + all_dropdowns,
            outputs=all_outputs,
        )

        for codec_name, comps in tab_components.items():
            comps["tab"].select(
                fn=lambda name=codec_name: name,
                inputs=[],
                outputs=[active_tab],
            )

    return demo


demo = build_ui()

if __name__ == "__main__":
    demo.launch()