import os
import re
from collections.abc import Iterator
from pathlib import Path
from threading import Thread

PREVIEW = os.getenv("TWIL_UI_PREVIEW") == "1"

if not PREVIEW:
    import spaces
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
else:

    class spaces:  # type: ignore[no-redef]
        @staticmethod
        def GPU(duration=None, **kwargs):
            def decorator(fn):
                return fn

            return decorator

import gradio as gr

MODEL_ID = "webAI-Official/TwIL-LM3"
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))
ASSETS = Path(__file__).resolve().parent / "brand"
AVATAR = ASSETS / "avatar.svg"

THINK_OPEN = chr(60) + "think" + chr(62)
THINK_CLOSE = chr(60) + chr(47) + "think" + chr(62)
_THINK_RE = re.compile(
    re.escape(THINK_OPEN) + r".*?(" + re.escape(THINK_CLOSE) + r"|$)",
    re.DOTALL,
)

SYSTEM_PROMPT = (
    "You are TwIL, a formal-logic reasoning model created by webAI Intelligence Lab. "
    "You specialise in entailment, first-order logic, semantic parsing, Lean, and proof critique. "
    f"Work through the problem in a {THINK_OPEN} block, then give a concise, precise answer."
)

CHIPS = [
    (
        "Analyze an argument",
        "Analyze this argument: identify the premises and the conclusion, say whether it is valid, and name any fallacy.\n\n",
    ),
    (
        "Check an entailment",
        "Do these premises entail the conclusion? Give a formal proof or a counterexample.\n\nPremise 1: \nPremise 2: \nConclusion: ",
    ),
    (
        "Translate to FOL",
        "Translate into first-order logic: Every student who studies hard passes at least one exam.",
    ),
]

SHOWCASE_TEMPLATE = (
    "Analyze this argument: extract the premises and the conclusion, translate them "
    "into first-order logic, state whether the argument is valid, and name any fallacy.\n\n"
    'Argument: "{argument}"'
)

SHOWCASE = [
    {
        "kicker": "Viral tweet",
        "quote": "The streets are wet this morning, and rain always makes the streets wet. So it obviously rained last night.",
        "hint": "Can wet streets prove rain?",
    },
    {
        "kicker": "School board speech",
        "quote": "If we let students retake this one exam, soon they will demand to retake every exam, and before long no grade at this school will mean anything.",
        "hint": "Does one retake doom every grade?",
    },
    {
        "kicker": "Supplement ad",
        "quote": "Our formula is 100% natural, and nature knows best. That means it is completely safe for your body.",
        "hint": "Does natural imply safe?",
    },
    {
        "kicker": "Campaign rally",
        "quote": "Either you support this bill or you don't care about public safety. You oppose the bill, so you don't care about public safety.",
        "hint": "Are those really the only options?",
    },
    {
        "kicker": "Toothpaste commercial",
        "quote": "4 out of 5 dentists we surveyed recommend SparkleDent, and experts know best. You should switch to SparkleDent today.",
        "hint": "What is the survey hiding?",
    },
    {
        "kicker": "HR policy memo",
        "quote": "Every employee who completes the security training receives a certificate. Sam completed the security training. Therefore Sam receives a certificate.",
        "hint": "One of these six actually holds up.",
    },
]

SHOWCASE_INTRO = """
<div class="twil-showcase-intro">
  <h2>Can TwIL spot the flaw?</h2>
  <p>Six arguments from the wild. One click and TwIL extracts the premises, translates
  them to first-order logic, and rules on validity. Five are broken — one holds up.</p>
</div>
"""


def _card_html(case: dict) -> str:
    return (
        '<div class="twil-card-body">'
        f'<div class="twil-card-kicker">{case["kicker"]}</div>'
        f'<div class="twil-card-quote">&ldquo;{case["quote"]}&rdquo;</div>'
        f'<div class="twil-card-hint">{case["hint"]}</div>'
        "</div>"
    )


WAITLIST_URL = "https://www.webai.com/waitlist"

