"""Cosmos T2-Accelerate-beta — Gradio Chat Demo

Standalone inference app generated by the Cosmos T2-Accelerate-beta universal training notebook.
It matches the notebook architecture: RoPE, RMSNorm, SwiGLU, GQA, and a configurable Engram memory path.
"""

import contextlib
import math
import re
import threading
from pathlib import Path

import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

MODEL_REPO_ID = "wop/Cosmos-T2-Accelerate-beta"
CHECKPOINT_NAME = "Cosmos-T2-Accelerate-beta.pt"
TOKENIZER_NAME = "Qwen/Qwen2.5-0.5B"
MODEL_NAME = "Cosmos T2-Accelerate-beta"
DEFAULT_SYSTEM_PROMPT = "Enable thinking features: INTUITION"
MAX_CTX_HARD = 1028
MAX_NEW_HARD = 1028
USE_KV_CACHE = True

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

# ---------------------------------------------------------------------------
# Architecture
# ---------------------------------------------------------------------------

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    def forward(self, x):
        rms = x.pow(2).mean(dim=-1, keepdim=True)
        x = x * torch.rsqrt(rms + self.eps)
        return x * self.weight

def rotate_half(x):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)

def apply_rope(q, k, cos, sin):
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)

class GQAAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, rope_base=10000, dropout=0.0):
        super().__init__()
        assert d_model % n_heads == 0
        assert n_heads % n_kv_heads == 0
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        self.dropout = dropout
        self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, rope_cos, rope_sin, past_kv=None, use_cache=False):
        batch, seq_len, _ = x.shape
        q = self.q_proj(x).view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
        q, k = apply_rope(q, k, rope_cos, rope_sin)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        present_kv = (k, v) if use_cache else None

        if self.n_kv_heads != self.n_heads:
            repeat = self.n_heads // self.n_kv_heads
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)

        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=(past_kv is None), dropout_p=self.dropout if self.training else 0.0
        )
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        out = self.o_proj(out)
        return (out, present_kv) if use_cache else out

