import math
from pathlib import Path
import struct
import wave

import gradio as gr


DEFAULT_MAX_SECONDS = 300
GROUNDING = "Temporal grounding (when?)"
DETECTION = "Event detection (does it occur?)"


def create_fake_examples() -> list[list[str]]:
    """Create short synthetic clips so the UI examples work without model data."""
    example_dir = Path("/tmp/spotsound_fake_examples")
    example_dir.mkdir(parents=True, exist_ok=True)
    specs = [
        ("fake_hair_dryer.wav", "hair dryer drying", (900, 1250)),
        ("fake_ambulance_siren.wav", "ambulance siren", (520, 690)),
        ("fake_train_horn.wav", "train horning", (180, 240)),
        ("fake_police_siren.wav", "police car siren", (600, 850)),
    ]

    examples: list[list[str]] = []
    sample_rate = 16_000
    duration = 2
    for filename, query, frequencies in specs:
        path = example_dir / filename
        with wave.open(str(path), "wb") as wav:
            wav.setnchannels(1)
            wav.setsampwidth(2)
            wav.setframerate(sample_rate)
            frames = bytearray()
            for index in range(sample_rate * duration):
                time = index / sample_rate
                frequency = frequencies[int(time * 3) % len(frequencies)]
                signal = 0.16 * math.sin(2 * math.pi * frequency * time)
                signal += 0.04 * math.sin(2 * math.pi * frequency * 1.5 * time)
                frames.extend(struct.pack("<h", int(signal * 32767)))
            wav.writeframes(frames)
        examples.append([str(path), query])
    return examples


FAKE_EXAMPLES = create_fake_examples()


def spot(
    audio_path: str | None,
    query: str,
    task: str = GROUNDING,
    max_audio_seconds: int = DEFAULT_MAX_SECONDS,
    max_new_tokens: int = 128,
):
    """UI-only stand-in for the original SpotSound inference callback."""
    return None, "", None


CSS = """
#col-container {
  max-width: 1180px;
  margin: 0 auto;
  height: calc(100vh - 32px);
  overflow: visible;
}
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(
    title="SpotSound — fine-grained audio temporal grounding", fill_height=True
) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🔍🔊 SpotSound — fine-grained audio temporal grounding"
        )
        gr.Markdown(
            """
Find *when* a described sound happens inside a long recording. SpotSound is a LoRA on top of
**Audio Flamingo 3** that interleaves explicit timestamps into the audio stream, so the model can
answer with precise start/end times for short events buried in dense background noise.

[Paper](https://huggingface.co/papers/2604.13023) · [Project page](https://loiesun.github.io/spotsound/) ·
[Code](https://github.com/LoieSun/SpotSound) · [Model](https://huggingface.co/Loie/SpotSound) ·
[Benchmark](https://huggingface.co/datasets/Loie/SpotSound-Bench)
"""
        )
        with gr.Row():
            with gr.Column(scale=1):
                audio_in = gr.Audio(
                    label="Audio recording",
                    type="filepath",
                    sources=["upload", "microphone"],
                )
                query_in = gr.Textbox(
                    label="Sound to look for",
                    placeholder="e.g. dog barking, police car siren, hair dryer drying…",
                    lines=1,
                )
                run_btn = gr.Button("Spot it", variant="primary")
                with gr.Accordion("Advanced settings", open=False):
                    task_in = gr.Radio(
                        choices=[GROUNDING, DETECTION],
                        value=GROUNDING,
                        label="Task",
                    )
                    max_seconds_in = gr.Slider(
                        30,
                        600,
                        value=DEFAULT_MAX_SECONDS,
                        step=30,
                        label="Truncate audio to (seconds)",
                    )
                    max_tokens_in = gr.Slider(
                        16, 512, value=128, step=16, label="Max new tokens"
                    )
            with gr.Column(scale=1):
                plot_out = gr.Image(
                    label="Predicted temporal window(s)", type="pil"
                )
                answer_out = gr.Textbox(label="Model answer", lines=4)
                segment_out = gr.Audio(
                    label="Spotted segment(s) — full clip if nothing was found",
                    type="numpy",
                )

        gr.Examples(
            examples=FAKE_EXAMPLES,
            inputs=[audio_in, query_in],
            outputs=[plot_out, answer_out, segment_out],
            fn=spot,
            cache_examples=False,
            label="Fake examples (UI only)",
        )

    inputs = [audio_in, query_in, task_in, max_seconds_in, max_tokens_in]
    outputs = [plot_out, answer_out, segment_out]
    run_btn.click(fn=spot, inputs=inputs, outputs=outputs, api_name="spot")
    query_in.submit(fn=spot, inputs=inputs, outputs=outputs, api_name=False)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        theme=gr.themes.Citrus(),
        css=CSS,
        mcp_server=True,
    )