HEADER = f"""
<div class="twil-header">
  <img class="twil-mark" src="/gradio_api/file=brand/AppiCon.svg" alt="webAI" />
  <div class="twil-header-copy">
    <div class="twil-wordmark">webAI</div>
    <div class="twil-product">TwIL-LM3 · Intelligence Lab</div>
  </div>
  <a class="twil-cta twil-cta-header" href="{WAITLIST_URL}" target="_blank" rel="noopener">Get early access</a>
</div>
"""

HEAD = """
<script>
document.addEventListener("keydown", (e) => {
  if (e.key !== "Enter" || e.shiftKey || e.isComposing) return;
  const ta = document.querySelector("#twil-input textarea");
  if (!ta || e.target !== ta) return;
  e.preventDefault();
  e.stopPropagation();
  document.querySelector("#send-btn")?.click();
}, true);
</script>
"""

PLACEHOLDER = f"""
<div class="twil-empty">
  <img class="twil-cube" src="/gradio_api/file=brand/webai-cube-256.webp" alt="" />
  <h1>What should we work on?</h1>
  <p>TwIL-LM3 is webAI's formal-logic model — entailment, FOL, Lean, rule induction.</p>
  <div class="twil-waitlist">
    <p>TwIL and the full family of webAI models are coming to the webAI app.
    Be the first to try expert intelligence on everyday devices.</p>
    <a class="twil-cta" href="{WAITLIST_URL}" target="_blank" rel="noopener">Join the waitlist</a>
  </div>
</div>
"""