class SwiGLUMLP(nn.Module):
    def __init__(self, d_model, hidden_dim, dropout=0.0):
        super().__init__()
        self.gate = nn.Linear(d_model, hidden_dim, bias=False)
        self.up = nn.Linear(d_model, hidden_dim, bias=False)
        self.down = nn.Linear(hidden_dim, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        return self.down(self.dropout(F.silu(self.gate(x)) * self.up(x)))

class EngramMemory(nn.Module):
    def __init__(self, d_model, bucket_count, memory_dim, order, pad_id=0, dropout=0.0):
        super().__init__()
        self.bucket_count = bucket_count
        self.order = order
        self.pad_id = pad_id
        self.bucket = nn.Embedding(bucket_count, memory_dim)
        self.query = nn.Linear(d_model, memory_dim, bias=False)
        self.project = nn.Linear(memory_dim, d_model, bias=False)
        self.gate = nn.Linear(d_model, d_model, bias=True)
        self.dropout = nn.Dropout(dropout)
        primes = [1, 1315423911, 2654435761, 97531, 433494437]
        self.register_buffer("primes", torch.tensor(primes[:order], dtype=torch.long), persistent=False)

    def hash_tokens(self, idx):
        batch, seq_len = idx.shape
        pad = torch.full((batch, self.order - 1), self.pad_id, device=idx.device, dtype=idx.dtype)
        history = torch.cat([pad, idx], dim=1)
        hashed = torch.zeros((batch, seq_len), device=idx.device, dtype=torch.long)
        for offset in range(self.order):
            slice_ = history[:, offset: offset + seq_len].long()
            hashed = (hashed * 1315423911 + slice_ * self.primes[offset]) % self.bucket_count
        return hashed

    def forward(self, x, idx):
        hashed = self.hash_tokens(idx)
        if hashed.size(1) != x.size(1):
            hashed = hashed[:, -x.size(1):]
        query = torch.tanh(self.query(x))
        memory = self.bucket(hashed) * query
        memory = self.project(memory)
        gate = torch.sigmoid(self.gate(x))
        return self.dropout(gate * memory)

class Block(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, d_ff, rope_base, dropout=0.0,
                 use_engram=False, engram_bucket_count=64, engram_dim=16, engram_order=3, pad_id=0):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.attn = GQAAttention(d_model, n_heads, n_kv_heads, rope_base=rope_base, dropout=dropout)
        self.norm2 = RMSNorm(d_model)
        self.engram = EngramMemory(d_model, engram_bucket_count, engram_dim, engram_order,
                                   pad_id=pad_id, dropout=dropout) if use_engram else None
        self.norm3 = RMSNorm(d_model)
        self.mlp = SwiGLUMLP(d_model, d_ff, dropout=dropout)

    def forward(self, x, idx, rope_cos, rope_sin):
        x = x + self.attn(self.norm1(x), rope_cos, rope_sin)
        if self.engram is not None:
            x = x + self.engram(self.norm2(x), idx)
        return x + self.mlp(self.norm3(x))

    def forward_cached(self, x, idx_context, rope_cos, rope_sin, past_kv=None):
        attn_out, present_kv = self.attn(self.norm1(x), rope_cos, rope_sin, past_kv=past_kv, use_cache=True)
        x = x + attn_out
        if self.engram is not None:
            x = x + self.engram(self.norm2(x), idx_context)
        x = x + self.mlp(self.norm3(x))
        return x, present_kv

class CosmosT2_Accelerate_LLM(nn.Module):
    def __init__(self, vocab_size, d_model=32, n_layers=6, n_heads=2, n_kv_heads=1,
                 d_ff=256, max_len=1028, rope_base=10000, dropout=0.05, use_engram=True,
                 engram_every=2, engram_bucket_count=64, engram_dim=16, engram_order=3, pad_id=0):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        self.max_len = max_len
        self.rope_base = rope_base
        self.pad_id = pad_id
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.ModuleList()
        for layer_index in range(n_layers):
            block_uses_engram = use_engram and ((layer_index + 1) % engram_every == 0)
            self.blocks.append(Block(
                d_model=d_model, n_heads=n_heads, n_kv_heads=n_kv_heads, d_ff=d_ff,
                rope_base=rope_base, dropout=dropout, use_engram=block_uses_engram,
                engram_bucket_count=engram_bucket_count, engram_dim=engram_dim,
                engram_order=engram_order, pad_id=pad_id,
            ))
        self.norm_f = RMSNorm(d_model)
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def build_rope(self, seq_len, device, dtype, start_pos=0):
        inv_freq = 1.0 / (self.rope_base ** (
            torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim
        ))
        positions = torch.arange(start_pos, start_pos + seq_len, device=device, dtype=torch.float32)
        freqs = torch.outer(positions, inv_freq)
        cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype)[None, None, :, :]
        sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype)[None, None, :, :]
        return cos, sin

    def forward(self, idx, targets=None):
        if idx.size(1) > self.max_len:
            idx = idx[:, -self.max_len:]
        seq_len = idx.size(1)
        rope_cos, rope_sin = self.build_rope(seq_len, idx.device, self.tok_emb.weight.dtype)
        x = self.tok_emb(idx)
        for block in self.blocks:
            x = block(x, idx, rope_cos, rope_sin)
        x = self.norm_f(x)
        logits = F.linear(x, self.tok_emb.weight)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss

    def trim_kv_cache(self, past_kv, max_tokens):
        if past_kv is None:
            return None
        max_tokens = max(0, int(max_tokens))
        trimmed = []
        for k, v in past_kv:
            if max_tokens == 0:
                k = k[:, :, :0, :].contiguous()
                v = v[:, :, :0, :].contiguous()
            elif k.size(2) > max_tokens:
                k = k[:, :, -max_tokens:, :].contiguous()
                v = v[:, :, -max_tokens:, :].contiguous()
            trimmed.append((k, v))
        return trimmed

    @torch.no_grad()
    def forward_cached(self, idx, past_kv=None, cache_pos=0, max_ctx=None, idx_context=None):
        self.eval()
        max_ctx = self.max_len if max_ctx is None else int(max_ctx)
        if past_kv is None:
            idx = idx[:, -max_ctx:]
            idx_context = idx
            cache_pos = 0
        else:
            keep_past = max(0, max_ctx - idx.size(1))
            past_kv = self.trim_kv_cache(past_kv, keep_past)
            idx_context = idx if idx_context is None else idx_context[:, -max_ctx:]

        seq_len = idx.size(1)
        rope_cos, rope_sin = self.build_rope(
            seq_len,
            idx.device,
            self.tok_emb.weight.dtype,
            start_pos=cache_pos,
        )
        x = self.tok_emb(idx)
        present_kv = []
        for layer_index, block in enumerate(self.blocks):
            layer_past = None if past_kv is None else past_kv[layer_index]
            x, layer_present = block.forward_cached(x, idx_context, rope_cos, rope_sin, past_kv=layer_past)
            present_kv.append(layer_present)
        x = self.norm_f(x)
        logits = F.linear(x, self.tok_emb.weight)
        return logits, present_kv, cache_pos + seq_len

    def sample_next(self, logits, temperature=0.8, top_k=50):
        if logits.dim() == 3:
            logits = logits[:, -1, :]
        if temperature <= 1e-6:
            return torch.argmax(logits, dim=-1, keepdim=True)
        logits = logits / temperature
        if top_k and top_k > 0:
            values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits = logits.masked_fill(logits < values[:, [-1]], float("-inf"))
        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)

    @torch.no_grad()
    def prefill_cache(self, idx, max_ctx=None):
        logits, past_kv, cache_pos = self.forward_cached(idx, past_kv=None, cache_pos=0, max_ctx=max_ctx)
        return logits[:, -1, :], past_kv, cache_pos

    @torch.no_grad()
    def decode_cached(self, idx, past_kv, cache_pos, idx_context, max_ctx=None):
        logits, past_kv, cache_pos = self.forward_cached(
            idx,
            past_kv=past_kv,
            cache_pos=cache_pos,
            max_ctx=max_ctx,
            idx_context=idx_context,
        )
        return logits[:, -1, :], past_kv, cache_pos

    @torch.no_grad()
    def generate_step(self, idx, temperature=0.8, top_k=50, max_ctx=None):
        max_ctx = self.max_len if max_ctx is None else max_ctx
        logits, _, _ = self.prefill_cache(idx[:, -max_ctx:], max_ctx=max_ctx)
        return self.sample_next(logits, temperature=temperature, top_k=top_k)

    @torch.no_grad()
    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, max_ctx=None, stop_ids=None):
        """Autoregressive generation with KV cache. Stops on stop_ids (defaults to STOP_IDS)."""
        self.eval()
        if stop_ids is None:
            stop_ids = globals().get("STOP_IDS")
        max_ctx = self.max_len if max_ctx is None else int(max_ctx)
        idx = idx[:, -max_ctx:]
        logits, past_kv, cache_pos = self.prefill_cache(idx, max_ctx=max_ctx)
        for step in range(max_new_tokens):
            nxt = self.sample_next(logits, temperature=temperature, top_k=top_k)
            if stop_ids and nxt.numel() == 1 and int(nxt.item()) in stop_ids:
                break
            idx = torch.cat([idx, nxt], dim=1)
            if step + 1 < max_new_tokens:
                logits, past_kv, cache_pos = self.decode_cached(
                    nxt, past_kv, cache_pos, idx[:, -max_ctx:], max_ctx=max_ctx)
        return idx

