"""WitnessBox — Gradio Space entrypoint.

Cross-examine Marcus Reid with your voice. Your *delivery* (perceived vocal
confidence) steers him; surface three contradictions and his voice cracks.

Boots anywhere: with WITNESSBOX_BACKEND unset it runs the offline mock end to
end (type your questions). Set WITNESSBOX_BACKEND=modal + Modal Space secrets
for live Whisper ASR / MiniCPM4.1-8B / VoxCPM2 and push-to-talk.
"""
from __future__ import annotations

import os

import numpy as np
import gradio as gr

import config
from witnessbox.backends import get_backends
from witnessbox.engine import WitnessBoxEngine
from witnessbox.witness import WITNESS_NAME, WITNESS_ROLE

CSS = """
/* ===== WitnessBox — a courtroom palette: parchment · ink · oxblood · brass =====
   Overrides Gradio's variables in BOTH light and dark mode so the look is
   consistent and there's not a drop of the default indigo anywhere. */
.gradio-container, .gradio-container.dark {
  --wb-ink:#2b2014; --wb-ink-soft:#6d5a3a;
  --wb-card:#f5edd9; --wb-card-2:#efe4cb;
  --wb-border:#c7b083; --wb-border-2:#a98c58;
  --wb-oxblood:#7c1f25; --wb-oxblood-hi:#94333a;
  --wb-brass:#9c7a3c; --wb-walnut:#241a10;

  --body-background-fill: radial-gradient(125% 95% at 50% -12%, #f1e7cf 0%, #e3d3ab 72%);
  --body-text-color: var(--wb-ink);
  --body-text-color-subdued: var(--wb-ink-soft);
  --background-fill-primary: var(--wb-card);
  --background-fill-secondary: var(--wb-card-2);
  --block-background-fill: var(--wb-card);
  --block-border-color: var(--wb-border);
  --block-label-background-fill: var(--wb-oxblood);
  --block-label-text-color: #f4e8c9;
  --block-title-text-color: var(--wb-ink);
  --block-info-text-color: var(--wb-ink-soft);
  --border-color-primary: var(--wb-border);
  --border-color-accent: var(--wb-brass);
  --button-primary-background-fill: linear-gradient(180deg,#8e2731,#6f1b21);
  --button-primary-background-fill-hover: linear-gradient(180deg,#9c333b,#7c1f25);
  --button-primary-text-color:#f6ecd2;
  --button-primary-border-color:#591319;
  --button-secondary-background-fill: linear-gradient(180deg,#ecdcb4,#dcc798);
  --button-secondary-background-fill-hover: linear-gradient(180deg,#f2e4c0,#e3d0a4);
  --button-secondary-text-color: var(--wb-ink);
  --button-secondary-border-color: var(--wb-border-2);
  --input-background-fill:#fbf6e8;
  --input-border-color: var(--wb-border);
  --input-placeholder-color:#9a8763;
  --color-accent: var(--wb-brass);
  --color-accent-soft:#ecdfbc;
  --link-text-color: var(--wb-oxblood);
  --link-text-color-hover: var(--wb-oxblood-hi);
  --block-shadow: 0 2px 7px rgba(60,40,15,.10);
  --block-radius:12px;
  font-family:'Iowan Old Style','Palatino Linotype','Palatino',Georgia,'Times New Roman',serif;
}
.gradio-container { background: var(--body-background-fill); }

/* ---- header ---- */
#wb-title {text-align:center; padding:14px 0 2px;}
#wb-title .wb-crest {font-size:1.5rem; line-height:1; opacity:.85;}
#wb-title h1 {color:var(--wb-ink)!important; font-variant:small-caps; letter-spacing:1.5px;
  font-weight:700; font-size:2.7rem; margin:.05em 0 .04em; line-height:1.05;}
#wb-title .wb-sub {color:var(--wb-ink-soft)!important; font-size:1.06rem; letter-spacing:.2px;}
#wb-title .wb-sub b {color:var(--wb-oxblood); font-weight:700;}
#wb-title .wb-rule {width:210px; height:2px; margin:11px auto 0;
  background:linear-gradient(90deg,transparent,var(--wb-brass),transparent);}

/* ---- small-caps block labels everywhere ---- */
.gradio-container .block-label, .gradio-container label > span:first-child {
  font-variant:small-caps; letter-spacing:.6px; font-weight:700;}

/* ---- chatbot = a deposition transcript on aged paper ---- */
#wb-chat {background:var(--wb-card)!important; border:1px solid var(--wb-border)!important;}
#wb-chat *, #wb-chat .message, #wb-chat .prose {color:var(--wb-ink)!important;}
#wb-chat .message {font-size:1.02rem!important; line-height:1.5!important; border-radius:10px!important;}
#wb-chat .user, #wb-chat [data-role="user"] {background:#efe3c6!important; border:1px solid var(--wb-border)!important;}
#wb-chat .bot, #wb-chat [data-role="assistant"] {background:#fbf6e8!important; border-left:3px solid var(--wb-oxblood)!important;}
#wb-chat .placeholder, #wb-chat .placeholder * {color:var(--wb-ink-soft)!important; opacity:.7;}

/* ---- banner ---- */
.wb-banner {text-align:center; font-size:1.15rem; font-variant:small-caps; letter-spacing:.6px;
  padding:12px; border-radius:10px; border:1px solid var(--wb-border);}

/* ---- custom HTML panels (stance / counters) ---- */
.wb-card {background:var(--wb-card); border:1px solid var(--wb-border); border-radius:10px; padding:13px 15px; box-shadow:0 1px 0 #fffdf6 inset;}
.wb-bar-track {background:#e0d3af; border-radius:7px; height:16px; overflow:hidden; border:1px solid var(--wb-border);}
.wb-bar-fill {height:100%; transition:width .4s ease;}
.wb-disclaimer {font-size:11px; color:var(--wb-ink-soft); font-style:italic;}
.wb-tier {font-variant:small-caps; font-weight:700; color:var(--wb-oxblood); letter-spacing:.5px;}

/* ---- portrait framing ---- */
#wb-portrait img {border:1px solid var(--wb-border-2); border-radius:10px;}

/* ---- evidence: a walnut 'exhibit' readout ---- */
#wb-evidence textarea {font-family:ui-monospace,Menlo,Consolas,monospace!important;
  background:var(--wb-walnut)!important; color:#ecdfae!important;
  border:1px solid var(--wb-brass)!important; border-radius:8px;}

/* ---- every component label = the same oxblood plate (no stray pink/purple) ---- */
.gradio-container .block-label, .gradio-container .block-label * {
  background:var(--wb-oxblood)!important; color:#f4e8c9!important; opacity:1!important;
  border-color:var(--wb-border)!important;}
.gradio-container .block-label svg, .gradio-container .block-label path {fill:#f4e8c9!important;}

/* ---- accordion header: dark, legible, small-caps ---- */
.gradio-container .label-wrap, .gradio-container .label-wrap > span {
  color:var(--wb-ink)!important; font-weight:700; font-variant:small-caps; letter-spacing:.5px;}
.gradio-container .label-wrap .icon, .gradio-container .label-wrap svg {color:var(--wb-oxblood)!important; opacity:1;}

/* Paint the page edge-to-edge in parchment in BOTH modes (no dark gutters). */
html, body, gradio-app {background:#e6d8b8 !important;}
"""