CSS = """
:root {
  --weba-canvas: #f4f4f4;
  --weba-fg: #161616;
  --weba-muted: #737373;
  --weba-secondary: #e8e8e8;
  --weba-border: rgba(22, 22, 22, 0.1);
  --weba-composer: #ffffff;
  --weba-primary: #232323;
  --weba-primary-fg: #fafafa;
}
.dark {
  --weba-canvas: #232323;
  --weba-fg: #fafafa;
  --weba-muted: #a8a8a8;
  --weba-secondary: #323232;
  --weba-border: rgba(255, 255, 255, 0.1);
  --weba-composer: #161616;
  --weba-primary: #e8e8e8;
  --weba-primary-fg: #232323;
}

* { box-shadow: none !important; text-shadow: none !important; }
html, body, #root, .gradio-container, .gradio-container > .main,
.gradio-container .contain, .fillable,
.gradio-container .column, .gradio-container .row, .contain, .wrapper {
  background: var(--weba-canvas) !important;
  box-shadow: none !important;
  filter: none !important;
}
/* Never paint over content while a job is running: the status tracker
   overlay must stay transparent or the chat looks blank mid-generation. */
.gradio-container [data-testid="status-tracker"],
.gradio-container .wrap.default {
  background: transparent !important;
}
html, body, .gradio-container, .gradio-container > .main, .fillable,
.contain, .app {
  height: 100% !important;
  min-height: 100vh !important;
  max-height: 100vh !important;
  max-width: none !important;
  margin: 0 !important;
  overflow: hidden !important;
}
.gradio-container, .main.fillable, .contain {
  display: flex !important;
  flex-direction: column !important;
  padding: 0 !important;
}
.gradio-container {
  font-family: ui-sans-serif, system-ui, -apple-system, "Segoe UI", sans-serif !important;
  color: var(--weba-fg) !important;
}
footer, .footer, .built-with, .settings, .settings-bar, .show-api,
.show-api-divider, .divider {
  display: none !important;
}

#app-shell {
  flex: 1 1 auto !important;
  height: 100% !important;
  max-height: 100% !important;
  min-height: 0 !important;
  display: flex !important;
  flex-direction: column !important;
  padding: 16px 24px 12px !important;
  box-sizing: border-box !important;
  overflow: hidden !important;
  background: var(--weba-canvas) !important;
}
#app-shell > .gap, #app-shell > div {
  background: transparent !important;
  flex-grow: 0 !important;
  height: auto !important;
  min-height: 0 !important;
}
#app-shell > #twil-chat,
#app-shell > .block:has(#twil-chat),
#app-shell > div:has(#twil-chat) {
  flex: 1 1 auto !important;
  min-height: 0 !important;
  height: auto !important;
}

.twil-header {
  display: flex;
  align-items: center;
  gap: 12px;
  padding: 4px 0 8px;
  flex-shrink: 0;
}
.twil-mark { width: 28px; height: 28px; border-radius: 6px; }
.twil-wordmark { font-weight: 600; font-size: 15px; letter-spacing: -0.02em; line-height: 1.2; color: var(--weba-fg); }
.twil-product { font-size: 12px; color: var(--weba-muted); line-height: 1.3; }

.twil-cta {
  display: inline-flex;
  align-items: center;
  height: 34px;
  padding: 0 18px;
  border-radius: 999px;
  background: var(--weba-primary);
  color: var(--weba-primary-fg) !important;
  font-size: 13px;
  font-weight: 500;
  text-decoration: none !important;
  border: none;
  white-space: nowrap;
  transition: opacity 0.15s ease;
}
.twil-cta:hover { opacity: 0.85; }
.twil-cta-header {
  /* Keep clear of the floating HF Space pill in the top-right corner */
  margin-left: 20px;
  height: 30px;
  padding: 0 14px;
  font-size: 12px;
}
.twil-waitlist {
  margin-top: 28px;
  display: flex;
  flex-direction: column;
  align-items: center;
  gap: 12px;
}
.twil-waitlist p {
  margin: 0;
  max-width: 30rem;
  font-size: 13px;
  line-height: 1.5;
  color: var(--weba-muted);
}

.twil-empty {
  height: 100%;
  min-height: 280px;
  display: flex;
  flex-direction: column;
  align-items: center;
  justify-content: center;
  text-align: center;
  padding: 12px;
}
.twil-cube { width: 96px; height: 96px; margin-bottom: 24px; }
.twil-empty h1 {
  font-size: 28px;
  font-weight: 500;
  letter-spacing: -0.03em;
  margin: 0 0 8px;
  color: var(--weba-fg);
}
.twil-empty p {
  margin: 0;
  max-width: 28rem;
  font-size: 15px;
  line-height: 1.5;
  color: var(--weba-muted);
}

#twil-chat, #twil-chat > .wrapper, #twil-chat .bubble-wrap,
#twil-chat .message-wrap, #twil-chat .placeholder-content {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  flex: 1 1 auto !important;
  min-height: 0 !important;
}
#twil-chat {
  max-width: 768px !important;
  width: 100% !important;
  margin: 0 auto !important;
  flex: 1 1 auto !important;
}
#twil-chat .message.user, #twil-chat .user, #twil-chat .bubble.user {
  background: var(--weba-secondary) !important;
  color: var(--weba-fg) !important;
  border: none !important;
  box-shadow: none !important;
  border-radius: 12px !important;
}
#twil-chat .message.bot, #twil-chat .bot, #twil-chat .bubble.bot {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
}
#twil-chat [aria-label="Delete"], #twil-chat [aria-label="Copy"],
#twil-chat [aria-label="Share"], #twil-chat .message-buttons {
  display: none !important;
}
#twil-chat details, #twil-chat .thought, #twil-chat .md.thought {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  color: var(--weba-muted) !important;
}

#composer-row {
  max-width: 768px !important;
  width: 100% !important;
  margin: 8px auto 0 !important;
  display: flex !important;
  flex-direction: row !important;
  flex-wrap: nowrap !important;
  align-items: center !important;
  gap: 10px !important;
  flex: 0 0 auto !important;
  min-height: 56px !important;
  background: transparent !important;
}
#composer-row > .block:has(#twil-input),
#twil-input {
  flex: 1 1 0% !important;
  min-width: 0 !important;
  width: auto !important;
  height: 52px !important;
  min-height: 52px !important;
  max-height: 52px !important;
  background: var(--weba-composer) !important;
  border: 1px solid var(--weba-border) !important;
  border-radius: 999px !important;
  box-shadow: none !important;
  overflow: hidden !important;
}
#twil-input label,
#twil-input .input-container {
  display: flex !important;
  align-items: stretch !important;
  width: 100% !important;
  height: 100% !important;
  min-height: 52px !important;
  margin: 0 !important;
  padding: 0 !important;
  border: none !important;
  background: transparent !important;
  box-sizing: border-box !important;
}
#twil-input textarea,
#twil-input input,
#twil-input .scroll-hide,
#twil-input [data-testid="textbox"] {
  display: block !important;
  background: transparent !important;
  color: var(--weba-fg) !important;
  border: none !important;
  box-shadow: none !important;
  width: 100% !important;
  height: 52px !important;
  min-height: 52px !important;
  max-height: 52px !important;
  padding: 16px 18px !important;
  font-size: 16px !important;
  line-height: 20px !important;
  overflow-y: auto !important;
  overflow-x: hidden !important;
  resize: none !important;
  white-space: pre-wrap !important;
  field-sizing: fixed !important;
  pointer-events: auto !important;
}
#twil-input button, #twil-input .show-count, #twil-input .icon-button,
#twil-input .sr-only, #twil-input span[data-testid="block-info"] {
  display: none !important;
}
#composer-row > .form,
#composer-row > .block:has(#send-btn) {
  flex: 0 0 40px !important;
  width: 40px !important;
  min-width: 40px !important;
  max-width: 40px !important;
  background: transparent !important;
  border: none !important;
  padding: 0 !important;
  box-shadow: none !important;
}
#send-btn {
  flex: 0 0 40px !important;
  width: 40px !important;
  min-width: 40px !important;
  max-width: 40px !important;
  height: 40px !important;
  border-radius: 999px !important;
  background: var(--weba-secondary) !important;
  color: var(--weba-fg) !important;
  border: none !important;
  box-shadow: none !important;
  align-self: center !important;
  padding: 0 !important;
}
#chips-row {
  max-width: 768px !important;
  width: 100% !important;
  margin: 0 auto !important;
  display: flex !important;
  flex-wrap: wrap !important;
  justify-content: center !important;
  gap: 8px !important;
  flex-shrink: 0 !important;
  background: transparent !important;
  padding: 4px 0 8px !important;
}
#chips-row button {
  height: 32px !important;
  padding: 0 14px !important;
  border-radius: 999px !important;
  border: 1px solid var(--weba-border) !important;
  background: transparent !important;
  color: var(--weba-muted) !important;
  box-shadow: none !important;
  font-size: 13px !important;
  font-weight: 400 !important;
}
#chips-row button:hover {
  background: var(--weba-secondary) !important;
  color: var(--weba-fg) !important;
}

@keyframes twil-pulse {
  0%, 100% { opacity: 0.4; }
  50% { opacity: 1; }
}
#twil-chat .pending, #twil-chat [data-status="pending"] {
  animation: twil-pulse 1.1s ease-in-out infinite;
}

#params-box,
#params-box.block,
#params-box .wrap,
#params-box .styler,
#params-box .gap,
#accordion-content {
  max-width: 768px !important;
  margin: 4px auto 0 !important;
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  flex: 0 0 auto !important;
  flex-grow: 0 !important;
  height: auto !important;
  min-height: 0 !important;
  max-height: none !important;
  overflow: visible !important;
}
#params-box .label-wrap, #params-box .icon {
  color: var(--weba-muted) !important;
  font-size: 12px !important;
  box-shadow: none !important;
  background: transparent !important;
  border: none !important;
}

.twil-disclaimer {
  margin: 8px 0 0;
  text-align: center;
  font-size: 11px;
  color: var(--weba-muted);
  flex-shrink: 0;
}

/* Tabs shell: keep the chat layout filling the viewport */
#twil-tabs {
  flex: 1 1 auto !important;
  min-height: 0 !important;
  display: flex !important;
  flex-direction: column !important;
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
}
#twil-tabs > .tab-wrapper,
#twil-tabs .tab-container,
#twil-tabs .tab-nav {
  flex: 0 0 auto !important;
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  justify-content: center !important;
}
#twil-tabs .tab-container button,
#twil-tabs .tab-nav button {
  background: transparent !important;
  border: none !important;
  border-bottom: 2px solid transparent !important;
  border-radius: 0 !important;
  box-shadow: none !important;
  color: var(--weba-muted) !important;
  font-size: 13px !important;
  padding: 6px 14px !important;
}
#twil-tabs .tab-container button.selected,
#twil-tabs .tab-nav button.selected {
  color: var(--weba-fg) !important;
  border-bottom-color: var(--weba-fg) !important;
}
#twil-tabs .tabitem {
  flex: 1 1 auto !important;
  min-height: 0 !important;
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  padding: 0 !important;
}
#tab-chat > .gap, #tab-chat > .column,
#tab-chat, #tab-chat .column:has(> div > #twil-chat) {
  height: 100% !important;
  min-height: 0 !important;
}
#tab-chat > .gap, #tab-chat > .column {
  display: flex !important;
  flex-direction: column !important;
}
#tab-chat .block:has(#twil-chat),
#tab-chat div:has(> #twil-chat) {
  flex: 1 1 auto !important;
  min-height: 0 !important;
  height: auto !important;
}
#tab-showcase {
  overflow-y: auto !important;
}

/* Showcase cards */
.twil-showcase-intro {
  max-width: 768px;
  margin: 8px auto 4px;
  text-align: center;
}
.twil-showcase-intro h2 {
  font-size: 22px;
  font-weight: 500;
  letter-spacing: -0.02em;
  margin: 0 0 6px;
  color: var(--weba-fg);
}
.twil-showcase-intro p {
  margin: 0 0 8px;
  font-size: 14px;
  line-height: 1.5;
  color: var(--weba-muted);
}
.twil-card-row {
  max-width: 900px !important;
  width: 100% !important;
  margin: 6px auto !important;
  display: flex !important;
  gap: 12px !important;
  align-items: stretch !important;
  background: transparent !important;
}
.twil-card {
  background: var(--weba-composer) !important;
  border: 1px solid var(--weba-border) !important;
  border-radius: 16px !important;
  box-shadow: none !important;
  padding: 16px !important;
  gap: 10px !important;
  display: flex !important;
  flex-direction: column !important;
  justify-content: space-between !important;
  min-height: 170px !important;
}
.twil-card > * { background: transparent !important; }
.twil-card-kicker {
  font-size: 11px;
  font-weight: 600;
  letter-spacing: 0.06em;
  text-transform: uppercase;
  color: var(--weba-muted);
  margin-bottom: 8px;
}
.twil-card-quote {
  font-size: 14px;
  line-height: 1.45;
  color: var(--weba-fg);
  margin-bottom: 8px;
}
.twil-card-hint {
  font-size: 12px;
  color: var(--weba-muted);
  font-style: italic;
}
.twil-card-btn {
  height: 32px !important;
  border-radius: 999px !important;
  border: 1px solid var(--weba-border) !important;
  background: var(--weba-secondary) !important;
  color: var(--weba-fg) !important;
  box-shadow: none !important;
  font-size: 13px !important;
  font-weight: 500 !important;
  align-self: flex-start !important;
  padding: 0 16px !important;
  width: auto !important;
  min-width: 0 !important;
  flex: 0 0 auto !important;
}
.twil-card-btn:hover {
  background: var(--weba-fg) !important;
  color: var(--weba-canvas) !important;
}
"""