# ---------------------------------------------------------------------------
# Checkpoint loader
# ---------------------------------------------------------------------------

def load_checkpoint(tokenizer):
    ckpt_path = Path(CHECKPOINT_NAME)
    if not ckpt_path.exists():
        print(f"Downloading {CHECKPOINT_NAME} from {MODEL_REPO_ID}")
        ckpt_path = Path(hf_hub_download(repo_id=MODEL_REPO_ID, filename=CHECKPOINT_NAME))
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    cfg = ckpt.get("config", {})
    resolved = {
        "vocab_size":          cfg.get("vocab_size", len(tokenizer)),
        "d_model":             cfg.get("d_model", 32),
        "n_layers":            cfg.get("n_layers", 6),
        "n_heads":             cfg.get("n_heads", 2),
        "n_kv_heads":          cfg.get("n_kv_heads", 1),
        "d_ff":                cfg.get("d_ff", 256),
        "max_len":             cfg.get("max_len", 1028),
        "rope_base":           cfg.get("rope_base", 10000),
        "dropout":             0.0,
        "use_engram":          cfg.get("use_engram", True),
        "engram_every":        cfg.get("engram_every", 2),
        "engram_bucket_count": cfg.get("engram_bucket_count", 64),
        "engram_dim":          cfg.get("engram_dim", 16),
        "engram_order":        cfg.get("engram_order", 3),
        "pad_id":              tokenizer.pad_token_id,
    }
    print(f"Model config: {resolved}")
    model = CosmosT2_Accelerate_LLM(**resolved)
    state = ckpt.get("model_state", ckpt)
    state = {k.replace("module.", "", 1): v for k, v in state.items()}
    missing, unexpected = model.load_state_dict(state, strict=False)
    if missing:    print(f"Missing keys:    {missing}")
    if unexpected: print(f"Unexpected keys: {unexpected}")
    model = model.to(DEVICE).to(DTYPE).eval()
    return model