# --------------------------------------------------------------------------- #
# render helpers
# --------------------------------------------------------------------------- #
def _bar(label: str, pct: float, color: str, sub: str = "") -> str:
    pct = max(0, min(100, int(round(pct))))
    return (
        f"<div class='wb-card' style='margin-bottom:8px'>"
        f"<div style='display:flex;justify-content:space-between'>"
        f"<b>{label}</b><span>{pct}</span></div>"
        f"<div class='wb-bar-track'><div class='wb-bar-fill' style='width:{pct}%;background:{color}'></div></div>"
        f"{f'<div class=wb-disclaimer>{sub}</div>' if sub else ''}</div>"
    )


def _stance_html(stance) -> str:
    color = {"CONFIDENT": "#2f7d3b", "NEUTRAL": "#b08900", "HESITANT": "#9c3b2f"}.get(stance.tier, "#b08900")
    sub = "Perceived delivery — NOT a lie detector. Reads pauses &amp; pace, not truth."
    head = f"<div class='wb-tier'>Delivery&nbsp;·&nbsp;{stance.tier}</div>"
    return head + _bar("Perceived confidence", stance.confidence, color, sub)


def _counters_html(status: dict) -> str:
    catches = f"<div class='wb-card' style='margin-bottom:8px'><b>Contradictions</b> " \
              f"<span style='float:right'>{status['catches']} / {status['catches_to_win']}</span></div>"
    cred = _bar("Your standing with the bench", status["credibility"], "#43607f")
    comp = _bar(f"Witness composure · {status['witness_tier']}", status["composure"], "#7a4a2f")
    return catches + cred + comp


