"""Gradio + WebRTC speech-to-speech chat for the German v6 voice model.

Loads the CPU-resident model from `peitho_model` (base LFM2.5-Audio-1.5B with
the German v6 fine-tune overlaid) and runs realtime generation inside an
`@spaces.GPU` slice on ZeroGPU.

Notes on prior experiments (kept here as a record):
- Self-transcribe pre-pass: tried a "Transkribiere wörtlich" system prompt to
  surface what the model heard. The model ignored the instruction and just
  chatted, so the transcript label was useless and got removed.
- Running the stock base model: produced robotic German audio and
  English/German token soup. The base model is English-only for S2S per its
  model card, so the German v6 fine-tune is required for German audio output.
"""
from __future__ import annotations

import os
import time
from queue import Queue
from threading import Thread
from typing import Any

# Disable moshi/Mimi torch.compile + CUDA graphs (incompatible with ZeroGPU
# torch patching) before liquid_audio is imported anywhere.
os.environ.setdefault("NO_TORCH_COMPILE", "1")
os.environ.setdefault("NO_CUDA_GRAPH", "1")

import spaces
import gradio as gr
import httpx
import numpy as np
import torch
from fastrtc import (
    AdditionalOutputs,
    ReplyOnPause,
    WebRTC,
)
from briefing import (
    FEWSHOT_EXAMPLES,
    build_system_prompt,
    build_user_grounding,
    format_demo_questions_html,
    format_hero_html,
    format_schedule_html,
)
from liquid_audio import ChatState
from peitho_model import lfm2_audio, mimi, proc

GPU_DURATION_SECONDS = 120

AUDIO_EOS_TOKEN = 2048
TURN_TTL_SECONDS = 600
TURN_FETCH_ATTEMPTS = 3
TURN_FETCH_DELAY_SECONDS = 2.0
CLOUDFLARE_TURN_URL_TEMPLATE = (
    "https://rtc.live.cloudflare.com/v1/turn/keys/{key_id}/credentials/generate-ice-servers"
)
STUN_FALLBACK_RTC: dict[str, Any] = {
    "iceServers": [
        {"urls": "stun:stun.l.google.com:19302"},
        {"urls": "stun:stun1.l.google.com:19302"},
    ]
}


def read_cloudflare_turn_secrets() -> tuple[str, str]:
    """Read Cloudflare TURN key id + api token from Space secrets (either name)."""
    key_id = (
        os.getenv("CLOUDFLARE_TURN_KEY_ID", "") or os.getenv("TURN_KEY_ID", "")
    ).strip()
    api_token = (
        os.getenv("CLOUDFLARE_TURN_KEY_API_TOKEN", "") or os.getenv("TURN_KEY_API_TOKEN", "")
    ).strip()
    return key_id, api_token


def normalize_rtc_configuration(payload: dict[str, Any]) -> dict[str, Any]:
    """Coerce Cloudflare's payload into a browser-friendly RTCPeerConnection config."""
    ice_servers = payload.get("iceServers")
    if ice_servers is None:
        raise ValueError(f"Unexpected TURN payload keys: {sorted(payload.keys())}")
    if isinstance(ice_servers, dict):
        ice_servers = [ice_servers]
    return {"iceServers": ice_servers}


def fetch_cloudflare_turn(key_id: str, api_token: str) -> dict[str, Any]:
    """Cloudflare Calls TURN via your own keys (rtc.live.cloudflare.com is up)."""
    response = httpx.post(
        CLOUDFLARE_TURN_URL_TEMPLATE.format(key_id=key_id),
        headers={
            "Authorization": f"Bearer {api_token}",
            "Content-Type": "application/json",
        },
        json={"ttl": TURN_TTL_SECONDS},
        timeout=30.0,
    )
    response.raise_for_status()
    return normalize_rtc_configuration(response.json())