print(f"Device: {DEVICE} | dtype: {DTYPE}")
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
print(f"Tokenizer: {TOKENIZER_NAME}")
model = load_checkpoint(tokenizer)
n_params = sum(p.numel() for p in model.parameters())
_n_layers = model.n_layers if hasattr(model, "n_layers") else "?"
print(f"Loaded {n_params / 1e6:.2f}M parameters")
EOS_ID = tokenizer.eos_token_id

def _resolve_stop_ids(tok):
    ids = set()
    for t in ("<|im_end|>", "<|endoftext|>"):
        i = tok.convert_tokens_to_ids(t)
        if isinstance(i, int) and i >= 0 and i != tok.unk_token_id:
            ids.add(i)
    if tok.eos_token_id is not None:
        ids.add(tok.eos_token_id)
    return ids

STOP_IDS = _resolve_stop_ids(tokenizer)

# ---------------------------------------------------------------------------
# Text → HTML rendering  (preserves <think> blocks)
# ---------------------------------------------------------------------------

def render_message_html(text: str, is_streaming: bool = False) -> str:
    """
    Convert raw model output to HTML.
    - <think>...</think>  →  collapsible details block
    - Incomplete <think>  →  open/pulsing block while streaming
    - Remaining text     →  escaped, newlines → <br>
    """
    # Split on think tags
    # Patterns: complete <think>...</think> or dangling <think>...
    parts = []
    cursor = 0
    pattern = re.compile(r'<think>(.*?)</think>', re.DOTALL)

    for m in pattern.finditer(text):
        # text before this think block
        before = text[cursor:m.start()]
        if before:
            parts.append(('text', before))
        parts.append(('think_done', m.group(1)))
        cursor = m.end()

    tail = text[cursor:]
    # check for an open, unclosed <think>
    open_match = re.search(r'<think>(.*)', tail, re.DOTALL)
    if open_match:
        before_open = tail[:open_match.start()]
        if before_open:
            parts.append(('text', before_open))
        parts.append(('think_open', open_match.group(1)))
    else:
        if tail:
            parts.append(('text', tail))

    html_parts = []
    for kind, content in parts:
        if kind == 'text':
            escaped = content.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
            escaped = escaped.replace('\n', '<br>')
            html_parts.append(f'<span class="msg-text">{escaped}</span>')
        elif kind == 'think_done':
            inner = content.strip().replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
            inner = inner.replace('\n', '<br>')
            html_parts.append(
                f'<details class="think-block">'
                f'<summary>💭 Thinking</summary>'
                f'<div class="think-content">{inner}</div>'
                f'</details>'
            )
        elif kind == 'think_open':
            inner = content.strip().replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
            inner = inner.replace('\n', '<br>')
            pulse = ' think-streaming' if is_streaming else ''
            html_parts.append(
                f'<details class="think-block{pulse}" open>'
                f'<summary>💭 Thinking{"…" if is_streaming else ""}</summary>'
                f'<div class="think-content">{inner}</div>'
                f'</details>'
            )

    return ''.join(html_parts)