def _parse_mic(mic):
    if mic is None:
        return None, None
    sr, data = mic
    y = np.asarray(data)
    if y.dtype.kind in "iu":
        y = y.astype(np.float32) / max(1, np.iinfo(y.dtype).max)
    else:
        y = y.astype(np.float32)
    if y.ndim > 1:
        y = y.mean(axis=1)
    return y, int(sr)


def _concat(a, b, sr):
    if a is None:
        return b
    if b is None:
        return a
    gap = np.zeros(int(0.5 * sr), dtype=np.float32)
    return np.concatenate([a.astype(np.float32), gap, b.astype(np.float32)])


def _banner(kind: str, text: str) -> str:
    colors = {"win": "#2f7d3b;color:#fff", "lose": "#7a2f2f;color:#fff", "info": "#e9dfc3;color:#5a4220"}
    bg = colors.get(kind, colors["info"])
    return f"<div class='wb-banner' style='background:{bg}'>{text}</div>"


# --------------------------------------------------------------------------- #
# callbacks
# --------------------------------------------------------------------------- #
def on_start(engine):
    engine = WitnessBoxEngine(get_backends())
    intro = engine.start()
    chat = [
        {"role": "assistant", "content": f"⚖️ *The Court:* {intro['narration']}"},
        {"role": "assistant", "content": f"**{WITNESS_NAME}:** {intro['opening_text']}"},
    ]
    opening_audio = intro["opening_audio"]  # (sr, np) or None
    footer = f"Backend: **{intro['backend']}** — {intro['backend_note']}"
    from witnessbox.stance import _neutral
    return (
        engine,
        chat,
        gr.update(value=opening_audio),
        _stance_html(_neutral("awaiting your first question")),
        _counters_html(intro["status"]),
        gr.update(value="", visible=False),
        _banner("info", "Examination open. Mind how you say it — he listens for doubt."),
        footer,
        gr.update(interactive=True),       # ask button
        gr.update(visible=False),          # begin button
        gr.update(interactive=True),       # mic
        gr.update(interactive=True),       # typed
    )


def on_ask(engine, mic, typed):
    if engine is None:
        return (engine, gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(),
                _banner("info", "Press “Call the witness” to begin."), gr.skip())

    try:
        y, sr = _parse_mic(mic)
        result = engine.take_turn(audio=y, sr=sr, typed_text=typed)
    except Exception as exc:  # surface the real cause in the UI, don't silently toast
        import traceback
        traceback.print_exc()
        return (engine, gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(),
                _banner("lose", f"Turn failed — {type(exc).__name__}: {exc}"), gr.skip())

    # Rebuild the chat from the transcript (engine keeps it consistent with what
    # is actually spoken, including the break line on the winning turn).
    chat = []
    for rec in engine.state.transcript:
        tag = f"_[{rec.stance_tier.lower()}]_ " if rec.stance_tier != "NEUTRAL" else ""
        chat.append({"role": "user", "content": f"{tag}{rec.examiner_text}"})
        chat.append({"role": "assistant", "content": f"**{WITNESS_NAME}:** {rec.witness_text}"})
    # Terminal turns (e.g. no clear audio) don't enter the transcript — show them
    # anyway so the player always gets visible feedback.
    last_bot = chat[-1]["content"] if chat else ""
    if result.witness_text and result.witness_text not in last_bot:
        if result.examiner_text:
            chat.append({"role": "user", "content": result.examiner_text})
        chat.append({"role": "assistant", "content": f"**{WITNESS_NAME}:** {result.witness_text}"})

    # witness audio (+ epilogue concatenated on win/lose for a single dramatic play)
    audio_val = None
    if result.witness_audio is not None:
        merged = _concat(result.witness_audio, result.epilogue_audio, result.audio_sr)
        audio_val = (result.audio_sr, merged)

    # banner
    if result.events.won:
        banner = _banner("win", "🩻 He breaks. Three contradictions on the record — you win.")
    elif result.events.lost:
        banner = _banner("lose", "The bench excuses the witness. You’ve lost the room.")
    elif result.events.near_miss:
        banner = _banner("info", "He flinched. You’re circling something — name the specific fact.")
    else:
        banner = _banner("info", f"Stance read: {result.stance.tier.title()}.")

    evidence_update = (
        gr.update(value=result.evidence, visible=True)
        if result.evidence else gr.update()
    )
    return (
        engine,
        chat,
        gr.update(value=audio_val),
        _stance_html(result.stance),
        _counters_html(result.status),
        evidence_update,
        banner,
        gr.update(value=""),   # clear typed box
    )


