"""
PlotWeaver Hausa Voice AI Agent — HuggingFace Spaces
======================================================
Verified against gradio 5.49.1. Key constraints this file respects:

  * Models load LAZILY. The Space must show UI within seconds, not after a
    3.5GB download. Nothing heavy is imported or loaded at module scope.
  * CPU-SAFE DEFAULTS. whisper-large-v3 needs ~6GB and takes ~60s per utterance
    on the 2-vCPU free tier — unusable. Default is whisper-small; large-v3 is
    one env var away when you attach a GPU.
  * gr.skip() on stream ticks that produce no audio, otherwise the player
    restarts on every 0.5s tick and the agent's reply stutters.
  * Streaming partials OFF by default on CPU: an extra decode every 900ms
    saturates the box and makes the demo feel worse, not better.
"""

import os
import sys
import logging
from datetime import datetime

import gradio as gr
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("app")

from pipeline import HausaVoiceAIPipeline
from nlu import NLU
from orchestrator import Orchestrator

# ── Optional integrations (never block startup) ──────────────────────────────
try:
    from integrations.crm import CRMClient
    _crm = CRMClient()
    INTEGRATIONS_OK = True
except Exception as e:
    logger.warning(f"Integrations unavailable ({e}) — demo mode.")
    _crm, INTEGRATIONS_OK = None, False

# ── Config ───────────────────────────────────────────────────────────────────
HAS_GPU       = os.getenv("SPACES_GPU", "") or os.getenv("CUDA_VISIBLE_DEVICES", "")
SHOW_PARTIALS = os.getenv("SHOW_PARTIALS", "0") == "1"   # off by default on CPU
STREAM_EVERY  = float(os.getenv("STREAM_EVERY", "0.5"))

ai_pipeline = HausaVoiceAIPipeline()
dm          = Orchestrator(crm=_crm, nlu=NLU())

DEMO_PROMPTS = [
    ("Compound request (balance + transfer)",
     "Duba asusuna sannan ka aika 35000 zuwa Amina"),
    ("Entity prefill — recipient named",
     "Ina son aika kuɗi zuwa Abu"),
    ("Branch info, NOT block-card",
     "Ina ne reshenku mafi kusa domin in karɓi katin ATM"),
    ("Report a problem", "Ina da matsala da asusuna"),
    ("Escalate to human", "Ina son magana da mutum"),
]


# ── State ────────────────────────────────────────────────────────────────────

def new_state():
    return {"conv": dm.new_session(), "history": [], "sasr": None,
            "partial": "", "warm": False}


def render(history):
    if not history:
        return ('<div class="empty">Press <b>Warm up models</b>, then speak or '
                'tap a demo prompt.</div>')
    out = []
    for m in history:
        me = m["role"] == "user"
        out.append(
            f'<div class="row {"r-me" if me else "r-ag"}">'
            f'<div class="bub {"b-me" if me else "b-ag"}">'
            f'<div class="lbl">{"YOU" if me else "AGENT"}</div>'
            f'<div class="ha">{m["hausa"]}</div>'
            f'<div class="en">{m["english"]}</div>'
            f'<div class="tm">{m["time"]}</div></div></div>')
    return "".join(out)


# ── Core turn ────────────────────────────────────────────────────────────────

def run_turn(hausa_text, state, asr_ms=0.0):
    import time
    t = time.perf_counter()
    english = ai_pipeline.hausa_to_english(hausa_text)
    mt_in = (time.perf_counter() - t) * 1000

    t = time.perf_counter()
    reply_en, conv, escalated = dm.respond(english, hausa_text, state["conv"])
    state["conv"] = conv
    nlu_ms = (time.perf_counter() - t) * 1000

    t = time.perf_counter()
    reply_ha = ai_pipeline.english_to_hausa(reply_en)
    mt_out = (time.perf_counter() - t) * 1000

    t = time.perf_counter()
    sr, wav = ai_pipeline.hausa_text_to_audio(reply_ha)
    tts_ms = (time.perf_counter() - t) * 1000

    now = datetime.now().strftime("%H:%M:%S")
    state["history"] += [
        {"role": "user",  "hausa": hausa_text, "english": english,  "time": now},
        {"role": "agent", "hausa": reply_ha,   "english": reply_en, "time": now},
    ]

    total = asr_ms + mt_in + nlu_ms + mt_out + tts_ms
    openq = [x for x in conv.tasks
             if x.status in ("pending", "collecting", "confirming")]
    status = (f"turn {conv.turn} · queue {len(openq)} open / "
              f"{sum(1 for x in conv.tasks if x.status=='done')} done"
              f"{' · ' + conv.active_task.intent if conv.active_task else ''}"
              f"  |  ASR {asr_ms:.0f} · MT {mt_in+mt_out:.0f} · "
              f"NLU {nlu_ms:.0f} · TTS {tts_ms:.0f} · total {total:.0f} ms")
    if escalated:
        status += "   ESCALATED"
    return (sr, wav), status