THEME = gr.themes.Base(
    primary_hue="zinc",
    secondary_hue="zinc",
    neutral_hue="zinc",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    radius_size=gr.themes.sizes.radius_lg,
).set(
    body_background_fill="#f4f4f4",
    body_background_fill_dark="#232323",
    body_text_color="#161616",
    body_text_color_dark="#fafafa",
    background_fill_primary="transparent",
    background_fill_primary_dark="transparent",
    background_fill_secondary="#e8e8e8",
    background_fill_secondary_dark="#323232",
    border_color_primary="rgba(22,22,22,0.1)",
    border_color_primary_dark="rgba(255,255,255,0.1)",
    block_background_fill="transparent",
    block_background_fill_dark="transparent",
    block_border_width="0px",
    block_shadow="none",
    block_shadow_dark="none",
    shadow_drop="none",
    shadow_drop_lg="none",
    input_background_fill="#ffffff",
    input_background_fill_dark="#161616",
    input_border_color="rgba(22,22,22,0.1)",
    input_border_color_dark="rgba(255,255,255,0.1)",
    input_shadow="none",
    button_secondary_background_fill="#e8e8e8",
    button_secondary_background_fill_dark="#323232",
    button_secondary_text_color="#161616",
    button_secondary_text_color_dark="#fafafa",
    button_secondary_shadow="none",
)