def resolve_rtc_configuration() -> dict[str, Any]:
    """Lazy TURN fetch via your own Cloudflare keys; STUN fallback so boot never crashes.

    The free fastrtc/HF community TURN servers (turn.fastrtc.org,
    fastrtc-turn-server-login.hf.space) are currently down (fastrtc issue #429),
    so we use Cloudflare Calls directly with the Space secrets
    CLOUDFLARE_TURN_KEY_ID and CLOUDFLARE_TURN_KEY_API_TOKEN.
    """
    key_id, api_token = read_cloudflare_turn_secrets()
    if key_id and api_token:
        for attempt in range(TURN_FETCH_ATTEMPTS):
            try:
                config = fetch_cloudflare_turn(key_id, api_token)
                print("TURN credentials ready via Cloudflare.")
                return config
            except Exception as exc:
                print(f"TURN fetch attempt {attempt + 1}/{TURN_FETCH_ATTEMPTS}: {exc}")
                if attempt + 1 < TURN_FETCH_ATTEMPTS:
                    time.sleep(TURN_FETCH_DELAY_SECONDS)
    else:
        print(
            "Cloudflare TURN secrets missing. Set CLOUDFLARE_TURN_KEY_ID and "
            "CLOUDFLARE_TURN_KEY_API_TOKEN in Space settings for realtime voice."
        )
    print("Using STUN-only WebRTC (no TURN relay).")
    return STUN_FALLBACK_RTC


def ensure_cuda() -> None:
    """Move model + codec + processor to CUDA. Safe to call every GPU slice.

    ZeroGPU may run each `@spaces.GPU` call in a fresh process that starts from
    the CPU-loaded weights, so we check the live device and re-place if needed
    instead of caching a one-shot flag.
    """
    if next(lfm2_audio.parameters()).device.type == "cuda":
        return
    proc.to("cuda")
    mimi.to("cuda")
    lfm2_audio.to("cuda")


def build_grounded_chat(audio: tuple[int, np.ndarray]) -> ChatState:
    """Fresh, fully grounded chat for one spoken question (no cross-turn state).

    Built inside the GPU call so its tensors allocate on CUDA. We re-prime the
    system prompt + few-shot examples every utterance, which is exactly what we
    want for an independent schedule question.
    """
    rate, wav = audio
    chat = ChatState(proc)
    chat.new_turn("system")
    chat.add_text(build_system_prompt())
    chat.end_turn()
    for example_question, example_answer in FEWSHOT_EXAMPLES:
        chat.new_turn("user")
        chat.add_text(example_question)
        chat.end_turn()
        chat.new_turn("assistant")
        chat.add_text(example_answer)
        chat.end_turn()
    chat.new_turn("user")
    chat.add_text(build_user_grounding())
    chat.add_audio(torch.tensor(wav / 32_768, dtype=torch.float), rate)
    chat.end_turn()
    chat.new_turn("assistant")
    return chat


def chat_producer(
    q: "Queue[torch.Tensor | None]",
    chat: ChatState,
    temp: float | None,
    topk: int | None,
) -> None:
    print(f"Starting v6 generation with state {chat}.")
    with torch.no_grad(), mimi.streaming(1):
        for t in lfm2_audio.generate_interleaved(
            **chat,
            max_new_tokens=1024,
            audio_temperature=temp,
            audio_top_k=topk,
        ):
            q.put(t)
            if t.numel() > 1:
                if (t == AUDIO_EOS_TOKEN).any():
                    continue
                wav_chunk = mimi.decode(t[None, :, None])[0]
                q.put(wav_chunk)
    q.put(None)