# ── Handlers ─────────────────────────────────────────────────────────────────

def warm_up(state, progress=gr.Progress()):
    """Download + load models with visible progress, so the first real
    utterance isn't a 3-minute silence."""
    if state is None:
        state = new_state()
    try:
        progress(0.05, desc="Loading ASR …")
        ai_pipeline._load_asr()
        progress(0.55, desc="Loading translation …")
        ai_pipeline._load_nllb()
        progress(0.85, desc="Loading speech synthesis …")
        ai_pipeline._load_tts()
        progress(1.0, desc="Ready")
        state["warm"] = True
        gr.Info("Models loaded. You can speak now.")
        return f"Ready · {ai_pipeline.describe_models()}", state
    except Exception as e:
        logger.exception("warm-up failed")
        gr.Warning(f"Model load failed: {e}")
        return f"Model load failed: {e}", state


def on_text(text, state):
    if state is None:
        state = new_state()
    if not text or not text.strip():
        return gr.skip(), render(state["history"]), "Type or say something first.", state
    try:
        audio, status = run_turn(text.strip(), state)
        return audio, render(state["history"]), status, state
    except Exception as e:
        logger.exception("turn failed")
        gr.Warning(str(e))
        return gr.skip(), render(state["history"]), f"Error: {e}", state


def on_record(audio, state):
    """Push-to-talk: one complete recording."""
    if state is None:
        state = new_state()
    if audio is None:
        return gr.skip(), render(state["history"]), "No audio recorded.", state
    sr, arr = audio
    try:
        import time
        t = time.perf_counter()
        hausa = ai_pipeline.transcribe(arr, sr)
        asr_ms = (time.perf_counter() - t) * 1000
        if not hausa.strip():
            return gr.skip(), render(state["history"]), "Didn't catch that.", state
        out, status = run_turn(hausa, state, asr_ms)
        return out, render(state["history"]), status, state
    except Exception as e:
        logger.exception("record failed")
        gr.Warning(str(e))
        return gr.skip(), render(state["history"]), f"Error: {e}", state


def on_stream(chunk, state):
    """
    Live mic. Fires every STREAM_EVERY seconds. VAD decides when the turn ends,
    so there is no send button.

    Returns gr.skip() for the audio output on ticks with no reply, otherwise
    the player restarts on every tick.
    """
    if state is None:
        state = new_state()
    if chunk is None:
        return gr.skip(), gr.skip(), gr.skip(), state.get("partial", ""), state

    if state["sasr"] is None:
        state["sasr"] = ai_pipeline.make_streaming_session(
            emit_partials=SHOW_PARTIALS)

    sr, arr = chunk
    try:
        events = state["sasr"].accept_audio(arr, sr)
    except Exception as e:
        logger.exception("stream failed")
        return gr.skip(), gr.skip(), f"Stream error: {e}", "", state

    audio_out, status, convo = gr.skip(), gr.skip(), gr.skip()

    for ev in events:
        if ev.kind == "speech_start":
            status = "listening — speech detected"
        elif ev.kind == "partial":
            state["partial"] = ev.text
            status = f"listening … {ev.duration_ms/1000:.1f}s"
        elif ev.kind == "bargein":
            audio_out = None
            status = "you interrupted — go ahead"
        elif ev.kind == "discarded":
            state["partial"] = ""
        elif ev.kind == "final":
            state["partial"] = ""
            state["sasr"].agent_speaking = True
            audio_out, status = run_turn(ev.text, state, ev.latency_ms)
            convo = render(state["history"])

    return audio_out, convo, status, state.get("partial", ""), state


def reset(state):
    return None, render([]), "New session.", "", new_state()