def build_prompt(history, user_msg, system_msg):
    messages = []
    if system_msg and system_msg.strip():
        messages.append({"role": "system", "content": system_msg.strip()})
    for item in history:
        if isinstance(item, dict) and "role" in item and "content" in item:
            # history stores raw text, not html
            messages.append({"role": item["role"], "content": item["content"]})
    messages.append({"role": "user", "content": user_msg})
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


# ---------------------------------------------------------------------------
# Per-session stop flag
# ---------------------------------------------------------------------------

_stop_flags: dict[str, threading.Event] = {}

def get_stop_event(session_id: str) -> threading.Event:
    if session_id not in _stop_flags:
        _stop_flags[session_id] = threading.Event()
    return _stop_flags[session_id]


# ---------------------------------------------------------------------------
# CSS
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
html, body { overflow-x: hidden; max-width: 100vw; }

.gradio-container {
    max-width: 980px !important;
    width: 100% !important;
    margin: auto;
    padding: 12px !important;
    box-sizing: border-box;
}

/* ---- header ---- */
#header-card {
    background: linear-gradient(135deg, #0d1117 0%, #161b22 100%);
    border: 1px solid #30363d;
    border-radius: 14px;
    padding: 20px 24px;
    margin-bottom: 12px;
    text-align: center;
}
#header-card h2 { color: #58a6ff; margin: 4px 0 8px; font-weight: 700; font-size: 1.3em; }
#header-card p  { color: #8b949e; margin: 3px 0; font-size: 0.88em; }
.badge {
    display: inline-block;
    background: #21262d;
    color: #c9d1d9;
    padding: 2px 9px;
    border-radius: 999px;
    font-size: 0.76em;
    margin: 2px 3px;
    border: 1px solid #30363d;
}
.warn {
    background: #2a1f0a;
    border: 1px solid #6b4d11;
    color: #f0c674;
    padding: 8px 12px;
    border-radius: 8px;
    font-size: 0.84em;
    margin-top: 8px;
    text-align: left;
}

/* ---- chatbox ---- */
#chatbox {
    border: 1px solid #30363d;
    border-radius: 12px;
    background: #0d1117;
    padding: 16px;
    min-height: 420px;
    max-height: 60vh;
    overflow-y: auto;
    overflow-x: hidden;
    display: flex;
    flex-direction: column;
    gap: 14px;
    /* no forced scroll — user controls it */
    scroll-behavior: smooth;
}

/* ---- messages ---- */
.msg-row { display: flex; gap: 10px; align-items: flex-start; max-width: 100%; }
.msg-row.user  { flex-direction: row-reverse; }
.msg-row.asst  { flex-direction: row; }