tokenizer = None
model = None
EOS_TOKEN_ID = None

if not PREVIEW:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        dtype=torch.bfloat16,
        attn_implementation="sdpa",
    ).to("cuda")
    model.eval()
    model.generation_config.use_cache = True
    EOS_TOKEN_ID = tokenizer.eos_token_id


def _strip_thinking(text: str) -> str:
    if not isinstance(text, str):
        return ""
    return _THINK_RE.sub("", text).strip()


def _parse_think(raw: str) -> tuple[str, str, bool]:
    """Return (think, answer, think_open) from a possibly streaming reply."""
    open_idx = raw.find(THINK_OPEN)
    if open_idx == -1:
        # A partially streamed opening tag is not an answer yet.
        if raw.strip() and THINK_OPEN.startswith(raw.strip()):
            return "", "", False
        return "", raw, False
    rest = raw[open_idx + len(THINK_OPEN) :]
    close_idx = rest.find(THINK_CLOSE)
    before = raw[:open_idx]
    if close_idx == -1:
        # Hide a partially streamed closing tag at the end of the think block.
        for i in range(len(THINK_CLOSE) - 1, 0, -1):
            if rest.endswith(THINK_CLOSE[:i]):
                rest = rest[:-i]
                break
        return rest.strip(), before.strip(), True
    think = rest[:close_idx].strip()
    answer = (before + rest[close_idx + len(THINK_CLOSE) :]).strip()
    return think, answer, False


