from __future__ import annotations

import gradio as gr

from config import CONFIG
from assistant import bootstrap, chat, stt, tts, ui

UI = CONFIG.pack["ui"]
PROFILE_EXAMPLE = CONFIG.pack["profile_example"]


with gr.Blocks(title="Amigo") as demo:
    state = gr.State([])

    gr.HTML(
        f'<div id="amigo-head"><h1>{UI["title"]}</h1>'
        f'<p>{UI["subtitle"]}</p></div>'
    )
    avatar = gr.HTML(ui.avatar("idle"))

    # The one action: a big tap-to-talk button right under the avatar. CSS turns
    # Gradio's audio widget into a single round button (see assistant/ui.py).
    mic = gr.Audio(
        sources=["microphone"], type="filepath",
        label=UI["mic"], elem_id="amigo-mic",
    )
    # Non-streaming: one clip per turn. A streaming=True output only plays on
    # the first turn (Gradio reuses the component's MediaSource and never
    # reinitializes it), so each turn sets a fresh full-reply clip instead.
    reply_audio = gr.Audio(
        label=UI["reply"], autoplay=True,
        interactive=False, elem_id="amigo-reply",
    )

    # The transcript, as a reading aid. A placeholder fills the empty space.
    chatbot = gr.Chatbot(
        height=280, label="", elem_id="amigo-chat", show_label=False,
        placeholder=UI["chat_placeholder"],
    )

    # Setup lives at the bottom, collapsed: the editable profile, prefilled with
    # a neutral example. Nothing is predefined; it feeds every turn's prompt.
    with gr.Accordion(UI["profile_label"], open=False):
        gr.Markdown(UI["profile_help"])
        profile_box = gr.Code(
            value=PROFILE_EXAMPLE, language="yaml", label="",
            elem_id="amigo-profile",
        )

    def on_listen():
        """Mic started: show the listening orb."""
        return ui.avatar("listening")

    def handle_turn(audio_path, history, profile_text):
        """Transcribe, stream the text live, then speak the full reply.

        Text streams token-by-token so he sees the answer forming; the voice
        plays once the reply is complete, which is a few seconds for a short
        answer. `profile_text` is the YAML from the editor above. Dict-form
        yields touch only the named components.
        """
        history = history or []
        yield {avatar: ui.avatar("thinking")}

        user_text = stt.transcribe(audio_path)
        if not user_text:
            yield {avatar: ui.avatar("idle", UI["not_heard"])}
            return

        history = history + [{"role": "user", "content": user_text}]
        history.append({"role": "assistant", "content": ""})

        full = ""
        for chunk in chat.respond(user_text, history[:-2], profile_text):
            full += chunk
            history[-1]["content"] = full
            yield {chatbot: history, avatar: ui.avatar("talking")}

        sr, samples = tts.synth(tts.clean_text(full))
        yield {chatbot: history, reply_audio: (sr, samples),
               avatar: ui.avatar("idle")}

    mic.start_recording(on_listen, outputs=[avatar])
    (
        mic.stop_recording(
            handle_turn,
            inputs=[mic, state, profile_box],
            outputs=[chatbot, reply_audio, avatar],
        )
        .then(lambda h: h, inputs=chatbot, outputs=state)
        # Clear the mic so it drops its recorded clip and shows the big record
        # button again. Without this the audio widget stays in playback/edit
        # mode on the last take, and he can't start a new turn. Runs on every
        # path (including "not heard"), so a failed turn never locks him out.
        .then(lambda: gr.update(value=None), outputs=mic)
    )


if __name__ == "__main__":
    # On a Space this downloads the models on first boot; locally it's a no-op.
    bootstrap.ensure_models()
    # Warm the model + embeddings so the first turn isn't cold.
    print("Calentando modelos…")
    chat.llm.warmup()
    chat.memory.warmup()
    print("Listo.")
    # Gradio 6: theme + css belong on launch().
    demo.launch(theme=ui.THEME, css=ui.CSS)