.msg-avatar {
    width: 30px; height: 30px; border-radius: 50%;
    display: flex; align-items: center; justify-content: center;
    font-size: 0.85em; flex-shrink: 0; margin-top: 2px;
}
.msg-row.user .msg-avatar { background: #1f6feb; color: #fff; }
.msg-row.asst .msg-avatar { background: #21262d; color: #58a6ff; border: 1px solid #30363d; }

.msg-bubble {
    max-width: 78%;
    padding: 10px 14px;
    border-radius: 14px;
    font-size: 0.93em;
    line-height: 1.6;
    word-break: break-word;
    overflow-wrap: break-word;
}
.msg-row.user .msg-bubble {
    background: #1f6feb;
    color: #fff;
    border-bottom-right-radius: 4px;
}
.msg-row.asst .msg-bubble {
    background: #161b22;
    color: #c9d1d9;
    border: 1px solid #30363d;
    border-bottom-left-radius: 4px;
}

/* ---- think blocks ---- */
.think-block {
    margin: 6px 0;
    border: 1px solid #30363d;
    border-radius: 8px;
    overflow: hidden;
    background: #0d1117;
}
.think-block summary {
    cursor: pointer;
    padding: 5px 10px;
    font-size: 0.82em;
    color: #8b949e;
    background: #161b22;
    user-select: none;
    list-style: none;
}
.think-block summary::-webkit-details-marker { display: none; }
.think-block summary::before { content: "▶ "; font-size: 0.7em; }
.think-block[open] summary::before { content: "▼ "; }
.think-content {
    padding: 8px 12px;
    font-size: 0.82em;
    color: #6e7681;
    font-family: monospace;
    white-space: pre-wrap;
    line-height: 1.5;
    max-height: 280px;
    overflow-y: auto;
}
@keyframes pulse-border {
    0%   { border-color: #388bfd44; }
    50%  { border-color: #388bfd; }
    100% { border-color: #388bfd44; }
}
.think-streaming {
    animation: pulse-border 1.4s ease-in-out infinite;
}

/* ---- token counter ---- */
#token-counter {
    font-size: 0.76em;
    color: #6e7681;
    text-align: right;
    padding: 2px 4px;
}

/* ---- input row ---- */
#input-row {
    display: flex;
    gap: 8px;
    align-items: flex-end;
    margin-top: 8px;
}
#user-input textarea {
    background: #0d1117 !important;
    border: 1px solid #30363d !important;
    border-radius: 10px !important;
    color: #c9d1d9 !important;
    resize: none;
    font-size: 0.93em;
}
#user-input textarea:focus {
    border-color: #388bfd !important;
    outline: none !important;
}
#send-btn, #stop-btn, #clear-btn {
    min-width: 72px !important;
    height: 42px !important;
    border-radius: 10px !important;
    font-size: 0.88em !important;
    font-weight: 600 !important;
}
#send-btn  { background: #1f6feb !important; color: #fff !important; }
#send-btn:hover  { background: #388bfd !important; }
#stop-btn  { background: #da3633 !important; color: #fff !important; }
#stop-btn:hover  { background: #f85149 !important; }
#clear-btn { background: #21262d !important; color: #c9d1d9 !important; border: 1px solid #30363d !important; }
#clear-btn:hover { background: #30363d !important; }

/* ---- params accordion ---- */
#params-accordion { margin-top: 10px; }

footer { visibility: hidden; }

@media (max-width: 640px) {
    .gradio-container { padding: 6px !important; }
    .msg-bubble { max-width: 90%; font-size: 0.88em; }
    #chatbox { min-height: 320px; max-height: 55vh; }
}
"""

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

_n_layers_str = str(_n_layers)

HEADER_HTML = f"""
<div id="header-card">
  <h2>{MODEL_NAME}</h2>
  <p>
    <span class="badge">{n_params / 1e6:.2f}M params</span>
    <span class="badge">{_n_layers_str} layers</span>
    <span class="badge">RoPE + RMSNorm + SwiGLU + GQA</span>
    <span class="badge">Engram on</span>
    <span class="badge">{DEVICE.upper()}</span>
  </p>
  <p>
    Trained on <a href="https://huggingface.co/datasets/wop/XXXXXL-chain-of-thought" target="_blank" style="color:#58a6ff">wop/XXXXXL-chain-of-thought</a>
    &nbsp;·&nbsp;
    <a href="https://huggingface.co/wop/Cosmos-T2-Accelerate-beta" target="_blank" style="color:#58a6ff">Model repo</a>
  </p>
  <div class="warn">⚠️ Research/demo model — small, may hallucinate freely. Keep temperature low for stable outputs.</div>