# --------------------------------------------------------------------------- #
# layout
# --------------------------------------------------------------------------- #
def build() -> gr.Blocks:
    theme = gr.themes.Soft(
        primary_hue=gr.themes.colors.red,      # oxblood, refined to exact hues in CSS
        secondary_hue=gr.themes.colors.amber,  # brass
        neutral_hue=gr.themes.colors.stone,    # warm paper greys, never blue-greys
    )
    with gr.Blocks(css=CSS, title="WitnessBox", theme=theme) as demo:
        engine_state = gr.State(None)
        gr.HTML(
            "<div id='wb-title'>"
            "<div class='wb-crest'>⚖️</div>"
            "<h1>WitnessBox</h1>"
            f"<div class='wb-sub'>Cross-examine <b>{WITNESS_NAME}</b>, {WITNESS_ROLE}. "
            "Your <b>voice</b> is the weapon.</div>"
            "<div class='wb-rule'></div>"
            "</div>"
        )
        banner = gr.HTML(_banner("info", "Call the witness to the stand."))

        with gr.Row():
            with gr.Column(scale=2):
                _portrait = "assets/marcus_reid.png"
                gr.Image(
                    value=_portrait if os.path.exists(_portrait) else None,
                    show_label=False, height=300, elem_id="wb-portrait",
                    show_download_button=False, container=True,
                )
                stance_html = gr.HTML(label="Delivery")
            with gr.Column(scale=4):
                chat = gr.Chatbot(type="messages", height=420, label="The Stand",
                                  elem_id="wb-chat")
                witness_audio = gr.Audio(label="Witness", autoplay=True, interactive=False)
            with gr.Column(scale=2):
                counters_html = gr.HTML()

        with gr.Accordion("🔎 Contradiction Engine (live verdict)", open=True):
            evidence = gr.Textbox(
                elem_id="wb-evidence", show_label=False, visible=False, lines=5,
                interactive=False,
            )
            gr.Markdown(
                "_Catches are decided by a deterministic engine over three planted "
                "contradictions — the language model never grades itself, so the "
                "verdict is reproducible._"
            )

        with gr.Row():
            mic = gr.Audio(sources=["microphone"], type="numpy", label="Question (push to talk)",
                           interactive=False)
            typed = gr.Textbox(label="…or type your question (primary in offline mock mode)",
                               interactive=False, scale=2,
                               placeholder="e.g. The wire cleared March 6th — before the board approved it on the 14th.")
        with gr.Row():
            begin_btn = gr.Button("Call the witness to the stand", variant="primary")
            ask_btn = gr.Button("Put it to him", variant="secondary", interactive=False)

        footer = gr.Markdown("")

        outs_start = [engine_state, chat, witness_audio, stance_html, counters_html,
                      evidence, banner, footer, ask_btn, begin_btn, mic, typed]
        begin_btn.click(on_start, [engine_state], outs_start)

        outs_ask = [engine_state, chat, witness_audio, stance_html, counters_html,
                    evidence, banner, typed]
        ask_btn.click(on_ask, [engine_state, mic, typed], outs_ask)
        typed.submit(on_ask, [engine_state, mic, typed], outs_ask)

    return demo


demo = build()

if __name__ == "__main__":
    # server_name=0.0.0.0 so the HF Space's startup self-check can reach the app.
    # ssr_mode=False: Gradio 5's experimental SSR layer caused flaky request errors
    # on the Space; the classic client-rendered path is reliable for this app.
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", 7860)),
        ssr_mode=False,
    )