@spaces.GPU(duration=GPU_DURATION_SECONDS)
def chat_response(
    audio: tuple[int, np.ndarray],
    _id: str,
    temp: float | None = 0.2,
    topk: int | None = 10,
):
    ensure_cuda()
    if temp == 0:
        temp = None
    if topk == 0:
        topk = None
    if temp is not None:
        temp = float(temp)
    if topk is not None:
        topk = int(topk)

    chat = build_grounded_chat(audio)

    q: "Queue[torch.Tensor | None]" = Queue()
    chat_thread = Thread(target=chat_producer, args=(q, chat, temp, topk))
    chat_thread.start()

    out_text: list[torch.Tensor] = []

    while True:
        t = q.get()
        if t is None:
            break
        if t.numel() == 1:
            out_text.append(t)
            cur_string = proc.text.decode(torch.cat(out_text)).removesuffix("<|text_end|>")
            yield AdditionalOutputs(cur_string)
        elif t.numel() == 8:
            continue
        elif t.numel() == 1920:
            np_chunk = (t.cpu().numpy() * 32_767).astype(np.int16)
            yield (24_000, np_chunk)
        else:
            raise RuntimeError(f"unexpected shape: {t.shape}")

    chat_thread.join()


def clear():
    gr.Info("Gespräch zurückgesetzt", duration=3)
    return ""


HEAD = (
    "<link rel='preconnect' href='https://fonts.googleapis.com'>"
    "<link rel='preconnect' href='https://fonts.gstatic.com' crossorigin>"
    "<link rel='stylesheet' href='https://fonts.googleapis.com/css2?"
    "family=Inter:wght@400;500;600;700&display=swap'>"
)