# ── UI ───────────────────────────────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Sora:wght@400;600;800&family=IBM+Plex+Mono:wght@400;500&display=swap');
:root{--bg:#0B0B0F;--pnl:#15151C;--brd:#26262F;--txt:#E8E6E3;--mut:#7A7A88;
      --acc:#FF8A3D;--acc2:#4ADE80;--blu:#60A5FA}
.gradio-container{background:var(--bg)!important;font-family:'Sora',sans-serif!important;
      color:var(--txt)!important;max-width:1280px!important}
.hdr{background:linear-gradient(135deg,#15151C,#1E1E28 60%,#15151C);
     border:1px solid var(--brd);border-radius:14px;padding:22px 26px;margin-bottom:14px}
.hdr h1{font-size:25px;font-weight:800;margin:0;letter-spacing:-.5px}
.hdr h1 b{color:var(--acc)}
.hdr p{color:var(--mut);font-size:12.5px;margin:6px 0 0;letter-spacing:.4px}
.tags{display:flex;gap:7px;margin-top:13px;flex-wrap:wrap}
.tag{background:rgba(255,138,61,.12);border:1px solid rgba(255,138,61,.3);
     color:var(--acc);padding:3px 11px;border-radius:99px;font-size:10.5px;
     font-weight:600;letter-spacing:.6px}
.tag.g{background:rgba(74,222,128,.12);border-color:rgba(74,222,128,.3);color:var(--acc2)}
.tag.b{background:rgba(96,165,250,.12);border-color:rgba(96,165,250,.3);color:var(--blu)}
.convo{background:var(--pnl);border:1px solid var(--brd);border-radius:13px;
       padding:15px;height:395px;overflow-y:auto}
.empty{color:var(--mut);text-align:center;margin-top:150px;font-size:13px}
.row{display:flex;margin-bottom:11px}
.r-me{justify-content:flex-end}.r-ag{justify-content:flex-start}
.bub{max-width:86%;padding:10px 13px;border-radius:13px;font-size:13.5px;line-height:1.5}
.b-me{background:rgba(255,138,61,.14);border:1px solid rgba(255,138,61,.26);
      border-bottom-right-radius:4px}
.b-ag{background:rgba(255,255,255,.045);border:1px solid var(--brd);
      border-bottom-left-radius:4px}
.lbl{font-size:9.5px;font-weight:700;letter-spacing:.9px;opacity:.55;margin-bottom:4px}
.ha{font-weight:600}
.en{font-size:11.5px;color:var(--mut);margin-top:3px;font-style:italic}
.tm{font-size:9.5px;color:#4A4A56;margin-top:4px;font-family:'IBM Plex Mono',monospace}
.note{background:rgba(96,165,250,.07);border:1px solid rgba(96,165,250,.22);
      border-radius:11px;padding:13px 15px;font-size:12.5px;color:#B9C4D4;line-height:1.65}
.note b{color:var(--blu)}
footer{display:none!important}
"""

with gr.Blocks(css=CSS, title="PlotWeaver · Hausa Voice AI",
               theme=gr.themes.Base()) as demo:

    st = gr.State(None)

    gr.HTML(f"""
    <div class="hdr">
      <h1>Plot<b>Weaver</b> — Hausa Voice AI Agent</h1>
      <p>Real-time conversational AI for 100M+ Hausa speakers · investor POC</p>
      <div class="tags">
        <span class="tag">Whisper ASR</span>
        <span class="tag">NLLB-200</span>
        <span class="tag">MMS-TTS</span>
        <span class="tag g">VAD endpointing</span>
        <span class="tag g">Multi-intent NLU</span>
        <span class="tag b">{'GPU' if HAS_GPU else 'CPU'} mode</span>
      </div>
    </div>""")

    with gr.Tabs():
        with gr.TabItem("Live Demo"):
            gr.HTML('<div class="note"><b>First run:</b> press '
                    '<b>Warm up models</b> — roughly 2–4 minutes while ~3.5GB '
                    'downloads and caches. It only happens once per Space '
                    'restart. On CPU, expect 5–15s per reply; a GPU brings '
                    'that under 3s.</div>')

            with gr.Row():
                with gr.Column(scale=3):
                    convo = gr.HTML(render([]), elem_classes=["convo"])
                    partial_box = gr.Textbox(label="Live transcript",
                                             interactive=False, lines=1)
                    status = gr.Textbox(
                        label="Status",
                        value=("Not warmed up." if INTEGRATIONS_OK else
                               "Not warmed up · demo mode (no integrations)."),
                        interactive=False, lines=2)

                with gr.Column(scale=2):
                    warm_btn = gr.Button("Warm up models", variant="primary")

                    reply = gr.Audio(label="Agent reply", autoplay=True,
                                     interactive=False)

                    with gr.Tab("Type"):
                        txt = gr.Textbox(
                            label="Hausa text",
                            placeholder="Duba asusuna sannan ka aika 35000 zuwa Amina",
                            lines=2)
                        send = gr.Button("Send", variant="primary")

                    with gr.Tab("Record"):
                        rec = gr.Audio(sources=["microphone", "upload"],
                                       type="numpy", label="Record, then Send")
                        send_rec = gr.Button("Send recording", variant="primary")

                    with gr.Tab("Live mic (VAD)"):
                        gr.Markdown(
                            "Speak naturally — the VAD ends your turn on a "
                            "pause. Heavy on CPU; use Type or Record if it lags.")
                        live = gr.Audio(sources=["microphone"], type="numpy",
                                        streaming=True, label="Live")

                    reset_btn = gr.Button("Reset conversation")

                    gr.Markdown("**Demo prompts**")
                    for label, phrase in DEMO_PROMPTS:
                        gr.Button(label, size="sm").click(
                            lambda p=phrase: p, outputs=[txt])

            warm_btn.click(warm_up, [st], [status, st])
            send.click(on_text, [txt, st], [reply, convo, status, st])
            txt.submit(on_text, [txt, st], [reply, convo, status, st])
            send_rec.click(on_record, [rec, st], [reply, convo, status, st])
            live.stream(on_stream, [live, st],
                        [reply, convo, status, partial_box, st],
                        stream_every=STREAM_EVERY, show_progress="hidden")
            reset_btn.click(reset, [st],
                            [reply, convo, status, partial_box, st])

        with gr.TabItem("How it works"):
            gr.Markdown("""
### Pipeline

`mic → VAD endpointing → Whisper (Hausa) → NLLB hau→eng → task-queue
orchestrator → NLLB eng→hau → MMS-TTS → audio`

### What the dialogue layer does

Most voice bots resolve one intent and one missing slot per turn, which breaks
the moment a caller combines requests or answers out of order. This one
decomposes each message into a **task queue**:

| Caller says | What happens |
|---|---|
| "check my balance **and also** send 35000 to Amina" | Both tasks queued and acknowledged; neither is silently dropped |
| "send money **to Abu**" | Recipient prefilled from the utterance — never re-asked |
| "where's your branch so I can **get** my ATM card" | Branch info. *Not* block-card — destructive intents need an explicit verb |
| "too small" (as a return reason) | Accepted as a slot answer, not a fallback dead-end |
| Transfer above balance | Refused before the confirmation prompt |
| Two unparseable turns | Human handoff **with the transcript attached** |

Money actions require confidence ≥ 0.75 *and* an explicit yes/no.

### Model configuration

Defaults are tuned for a free CPU Space. Override via Space secrets:

| Variable | Default | Notes |
|---|---|---|
| `ASR_FINAL_MODEL` | `openai/whisper-small` | `whisper-large-v3` on GPU |
| `SHOW_PARTIALS` | `0` | `1` only on GPU |
| `HF_TOKEN` | unset | Enables LLM-based NLU (big quality gain) |
| `CRM_PROVIDER` | `demo` | `zendesk` for real tickets |

Without `HF_TOKEN` the NLU uses rule-based decomposition: it passes the full
regression suite but only handles phrasings that were anticipated.
""")

        with gr.TabItem("Market"):
            gr.Markdown("""
### Why Hausa, why now

| | |
|---|---|
| **100M+** | Hausa speakers — the largest language in West Africa |
| **~63%** | Low literacy in core regions — voice *is* the interface |
| **$4.2B** | Projected African contact-centre spend by 2027 |
| **~0** | Production-grade Hausa voice agents in market |

**Target verticals** — telecoms (MTN, Airtel, Glo), fintech and mobile money
(Kuda, PalmPay, OPay), public services (NIMC), health IVR, e-commerce.

**Moat** — open-weights stack with no per-call API lock-in; fine-tuned Hausa
models as owned IP; deployable on-premise for data-residency requirements;
multi-channel (voice, WhatsApp, SIP) from day one.

**Roadmap** — Yorùbá, Igbo, Fulfulde, Kanuri. NLLB covers 200 languages and
MMS covers 1,000+ for TTS, so each new language is a fine-tune, not a rebuild.
""")

    gr.HTML('<div style="text-align:center;padding:16px;color:#3A3A46;'
            'font-size:11px">PlotWeaver · open-weights Hausa voice AI</div>')


if __name__ == "__main__":
    demo.queue(max_size=12).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860)),
    )