def _gpu_seconds(history, max_new_tokens=1024, *args, **kwargs):
    # ZeroGPU adds its own startup overhead on top of this request, and
    # anonymous visitors have a small quota — keep the ask modest.
    tokens = int(max_new_tokens or 1024)
    return min(90, max(25, 10 + tokens // 30))


def _history_for_model(history: list) -> list[dict]:
    conversation = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in history or []:
        role = msg.get("role", "user")
        meta = msg.get("metadata") or {}
        # Only skip thought/tool bubbles; Gradio attaches an empty metadata
        # dict to ordinary messages, so a bare truthiness check drops them all.
        is_thought = isinstance(meta, dict) and bool(meta.get("title"))
        if role == "system" or is_thought:
            continue
        content = msg.get("content", "")
        if isinstance(content, list):
            # Gradio 6 sends message content as a list of blocks.
            parts = []
            for block in content:
                if isinstance(block, str):
                    parts.append(block)
                elif isinstance(block, dict):
                    text = block.get("text") or block.get("content") or ""
                    if isinstance(text, str) and text:
                        parts.append(text)
            content = "\n".join(parts)
        if role == "assistant":
            content = _strip_thinking(content)
        if content:
            conversation.append({"role": role, "content": content})
    return conversation


def _stream_tokens(conversation: list[dict], max_new_tokens: int, temperature: float, top_p: float, enable_thinking: bool) -> Iterator[str]:
    if PREVIEW:
        import time

        demo = (
            f"{THINK_OPEN}\n"
            "Check the premises, then the conclusion. Premise 1 gives Rain -> Wet. "
            "Premise 2 observes Wet. Inferring Rain from Wet affirms the consequent, "
            "which is invalid. The argument is a classic Barbara syllogism otherwise.\n"
            f"{THINK_CLOSE}\n\n"
            "Answer: entailment."
        )
        acc = ""
        for ch in demo:
            acc += ch
            time.sleep(0.02)
            yield acc
        return

    encoded = tokenizer.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        enable_thinking=enable_thinking,
    )
    input_ids = encoded["input_ids"]
    attention_mask = encoded["attention_mask"]
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        attention_mask = attention_mask[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed the conversation to the last {MAX_INPUT_TOKEN_LENGTH} tokens.")

    input_ids = input_ids.to(model.device)
    attention_mask = attention_mask.to(model.device)
    streamer = TextIteratorStreamer(
        tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True
    )
    generate_kwargs = dict(
        input_ids=input_ids,
        attention_mask=attention_mask,
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        num_beams=1,
        use_cache=True,
        eos_token_id=EOS_TOKEN_ID,
        pad_token_id=tokenizer.pad_token_id or EOS_TOKEN_ID,
    )
    if temperature == 0:
        generate_kwargs["do_sample"] = False
    else:
        generate_kwargs["do_sample"] = True
        generate_kwargs["temperature"] = float(temperature)
        generate_kwargs["top_p"] = float(top_p)

    Thread(target=model.generate, kwargs=generate_kwargs, daemon=True).start()
    acc = ""
    for text in streamer:
        acc += text
        yield acc


def queue_message(message, history):
    """Paint the user turn immediately so the GPU wait is not a blank screen."""
    message = (message or "").strip()
    if not message:
        return "", history, gr.skip()
    history = [dict(m) if isinstance(m, dict) else m for m in (history or [])]
    history.append({"role": "user", "content": message})
    history.append(
        {
            "role": "assistant",
            "content": "",
            "metadata": {"title": "Logicizing…", "status": "pending"},
        }
    )
    return "", history, gr.update(visible=False)


@spaces.GPU(duration=_gpu_seconds)
def generate_reply(history, max_new_tokens, temperature, top_p, enable_thinking):
    history = [dict(m) if isinstance(m, dict) else m for m in (history or [])]
    if not any(m.get("role") == "user" for m in history):
        yield history
        return
    conversation = _history_for_model(history)

    # Stream only into a bubble at the END of the history. Reaching for any
    # older assistant message would overwrite a previous turn in place.
    think_msg = None
    last = history[-1] if history else None
    if (
        last is not None
        and last.get("role") == "assistant"
        and isinstance(last.get("metadata"), dict)
        and last["metadata"].get("status") == "pending"
    ):
        think_msg = last
        think_msg["content"] = ""
    elif last is not None and last.get("role") == "user":
        # Gradio strips the empty pending bubble in preprocessing; recreate it.
        think_msg = {"role": "assistant", "content": ""}
        history.append(think_msg)
    else:
        # No new user turn (e.g. an empty submit); don't regenerate old answers.
        yield history
        return
    think_msg["metadata"] = {"title": "Logicizing…", "status": "pending"}
    yield history

    answer_msg = {"role": "assistant", "content": ""}

    for raw in _stream_tokens(
        conversation, max_new_tokens, temperature, top_p, enable_thinking
    ):
        think, answer, think_open = _parse_think(raw)
        if think or think_open:
            think_msg["content"] = think
            think_msg["metadata"] = {
                "title": "Logicizing…" if think_open else "Logic",
                "status": "pending" if think_open else "done",
            }
        if answer:
            answer_msg["content"] = answer
            if not any(m is answer_msg for m in history):
                history.append(answer_msg)
        elif not think_open and not think and raw.strip():
            answer_msg["content"] = raw.strip()
            if not any(m is answer_msg for m in history):
                history.append(answer_msg)
        elif any(m is answer_msg for m in history):
            # A partial tag was mistaken for an answer earlier; retract it.
            history.remove(answer_msg)
            answer_msg["content"] = ""
        yield history

    if think_msg["content"]:
        think_msg["metadata"] = {"title": "Logic", "status": "done"}
        yield history
    elif think_msg in history and not think_msg["content"]:
        history.remove(think_msg)
        yield history


with gr.Blocks(fill_height=True, fill_width=True, elem_id="app-root") as demo:
    with gr.Column(elem_id="app-shell"):
        gr.HTML(HEADER)
        with gr.Tabs(elem_id="twil-tabs") as tabs:
            with gr.Tab("Chat", id="chat", elem_id="tab-chat"):
                chatbot = gr.Chatbot(
                    value=[],
                    placeholder=PLACEHOLDER,
                    label="",
                    show_label=False,
                    layout="bubble",
                    avatar_images=(None, str(AVATAR) if AVATAR.exists() else None),
                    allow_tags=["think"],
                    reasoning_tags=[(THINK_OPEN, THINK_CLOSE)],
                    buttons=None,
                    editable=False,
                    line_breaks=False,
                    elem_id="twil-chat",
                    scale=1,
                    height="100%",
                )
                with gr.Row(elem_id="chips-row") as chips_row:
                    chip_btns = [
                        gr.Button(label, variant="secondary", size="sm", elem_classes=["twil-chip"])
                        for label, _ in CHIPS
                    ]
                with gr.Row(elem_id="composer-row"):
                    prompt = gr.Textbox(
                        placeholder="Ask webAI anything…",
                        show_label=False,
                        container=False,
                        lines=2,
                        max_lines=2,
                        elem_id="twil-input",
                        scale=8,
                    )
                    send = gr.Button("↑", elem_id="send-btn", scale=0)
                with gr.Accordion("Parameters", open=False, elem_id="params-box"):
                    max_new_tokens = gr.Slider(256, 4096, value=1024, step=256, label="Max new tokens")
                    temperature = gr.Slider(0, 1.5, value=0, step=0.05, label="Temperature (0 = greedy)")
                    top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
                    enable_thinking = gr.Checkbox(value=True, label="Enable thinking")
                gr.HTML('<p class="twil-disclaimer">This is AI and it can make mistakes</p>')
            with gr.Tab("Showcase", id="showcase", elem_id="tab-showcase"):
                gr.HTML(SHOWCASE_INTRO)
                card_btns = []
                for start in range(0, len(SHOWCASE), 3):
                    with gr.Row(elem_classes=["twil-card-row"]):
                        for case in SHOWCASE[start : start + 3]:
                            with gr.Column(elem_classes=["twil-card"], scale=1, min_width=220):
                                gr.HTML(_card_html(case))
                                card_btns.append(
                                    gr.Button("Analyze it", size="sm", elem_classes=["twil-card-btn"])
                                )

    # Lock the composer while a turn is running: overlapping events make the
    # finishing stream overwrite the chat with its own stale history.
    lockable = [prompt, send, *card_btns]

    def _lock():
        return [gr.update(interactive=False)] * len(lockable)

    def _unlock():
        return [gr.update(interactive=True)] * len(lockable)

    def _wire(event, prep_fn=None, prep_outputs=None):
        chain = event(_lock, None, lockable, show_progress="hidden")
        if prep_fn is not None:
            chain = chain.then(prep_fn, outputs=prep_outputs, show_progress="hidden")
        chain = chain.then(
            queue_message,
            [prompt, chatbot],
            [prompt, chatbot, chips_row],
            show_progress="hidden",
        ).then(
            generate_reply,
            [chatbot, max_new_tokens, temperature, top_p, enable_thinking],
            chatbot,
            concurrency_limit=1,
            show_progress="hidden",
        )
        chain.then(_unlock, None, lockable, show_progress="hidden")

    _wire(send.click)
    _wire(prompt.submit)
    for btn, (_, template) in zip(chip_btns, CHIPS):
        btn.click(lambda t=template: t, outputs=prompt)

    for btn, case in zip(card_btns, SHOWCASE):
        case_prompt = SHOWCASE_TEMPLATE.format(argument=case["quote"])
        _wire(
            btn.click,
            prep_fn=lambda p=case_prompt: (p, gr.Tabs(selected="chat")),
            prep_outputs=[prompt, tabs],
        )

if __name__ == "__main__":
    demo.launch(
        theme=THEME,
        css=CSS,
        head=HEAD,
        allowed_paths=[str(ASSETS), str(ASSETS.parent)],
        ssr_mode=False,
        show_error=True,
    )