</div>
"""

with gr.Blocks(
    theme=gr.themes.Base(),
    css=CUSTOM_CSS,
    title=f"{MODEL_NAME} Demo",
) as demo:

    # ---- state ----
    # history: list of {"role": "user"|"assistant", "content": <raw text>}
    history_state = gr.State([])
    session_id    = gr.State("")   # unique per browser tab

    # ---- layout ----
    gr.HTML(HEADER_HTML)

    chatbox = gr.HTML(
        value='<div id="chatbox"><div style="color:#6e7681;text-align:center;margin:auto">Send a message to start…</div></div>',
        elem_id="chatbox-wrapper",
    )

    token_counter = gr.Markdown("", elem_id="token-counter")

    with gr.Row(elem_id="input-row"):
        user_input = gr.Textbox(
            placeholder="Ask Cosmos T2-Accelerate-beta anything… (Shift+Enter for newline)",
            lines=1, max_lines=6,
            show_label=False,
            elem_id="user-input",
            scale=8,
        )
        send_btn  = gr.Button("Send",  elem_id="send-btn",  variant="primary", scale=1)
        stop_btn  = gr.Button("Stop",  elem_id="stop-btn",  variant="stop",    scale=1)
        clear_btn = gr.Button("Clear", elem_id="clear-btn",                    scale=1)

    with gr.Accordion("⚙️ Generation parameters", open=False, elem_id="params-accordion"):
        with gr.Row():
            system_box   = gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, label="System prompt", lines=2, scale=3)
            temperature  = gr.Slider(0.0, 2.0,       value=0.1,           step=0.05, label="Temperature", scale=1)
            top_k        = gr.Slider(1,   200,        value=1,             step=1,    label="Top-K",        scale=1)
        with gr.Row():
            ctx_size     = gr.Slider(64,  MAX_CTX_HARD, value=MAX_CTX_HARD, step=64,  label="Context window")
            max_new      = gr.Slider(16,  MAX_NEW_HARD,  value=128,          step=16,  label="Max new tokens")

    # ---- rendering helper ----
    def history_to_html(history):
        if not history:
            return '<div id="chatbox"><div style="color:#6e7681;text-align:center;margin:auto">Send a message to start…</div></div>'
        rows = []
        for msg in history:
            role = msg["role"]
            raw  = msg["content"]
            if role == "user":
                escaped = raw.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;').replace('\n', '<br>')
                bubble  = f'<div class="msg-bubble">{escaped}</div>'
                avatar  = '<div class="msg-avatar">👤</div>'
                rows.append(f'<div class="msg-row user">{avatar}{bubble}</div>')
            else:
                html    = render_message_html(raw, is_streaming=False)
                bubble  = f'<div class="msg-bubble">{html}</div>'
                avatar  = '<div class="msg-avatar">✦</div>'
                rows.append(f'<div class="msg-row asst">{avatar}{bubble}</div>')
        inner = "\n".join(rows)
        # scroll-to-bottom js trick: invisible anchor at end
        return f'<div id="chatbox">{inner}<div id="chat-end"></div></div><script>document.getElementById("chat-end")?.scrollIntoView({{behavior:"smooth",block:"end"}});</script>'

    # ---- send / stream ----
    def do_send(message, history, sys_msg, temp, tk, ctx, mn, sid):
        if not message or not message.strip():
            yield history, history_to_html(history), "", ""
            return

        # Assign session id if blank
        if not sid:
            import uuid
            sid = str(uuid.uuid4())

        stop_evt = get_stop_event(sid)
        stop_evt.clear()

        # Append user message
        history = list(history) + [{"role": "user", "content": message.strip()}]
        # Show user message immediately, assistant bubble loading
        loading_html = history_to_html(history + [{"role": "assistant", "content": "▌"}])
        yield history, loading_html, "", sid

        # Build prompt
        prompt    = build_prompt(history[:-1], message.strip(), sys_msg)
        input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(DEVICE)
        ctx       = int(min(max(int(ctx), 8), MAX_CTX_HARD))
        if input_ids.shape[1] > ctx - 16:
            input_ids = input_ids[:, -(ctx - 16):]
        mn   = int(min(max(int(mn), 1), MAX_NEW_HARD))
        temp = float(max(min(temp, 2.0), 0.0))
        tk   = int(max(tk, 1))

        cur_ids   = input_ids
        generated = []
        partial   = ""
        n_tokens  = 0
        kv_cache  = None
        cache_pos = 0
        next_logits = None

        if USE_KV_CACHE:
            next_logits, kv_cache, cache_pos = model.prefill_cache(cur_ids, max_ctx=ctx)

        for step in range(mn):
            if stop_evt.is_set():
                break
            if USE_KV_CACHE:
                nxt = model.sample_next(next_logits, temperature=temp, top_k=tk)
            else:
                nxt = model.generate_step(cur_ids, temperature=temp, top_k=tk, max_ctx=ctx)
            cur_ids  = torch.cat([cur_ids, nxt], dim=1)
            token_id = int(nxt.item())
            if token_id in STOP_IDS:
                break
            generated.append(token_id)
            n_tokens += 1

            if USE_KV_CACHE and step + 1 < mn:
                next_logits, kv_cache, cache_pos = model.decode_cached(
                    nxt,
                    kv_cache,
                    cache_pos,
                    cur_ids[:, -ctx:],
                    max_ctx=ctx,
                )
            new_text = tokenizer.decode(generated, skip_special_tokens=False)
            if new_text != partial:
                partial = new_text
                stream_history = history + [{"role": "assistant", "content": partial}]
                # Build streaming html (open think block pulses)
                rows = []
                for msg in stream_history[:-1]:
                    role2 = msg["role"]
                    raw2  = msg["content"]
                    if role2 == "user":
                        esc = raw2.replace('&','&amp;').replace('<','&lt;').replace('>','&gt;').replace('\n','<br>')
                        rows.append(f'<div class="msg-row user"><div class="msg-avatar">👤</div><div class="msg-bubble">{esc}</div></div>')
                    else:
                        h = render_message_html(raw2, is_streaming=False)
                        rows.append(f'<div class="msg-row asst"><div class="msg-avatar">✦</div><div class="msg-bubble">{h}</div></div>')
                # Last (streaming) assistant message
                h_stream = render_message_html(partial, is_streaming=True)
                rows.append(f'<div class="msg-row asst"><div class="msg-avatar">✦</div><div class="msg-bubble">{h_stream}</div></div>')
                inner = "\n".join(rows)
                chat_html = f'<div id="chatbox">{inner}<div id="chat-end"></div></div><script>document.getElementById("chat-end")?.scrollIntoView({{behavior:"instant",block:"end"}});</script>'
                counter   = f"`{n_tokens}` tokens generated"
                yield history, chat_html, counter, sid

        # Finalise
        final_text = tokenizer.decode(generated, skip_special_tokens=False)
        history    = history + [{"role": "assistant", "content": final_text}]
        stopped    = " *(stopped)*" if stop_evt.is_set() else ""
        counter    = f"`{n_tokens}` tokens generated{stopped}"
        yield history, history_to_html(history), counter, sid

    def do_stop(sid):
        if sid:
            get_stop_event(sid).set()

    def do_clear():
        empty_html = '<div id="chatbox"><div style="color:#6e7681;text-align:center;margin:auto">Send a message to start…</div></div>'
        return [], empty_html, "", ""

    # ---- wire up ----
    gen_inputs  = [user_input, history_state, system_box, temperature, top_k, ctx_size, max_new, session_id]
    gen_outputs = [history_state, chatbox, token_counter, session_id]

    send_event = send_btn.click(
        fn=do_send,
        inputs=gen_inputs,
        outputs=gen_outputs,
        show_progress=False,
    )
    # Also trigger on Enter (without shift)
    user_input.submit(
        fn=do_send,
        inputs=gen_inputs,
        outputs=gen_outputs,
        show_progress=False,
    )
    # Clear input box after sending
    send_btn.click(fn=lambda: "", outputs=[user_input])
    user_input.submit(fn=lambda: "", outputs=[user_input])

    stop_btn.click(fn=do_stop, inputs=[session_id], cancels=[send_event])
    clear_btn.click(fn=do_clear, outputs=[history_state, chatbox, token_counter, user_input])

if __name__ == "__main__":
    demo.queue(max_size=16, default_concurrency_limit=1).launch(server_name="0.0.0.0", server_port=7860)