THEME = gr.themes.Soft(
    primary_hue=gr.themes.colors.blue,
    secondary_hue=gr.themes.colors.sky,
    neutral_hue=gr.themes.colors.slate,
    font=(gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"),
).set(
    body_background_fill="#eef1f5",
    body_background_fill_dark="#eef1f5",
    block_background_fill="#ffffff",
    block_border_width="1px",
    block_border_color="#dfe4ea",
    block_radius="12px",
    block_shadow="0 1px 2px rgba(16,24,40,.04)",
    button_primary_background_fill="#1c64f2",
    button_primary_background_fill_hover="#1a56db",
    button_primary_text_color="#ffffff",
    button_primary_border_color="#1c64f2",
    button_secondary_background_fill="#ffffff",
    button_secondary_background_fill_hover="#f4f6f9",
    button_secondary_text_color="#334155",
    button_secondary_border_color="#cbd5e1",
    button_large_radius="9px",
    button_small_radius="8px",
)

CSS = """
:root {--ink:#16202c; --muted:#64748b; --line:#dfe4ea; --blue:#1c64f2; --canvas:#eef1f5;}
footer, .show-api, .built-with {display:none !important;}
.gradio-container {max-width: 1080px !important; margin: 0 auto !important; padding: 0 16px 20px !important;}
body, .gradio-container {background: var(--canvas) !important;}
#app-shell {gap: 14px !important;}
.topbar {display:flex; align-items:center; justify-content:space-between;
  background:#ffffff; border:1px solid var(--line); border-radius:12px; padding:12px 16px;}
.brand {display:flex; align-items:center; gap:11px;}
.brand-logo {width:32px; height:32px; border-radius:8px; display:grid; place-items:center;
  background:var(--blue); color:#fff; font-size:20px; font-weight:700; line-height:1;}
.brand-text {display:flex; flex-direction:column; line-height:1.15;}
.brand-name {font-weight:700; color:var(--ink); font-size:16px;}
.brand-sub {color:var(--muted); font-size:12.5px; font-weight:500;}
.top-meta {color:var(--muted); font-size:13px; display:flex; align-items:center; gap:10px;}
.top-meta-tag {background:#e8f0fe; color:#1a56db; font-weight:600; font-size:11.5px;
  padding:3px 9px; border-radius:6px;}
.panel-title {color:var(--muted); font-size:12px; font-weight:700; letter-spacing:.04em;
  text-transform:uppercase; margin:2px 2px 8px;}
.panel {background:#fff; border:1px solid var(--line); border-radius:12px; overflow:hidden;}
.panel-head {display:flex; align-items:baseline; justify-content:space-between;
  padding:13px 16px; border-bottom:1px solid var(--line); font-weight:700; color:var(--ink); font-size:14.5px;}
.panel-meta {color:var(--muted); font-weight:500; font-size:12.5px;}
.sched {width:100%; border-collapse:collapse; font-size:14px;}
.sched th {text-align:left; color:var(--muted); font-weight:600; font-size:11.5px;
  text-transform:uppercase; letter-spacing:.03em; padding:9px 16px; background:#f7f9fc; border-bottom:1px solid var(--line);}
.sched td {padding:11px 16px; border-bottom:1px solid #eef1f5; color:var(--ink);}
.sched tbody tr:last-child td {border-bottom:none;}
.sched tbody tr:hover td {background:#f7f9fc;}
.sc-time {font-variant-numeric:tabular-nums; font-weight:700; color:var(--blue); width:64px;}
.sc-name {font-weight:600;}
.sc-age {color:var(--muted); width:54px;}
.sc-reason {color:#475569;}
.panel-note {padding:10px 16px; color:var(--muted); font-size:11.5px; background:#fafbfc; border-top:1px solid var(--line);}
.ask {margin-top:12px;}
.ask-label {color:var(--muted); font-size:11.5px; font-weight:700; letter-spacing:.04em;
  text-transform:uppercase; margin-bottom:8px;}
.chips {display:flex; flex-wrap:wrap; gap:8px;}
.chip {background:#fff; border:1px solid var(--line); color:#334155; padding:7px 12px;
  border-radius:8px; font-size:13px;}
#answer-box textarea {font-size:18px !important; line-height:1.5 !important; color:var(--ink) !important;
  background:#f7f9fc !important; border:1px solid var(--line) !important; border-radius:9px !important;
  min-height:96px !important;}
#answer-box textarea:focus {border-color:var(--blue) !important; box-shadow:0 0 0 3px rgba(28,100,242,.12) !important;}
#answer-box span[data-testid='block-info'], #answer-box label span {color:var(--muted) !important; font-weight:600 !important;}
.foot {color:var(--muted); font-size:12px; text-align:center; line-height:1.7; margin:8px 0 4px;}
.foot a {color:var(--blue); text-decoration:none;}
"""

with gr.Blocks(title="Praxis-Briefing", theme=THEME, css=CSS, head=HEAD) as demo:
    with gr.Column(elem_id="app-shell"):
        gr.HTML(format_hero_html())
        with gr.Row(equal_height=False):
            with gr.Column(scale=6):
                gr.HTML("<div class='panel-title'>Sprachabfrage</div>")
                webrtc = WebRTC(
                    modality="audio",
                    mode="send-receive",
                    full_screen=False,
                    rtc_configuration=resolve_rtc_configuration,
                    button_labels={"start": "Abfrage starten", "stop": "Stopp", "waiting": "…"},
                )
                text_out = gr.Textbox(
                    label="Antwort",
                    lines=3,
                    interactive=False,
                    elem_id="answer-box",
                    show_copy_button=True,
                )
                gr.HTML(format_demo_questions_html())
                clear_btn = gr.Button("Zurücksetzen", variant="secondary", size="sm")
            with gr.Column(scale=6):
                gr.HTML(format_schedule_html())
        gr.HTML(
            "<div class='foot'>LFM2.5-Audio-1.5B · auf Deutsch feinabgestimmt · "
            "Build-Small-Hackathon · <b>kein Medizinprodukt</b></div>"
        )

    webrtc.stream(
        ReplyOnPause(
            chat_response,  # type: ignore[arg-type]
            input_sample_rate=24_000,
            output_sample_rate=24_000,
            can_interrupt=False,
        ),
        inputs=[webrtc],
        outputs=[webrtc],
    )
    webrtc.on_additional_outputs(lambda s: s, outputs=[text_out])
    clear_btn.click(clear, outputs=[text_out])