import gradio as gr
import torch
import spaces
from huggingface_hub import hf_hub_download, HfApi
import torch.nn as nn
import torch.nn.functional as F
from transformers import GPT2TokenizerFast
from dataclasses import dataclass
import os
import time

# ── Config ────────────────────────────────────────────────────────────────────
HF_MODEL_REPO = "Bc-AI/nova-1-standard-checkpoints"
HF_TOKEN = os.environ.get("HF_TOKEN")

@dataclass
class NovaConfig:
    vocab_size:   int   = 50304
    d_model:      int   = 2048
    n_heads:      int   = 16
    n_kv_heads:   int   = 8
    n_layers:     int   = 24
    ffn_mult:     float = 8 / 3
    max_len:      int   = 2048
    mod_every_n:  int   = 2
    mod_capacity: float = 0.5
    training_phase: int = 2

    @property
    def dtype(self): return torch.bfloat16
    @property
    def head_dim(self): return self.d_model // self.n_heads
    @property
    def ffn_hidden(self):
        return ((int(self.d_model * self.ffn_mult) + 63) // 64) * 64

# ── Model ─────────────────────────────────────────────────────────────────────
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps   = eps
        self.scale = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        return F.rms_norm(x, self.scale.shape, self.scale, self.eps)

def precompute_freqs_cis(head_dim, max_len, theta=10_000.0, device=None):
    freqs = 1.0 / (theta ** (
        torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim))
    t     = torch.arange(max_len, dtype=torch.float32, device=device)
    freqs = torch.outer(t, freqs)
    return torch.cos(freqs), torch.sin(freqs)

def apply_rope(xq, xk, cos, sin):
    L = xq.shape[1]
    c = torch.cat([cos[:L], cos[:L]], -1).unsqueeze(0).unsqueeze(2)
    s = torch.cat([sin[:L], sin[:L]], -1).unsqueeze(0).unsqueeze(2)
    def rot(x):
        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
        return torch.cat([-x2, x1], dim=-1)
    xq_f, xk_f = xq.float(), xk.float()
    return ((xq_f*c + rot(xq_f)*s).to(xq.dtype),
            (xk_f*c + rot(xk_f)*s).to(xk.dtype))

class MoDRouter(nn.Module):
    def __init__(self, d_model, capacity=0.5):
        super().__init__()
        self.capacity = capacity
        self.router   = nn.Linear(d_model, 1, bias=False)
    def forward(self, x, layer_fn):
        B, L, D  = x.shape
        k        = max(1, int(L * self.capacity))
        weights  = torch.sigmoid(self.router(x).squeeze(-1))
        topk_w, topk_idx = torch.topk(weights, k, dim=1, sorted=False)
        g_idx    = topk_idx.unsqueeze(-1).expand(B, k, D)
        routed_x = torch.gather(x, 1, g_idx)
        delta    = torch.zeros_like(x)
        delta.scatter_(1, g_idx, layer_fn(routed_x) * topk_w.unsqueeze(-1))
        return delta

class GQA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.nh, self.nkv, self.hd = cfg.n_heads, cfg.n_kv_heads, cfg.head_dim
        self.ng  = cfg.n_heads // cfg.n_kv_heads
        D, Dkv   = cfg.n_heads*cfg.head_dim, cfg.n_kv_heads*cfg.head_dim
        self.q   = nn.Linear(cfg.d_model, D,   bias=False)
        self.k   = nn.Linear(cfg.d_model, Dkv, bias=False)
        self.v   = nn.Linear(cfg.d_model, Dkv, bias=False)
        self.o   = nn.Linear(D, cfg.d_model,   bias=False)
    def forward(self, x, cos, sin):
        B, L, _ = x.shape
        q = self.q(x).view(B, L, self.nh,  self.hd)
        k = self.k(x).view(B, L, self.nkv, self.hd)
        v = self.v(x).view(B, L, self.nkv, self.hd)
        q, k = apply_rope(q, k, cos, sin)
        q = q.transpose(1,2)
        k = k.transpose(1,2).repeat_interleave(self.ng, 1)
        v = v.transpose(1,2).repeat_interleave(self.ng, 1)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)
        return self.o(out.transpose(1,2).contiguous().view(B, L, -1))

class SwiGLU(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        h         = cfg.ffn_hidden
        self.gate = nn.Linear(cfg.d_model, h, bias=False)
        self.up   = nn.Linear(cfg.d_model, h, bias=False)
        self.down = nn.Linear(h, cfg.d_model, bias=False)
    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class NovaBlock(nn.Module):
    def __init__(self, cfg, layer_idx):
        super().__init__()
        self.use_mod   = (layer_idx > 0) and (layer_idx % cfg.mod_every_n == 1)
        self.attn_norm = RMSNorm(cfg.d_model)
        self.ffn_norm  = RMSNorm(cfg.d_model)
        self.attn      = GQA(cfg)
        self.ffn       = SwiGLU(cfg)
        if self.use_mod:
            self.attn_router = MoDRouter(cfg.d_model, cfg.mod_capacity)
            self.ffn_router  = MoDRouter(cfg.d_model, cfg.mod_capacity)
    def forward(self, x, cos, sin):
        if self.use_mod:
            x = x + self.attn_router(x, lambda h: self.attn(self.attn_norm(h), cos, sin))
            x = x + self.ffn_router( x, lambda h: self.ffn( self.ffn_norm(h)))
        else:
            x = x + self.attn(self.attn_norm(x), cos, sin)
            x = x + self.ffn( self.ffn_norm(x))
        return x

class Nova1(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg     = cfg
        self.embed   = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.layers  = nn.ModuleList([NovaBlock(cfg, i) for i in range(cfg.n_layers)])
        self.norm    = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.embed.weight
        cos, sin = precompute_freqs_cis(cfg.head_dim, cfg.max_len)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def forward(self, input_ids):
        B, L = input_ids.shape
        x    = self.embed(input_ids)
        cos, sin = self.rope_cos[:L], self.rope_sin[:L]
        for layer in self.layers:
            x = layer(x, cos, sin)
        return self.lm_head(self.norm(x))

# ── Tokenizer ─────────────────────────────────────────────────────────────────
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

CHATML_END = "<|im_end|>"
REAL_VOCAB = len(tokenizer)          # 50257 — ids above are untrained padding rows
EOS_ID     = tokenizer.eos_token_id  # 50256 <|endoftext|>

# ── Domain Tokens ─────────────────────────────────────────────────────────────
DOMAIN_TOKENS = {
    "General": "<|domain_general|>",
    "Code": "<|domain_code|>",
    "Math": "<|domain_math|>",
    "Reasoning": "<|domain_reasoning|>"
}

def auto_detect_domain(text: str) -> str:
    if not text:
        return "General"
    
    text_lower = text.lower()
    
    # 1. Math rules (from open-web-math rules)
    math_indicators = ["\\frac", "\\int", "\\sum", "sqrt", "≤", "≥", "∑", "∏", "→", "∈", "∀", "∃", "=", "+", "$", "^"]
    math_score = sum(text_lower.count(ind) for ind in math_indicators)
    
    # 2. Code rules (from code_quality_filter keywords)
    code_indicators = ["def ", "class ", "import ", "#include", "public ", "private ", "return ", "fn ", "func ", "const ", "let ", "var ", "{", "}"]
    code_score = sum(text_lower.count(ind) * 2 for ind in code_indicators)
    
    # 3. Reasoning rules (conversational instructions & CoT tags)
    reasoning_indicators = ["explain", "step-by-step", "therefore", "consequently", "because", "how to", "why is", "q:", "a:", "solve"]
    reasoning_score = sum(text_lower.count(ind) for ind in reasoning_indicators)
    
    scores = {
        "Math": math_score,
        "Code": code_score,
        "Reasoning": reasoning_score
    }
    
    best_domain = max(scores, key=scores.get)
    if scores[best_domain] > 1:
        return best_domain
    return "General"

# ── Model loading ─────────────────────────────────────────────────────────────
model = None
device = None
loaded_checkpoint_name = "None"

def find_latest_checkpoint():
    try:
        api = HfApi()
        files = list(api.list_repo_files(
            repo_id=HF_MODEL_REPO,
            repo_type="model",
            token=HF_TOKEN
        ))
        ckpts = sorted(
            [f for f in files if f.endswith(".model.pt")],
            key=lambda x: int(x.split("_s")[-1].split(".")[0]) if "_s" in x else 0
        )
        return ckpts[-1] if ckpts else None
    except Exception as e:
        print(f"Error finding checkpoint: {e}")
        return None

def load_model():
    global model, device, loaded_checkpoint_name

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    cfg = NovaConfig()

    print("🔍 Finding latest checkpoint...")
    latest = find_latest_checkpoint()

    if latest is None:
        raise ValueError("No checkpoints found!")

    loaded_checkpoint_name = latest
    print(f"📥 Loading {latest}...")

    local_path = hf_hub_download(
        repo_id=HF_MODEL_REPO,
        filename=latest,
        repo_type="model",
        token=HF_TOKEN
    )

    model = Nova1(cfg).to(device).to(cfg.dtype)
    sd = torch.load(local_path, map_location=device, weights_only=True)
    sd = {k.replace('_orig_mod.', ''): v for k, v in sd.items()}
    model.load_state_dict(sd, strict=False)
    model.eval()
    print(f"✅ Nova-1 loaded | Device: {device}")

load_model()

# ── Generation ────────────────────────────────────────────────────────────────
def build_chatml_prompt(history, system_prompt, user_message, domain_token):
    """history is list of (user_str, assistant_str) tuples"""
    prompt = f"{domain_token}<|im_start|>system\n{system_prompt}{CHATML_END}\n"
    for user_msg, asst_msg in history:
        prompt += f"<|im_start|>user\n{user_msg}{CHATML_END}\n"
        if asst_msg:
            prompt += f"<|im_start|>assistant\n{asst_msg}{CHATML_END}\n"
    prompt += f"<|im_start|>user\n{user_message}{CHATML_END}\n"
    prompt += f"<|im_start|>assistant\n"
    return prompt

@spaces.GPU
@torch.no_grad()
def generate(
    message,
    history,
    system_prompt,
    max_new_tokens,
    temperature,
    top_p,
    repetition_penalty,
    domain_choice
):
    global model, device

    if model is None:
        load_model()

    if domain_choice == "Auto":
        resolved_domain = auto_detect_domain(message)
    else:
        resolved_domain = domain_choice

    domain_token = DOMAIN_TOKENS[resolved_domain]
    prompt = build_chatml_prompt(history, system_prompt, message, domain_token)

    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    max_ctx = model.cfg.max_len - max_new_tokens
    if input_ids.shape[1] > max_ctx:
        input_ids = input_ids[:, -max_ctx:]

    generated = input_ids.clone()
    past_tokens = []
    output_text = ""

    for _ in range(max_new_tokens):
        with torch.amp.autocast('cuda', dtype=torch.bfloat16):
            logits = model(generated)

        next_logits = logits[0, -1, :].float()
        next_logits[REAL_VOCAB:] = float("-inf")   # mask untrained padding rows

        if temperature > 0:
            next_logits = next_logits / temperature

        if repetition_penalty != 1.0 and past_tokens:
            for token_id in set(past_tokens[-64:]):
                if next_logits[token_id] < 0:
                    next_logits[token_id] *= repetition_penalty
                else:
                    next_logits[token_id] /= repetition_penalty

        probs = F.softmax(next_logits, dim=-1)
        if top_p < 1.0:
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cumsum = torch.cumsum(sorted_probs, dim=0)
            mask = cumsum - sorted_probs > top_p
            sorted_probs[mask] = 0
            sorted_probs /= sorted_probs.sum()
            next_token = sorted_idx[torch.multinomial(sorted_probs, 1)]
        else:
            next_token = torch.argmax(probs).unsqueeze(0)

        past_tokens.append(next_token.item())
        generated = torch.cat([generated, next_token.unsqueeze(0)], dim=1)

        new_text = tokenizer.decode(past_tokens, skip_special_tokens=False)

        if CHATML_END in new_text or "<|im_start|>" in new_text:
            output_text = new_text.split(CHATML_END)[0].split("<|im_start|>")[0].strip()
            break

        output_text = new_text.strip()
        yield output_text

    yield output_text

# ── Raw completion generation (base model mode) ───────────────────────────────
@spaces.GPU
@torch.no_grad()
def generate_raw(
    prompt,
    max_new_tokens,
    temperature,
    top_p,
    repetition_penalty,
    seed,
    stop_raw,
    stop_at_eos,
    domain_choice
):
    global model, device

    if model is None:
        load_model()

    # Domain token resolution
    if domain_choice == "Auto":
        resolved_domain = auto_detect_domain(prompt)
    else:
        resolved_domain = domain_choice

    domain_token = DOMAIN_TOKENS[resolved_domain]

    # Prepend domain token if not present
    if not any(prompt.strip().startswith(tok) for tok in DOMAIN_TOKENS.values()):
        prompt_with_domain = f"{domain_token}{prompt}"
    else:
        prompt_with_domain = prompt

    stop_strings = []
    for line in (stop_raw or "").split("\n"):
        line = line.strip()
        if line:
            stop_strings.append(line.replace("\\n", "\n").replace("\\t", "\t"))

    generator = None
    if seed is not None and int(seed) >= 0:
        generator = torch.Generator(device=device).manual_seed(int(seed))

    ids = tokenizer.encode(prompt_with_domain) if prompt_with_domain and prompt_with_domain.strip() else []
    if not ids:
        ids = [EOS_ID]
    max_ctx = model.cfg.max_len - int(max_new_tokens)
    ids = ids[-max(max_ctx, 1):]
    prompt_tokens = len(ids)

    generated   = torch.tensor([ids], dtype=torch.long, device=device)
    past_tokens = []
    output_text = ""
    t0     = time.time()
    finish = "max tokens"

    def stats(n, done=None):
        dt = max(time.time() - t0, 1e-6)
        s  = f"`{n}` tokens · `{n/dt:.1f}` tok/s · ctx `{prompt_tokens + n}/{model.cfg.max_len}`"
        s += f" · domain: **{resolved_domain.lower()}** (`{domain_token}`)"
        s += f" · stopped: **{done}**" if done else " · generating…"
        return s

    for _ in range(int(max_new_tokens)):
        with torch.amp.autocast('cuda', dtype=torch.bfloat16):
            logits = model(generated)

        next_logits = logits[0, -1, :].float()
        next_logits[REAL_VOCAB:] = float("-inf")   # mask untrained padding rows

        if temperature > 0:
            next_logits = next_logits / temperature

        if repetition_penalty != 1.0 and past_tokens:
            for token_id in set(past_tokens[-64:]):
                if next_logits[token_id] < 0:
                    next_logits[token_id] *= repetition_penalty
                else:
                    next_logits[token_id] /= repetition_penalty

        probs = F.softmax(next_logits, dim=-1)
        if top_p < 1.0:
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cumsum = torch.cumsum(sorted_probs, dim=0)
            mask = cumsum - sorted_probs > top_p
            sorted_probs[mask] = 0
            sorted_probs /= sorted_probs.sum()
            next_token = sorted_idx[torch.multinomial(sorted_probs, 1, generator=generator)]
        else:
            next_token = torch.argmax(probs).unsqueeze(0)

        tid = next_token.item()

        if stop_at_eos and tid == EOS_ID:
            finish = "<|endoftext|>"
            break

        past_tokens.append(tid)
        generated = torch.cat([generated, next_token.unsqueeze(0)], dim=1)
        output_text = tokenizer.decode(past_tokens)

        hit = next((s for s in stop_strings if s in output_text), None)
        if hit is not None:
            output_text = output_text.split(hit)[0]
            finish = f"stop sequence {hit!r}"
            break

        if generated.shape[1] >= model.cfg.max_len:
            finish = "context full"
            break

        yield [(prompt, "prompt"), (output_text, "completion")], output_text, stats(len(past_tokens))

    yield ([(prompt, "prompt"), (output_text, "completion")],
           output_text, stats(len(past_tokens), finish))

# ── Custom CSS ────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');

* { font-family: 'Inter', sans-serif !important; }

body, .gradio-container {
    background: #0a0a0f !important;
    color: #e2e8f0 !important;
}

.nova-header {
    text-align: center;
    padding: 2rem 1rem 1.5rem;
    background: linear-gradient(135deg, #0a0a0f 0%, #0f0f1a 50%, #0a0a0f 100%);
    border-bottom: 1px solid rgba(139, 92, 246, 0.2);
    margin-bottom: 1.5rem;
    border-radius: 16px;
}

.nova-title {
    font-size: 2.8rem;
    font-weight: 700;
    background: linear-gradient(135deg, #a78bfa, #60a5fa, #34d399);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin: 0;
    letter-spacing: -0.5px;
}

.nova-subtitle {
    color: #94a3b8;
    font-size: 0.95rem;
    margin-top: 0.5rem;
}

.nova-badge {
    display: inline-block;
    background: rgba(139, 92, 246, 0.15);
    border: 1px solid rgba(139, 92, 246, 0.4);
    color: #a78bfa;
    padding: 0.25rem 0.75rem;
    border-radius: 999px;
    font-size: 0.75rem;
    font-weight: 600;
    margin: 0.2rem;
    letter-spacing: 0.5px;
}

.warning-badge {
    background: rgba(251, 191, 36, 0.1) !important;
    border: 1px solid rgba(251, 191, 36, 0.3) !important;
    color: #fbbf24 !important;
}

/* Chatbot container */
.chatbot-wrap {
    background: #0d0d17 !important;
    border: 1px solid rgba(139, 92, 246, 0.25) !important;
    border-radius: 16px !important;
    box-shadow: 0 0 40px rgba(139, 92, 246, 0.08), 0 4px 24px rgba(0,0,0,0.4) !important;
    overflow: hidden !important;
}

/* Input box */
.msg-input textarea {
    background: #0d0d17 !important;
    border: 1px solid rgba(139, 92, 246, 0.3) !important;
    border-radius: 12px !important;
    color: #e2e8f0 !important;
    font-size: 0.95rem !important;
    padding: 0.75rem 1rem !important;
    transition: border-color 0.2s ease !important;
}

.msg-input textarea:focus {
    border-color: rgba(139, 92, 246, 0.7) !important;
    box-shadow: 0 0 0 3px rgba(139, 92, 246, 0.12) !important;
    outline: none !important;
}

.msg-input textarea::placeholder { color: #475569 !important; }

/* Send button */
.send-btn button {
    background: linear-gradient(135deg, #7c3aed, #4f46e5) !important;
    border: none !important;
    border-radius: 12px !important;
    color: white !important;
    font-weight: 600 !important;
    transition: all 0.2s ease !important;
    box-shadow: 0 4px 15px rgba(124, 58, 237, 0.3) !important;
}

.send-btn button:hover {
    background: linear-gradient(135deg, #8b5cf6, #6366f1) !important;
    box-shadow: 0 4px 20px rgba(124, 58, 237, 0.5) !important;
    transform: translateY(-1px) !important;
}

/* Secondary buttons */
.secondary-btn button {
    background: rgba(30, 30, 50, 0.8) !important;
    border: 1px solid rgba(139, 92, 246, 0.2) !important;
    border-radius: 10px !important;
    color: #94a3b8 !important;
    transition: all 0.2s ease !important;
}

.secondary-btn button:hover {
    background: rgba(139, 92, 246, 0.15) !important;
    border-color: rgba(139, 92, 246, 0.4) !important;
    color: #a78bfa !important;
}

/* Settings panel */
.settings-panel {
    background: #0d0d17 !important;
    border: 1px solid rgba(139, 92, 246, 0.2) !important;
    border-radius: 16px !important;
    padding: 1.25rem !important;
    box-shadow: 0 4px 24px rgba(0,0,0,0.3) !important;
}

.settings-panel label {
    color: #94a3b8 !important;
    font-size: 0.82rem !important;
    font-weight: 500 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.5px !important;
}

.settings-panel textarea {
    background: #0a0a0f !important;
    border: 1px solid rgba(139, 92, 246, 0.2) !important;
    border-radius: 10px !important;
    color: #e2e8f0 !important;
    font-size: 0.875rem !important;
}

.settings-panel textarea:focus {
    border-color: rgba(139, 92, 246, 0.5) !important;
    box-shadow: 0 0 0 2px rgba(139, 92, 246, 0.1) !important;
}

/* Model info card */
.model-info-card {
    background: linear-gradient(135deg, rgba(139, 92, 246, 0.08), rgba(96, 165, 250, 0.05));
    border: 1px solid rgba(139, 92, 246, 0.2);
    border-radius: 12px;
    padding: 1rem;
    margin-top: 1rem;
}

.model-info-card p {
    color: #94a3b8 !important;
    font-size: 0.82rem !important;
    margin: 0.2rem 0 !important;
    line-height: 1.6 !important;
}

.model-info-card strong { color: #a78bfa !important; }

/* Examples */
table.examples td {
    background: rgba(139, 92, 246, 0.06) !important;
    border: 1px solid rgba(139, 92, 246, 0.15) !important;
    border-radius: 8px !important;
    color: #94a3b8 !important;
    font-size: 0.85rem !important;
    padding: 0.5rem 0.75rem !important;
    cursor: pointer !important;
    transition: all 0.2s !important;
}

table.examples td:hover {
    background: rgba(139, 92, 246, 0.15) !important;
    color: #c4b5fd !important;
    border-color: rgba(139, 92, 246, 0.4) !important;
}

/* Scrollbar */
::-webkit-scrollbar { width: 6px; }
::-webkit-scrollbar-track { background: #0a0a0f; }
::-webkit-scrollbar-thumb {
    background: rgba(139, 92, 246, 0.4);
    border-radius: 3px;
}
::-webkit-scrollbar-thumb:hover { background: rgba(139, 92, 246, 0.7); }

input[type="range"] { accent-color: #7c3aed !important; }

.footer-text {
    text-align: center;
    color: #334155;
    font-size: 0.75rem;
    padding: 1rem;
    margin-top: 1rem;
    border-top: 1px solid rgba(139, 92, 246, 0.1);
}

/* ── Tabs ── */
.tab-nav {
    border-bottom: 1px solid rgba(139, 92, 246, 0.2) !important;
    margin-bottom: 1rem !important;
}
.tab-nav button {
    color: #64748b !important;
    font-weight: 600 !important;
    font-size: 0.9rem !important;
    border-bottom: 2px solid transparent !important;
}
.tab-nav button.selected {
    color: #a78bfa !important;
    border-bottom-color: #7c3aed !important;
}

/* ── Raw completion playground ── */
#raw-prompt textarea, #raw-completion textarea {
    font-family: 'JetBrains Mono', ui-monospace, monospace !important;
    font-size: 0.9rem !important;
    line-height: 1.65 !important;
    background: #0d0d17 !important;
    border: 1px solid rgba(139, 92, 246, 0.3) !important;
    border-radius: 12px !important;
    color: #e2e8f0 !important;
    padding: 0.75rem 1rem !important;
}
#raw-prompt textarea:focus, #raw-completion textarea:focus {
    border-color: rgba(139, 92, 246, 0.7) !important;
    box-shadow: 0 0 0 3px rgba(139, 92, 246, 0.12) !important;
}
#raw-output {
    font-family: 'JetBrains Mono', ui-monospace, monospace !important;
    font-size: 0.9rem !important;
    line-height: 1.65 !important;
    background: #0d0d17 !important;
    border: 1px solid rgba(139, 92, 246, 0.25) !important;
    border-radius: 12px !important;
    padding: 0.75rem 1rem !important;
}
.raw-stats {
    text-align: right;
    color: #475569 !important;
    font-size: 0.78rem !important;
    margin: 0.25rem 0 0.5rem;
}
.raw-stats code {
    color: #a78bfa !important;
    background: rgba(124, 58, 237, 0.1) !important;
    padding: 0.05rem 0.35rem !important;
    border-radius: 4px !important;
}
"""

# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(title="Nova-1 | Bc-AI") as demo:

    gr.HTML(f"""
    <div class="nova-header">
        <h1 class="nova-title">✦ Nova-1</h1>
        <p class="nova-subtitle">Custom Language Model by Bc-AI — raw pretrained base model</p>
        <div style="margin-top: 0.75rem;">
            <span class="nova-badge">1.217B Params</span>
            <span class="nova-badge">2048 Context</span>
            <span class="nova-badge">Phase 2</span>
            <span class="nova-badge">~8B Tokens</span>
            <span class="nova-badge">GQA + MoD + SwiGLU</span>
            <span class="nova-badge warning-badge">⚠️ Raw Pretrain — Not Instruction Tuned</span>
        </div>
        <p style="color: #475569; font-size: 0.78rem; margin-top: 0.75rem;">
            Checkpoint: <code style="color: #7c3aed; background: rgba(124,58,237,0.1);
            padding: 0.1rem 0.4rem; border-radius: 4px;">{loaded_checkpoint_name}</code>
        </p>
    </div>
    """)

    with gr.Tabs():

        # ══ Tab 1: Raw completion (primary — it's a base model) ══════════════
        with gr.Tab("✍️ Completion"):

            with gr.Row(equal_height=True):

                with gr.Column(scale=3):
                    raw_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Start any text — Nova-1 will continue it…",
                        lines=8,
                        max_lines=20,
                        elem_id="raw-prompt",
                        value="The history of the transistor begins in",
                    )

                    with gr.Row():
                        raw_gen  = gr.Button("✦ Generate", variant="primary",
                                             elem_classes=["send-btn"], scale=2)
                        raw_cont = gr.Button("↳ Continue",
                                             elem_classes=["secondary-btn"], scale=1)
                        raw_stop = gr.Button("■ Stop", variant="stop", scale=1)
                        raw_clr  = gr.Button("🗑 Clear",
                                             elem_classes=["secondary-btn"], scale=1)

                    raw_stats = gr.Markdown("", elem_classes=["raw-stats"])

                    raw_output = gr.HighlightedText(
                        label="Output  (grey = prompt · purple = completion)",
                        elem_id="raw-output",
                        show_legend=False,
                        combine_adjacent=True,
                        color_map={"prompt": "#262a3a", "completion": "#4c1d95"},
                    )

                    with gr.Accordion("Completion only (copyable)", open=False):
                        raw_completion = gr.Textbox(
                            show_label=False,
                            lines=8,
                            elem_id="raw-completion",
                        )

                    gr.Examples(
                        examples=[
                            ["The history of the transistor begins in"],
                            ["def quicksort(arr):\n    "],
                            ["Q: Why is the sky blue?\nA:"],
                            ["Once upon a time, at the edge of a dying star, there lived"],
                            ["# Chapter 1\n\nThe old lighthouse stood"],
                            ["Paris is the capital of France. Rome is the capital of Italy. Tokyo is the capital of"],
                        ],
                        inputs=raw_prompt,
                        label="Example prompts",
                    )

                with gr.Column(scale=1, min_width=260):
                    with gr.Group(elem_classes=["settings-panel"]):
                        gr.HTML("<p style='color:#a78bfa; font-weight:600; font-size:0.9rem;"
                                "margin-bottom:1rem;'>⚙️ Sampling</p>")

                        raw_domain = gr.Dropdown(
                            choices=["Auto", "General", "Code", "Math", "Reasoning"],
                            value="Auto",
                            label="Domain Token",
                            info="Appends domain prefix dynamically"
                        )
                        raw_max  = gr.Slider(16, 512, value=200, step=16,
                                             label="Max New Tokens")
                        raw_temp = gr.Slider(0.1, 2.0, value=0.8, step=0.05,
                                             label="Temperature")
                        raw_topp = gr.Slider(0.1, 1.0, value=0.95, step=0.05,
                                             label="Top-p")
                        raw_rep  = gr.Slider(1.0, 1.5, value=1.1, step=0.05,
                                             label="Repetition Penalty")

                    with gr.Accordion("Advanced", open=False):
                        raw_seed = gr.Number(value=-1,
                                             label="Seed (-1 = random)")
                        raw_eos  = gr.Checkbox(value=True,
                                               label="Stop at <|endoftext|>")
                        raw_stops = gr.Textbox(
                            label="Stop sequences (one per line)",
                            placeholder="### or \\n\\n",
                            lines=2,
                        )

                    gr.HTML("""
                    <div class="model-info-card">
                        <p><strong>Base model tips</strong></p>
                        <p>· It <em>continues</em> text — it doesn't follow instructions.</p>
                        <p>· Few-shot works best: show 2–3 examples, then cut off mid-pattern.</p>
                        <p>· <strong>Continue</strong> appends the output to the prompt and keeps going.</p>
                        <p>· Lower temp + higher rep penalty → more factual-looking prose.</p>
                    </div>
                    """)

        # ══ Tab 2: Chat (ChatML, experimental) ═══════════════════════════════
        with gr.Tab("💬 Chat (ChatML)"):

            gr.HTML("""
            <p style="color:#fbbf24; font-size:0.8rem; margin:0 0 0.75rem;
                      background:rgba(251,191,36,0.07); border:1px solid rgba(251,191,36,0.25);
                      border-radius:10px; padding:0.6rem 0.9rem;">
                ⚠️ Nova-1 is <strong>not instruction tuned</strong>. This tab wraps your message
                in a ChatML template — expect incoherent or off-topic replies.
                The Completion tab shows what the model actually does well.
            </p>
            """)

            with gr.Row(equal_height=True):

                # ── Left: Chat ────────────────────────────────────────────────
                with gr.Column(scale=3):
                    chatbot = gr.Chatbot(
                        label="",
                        height=520,
                        elem_classes=["chatbot-wrap"],
                        avatar_images=("👤", "✦"),
                        show_label=False,
                    )

                    with gr.Row():
                        msg = gr.Textbox(
                            placeholder="Message Nova-1...",
                            show_label=False,
                            scale=5,
                            container=False,
                            elem_classes=["msg-input"],
                            lines=1,
                            max_lines=5,
                        )
                        send_btn = gr.Button(
                            "Send ↗",
                            variant="primary",
                            scale=1,
                            elem_classes=["send-btn"],
                            min_width=80
                        )

                    with gr.Row():
                        clear_btn = gr.Button("🗑 Clear", elem_classes=["secondary-btn"], size="sm")
                        retry_btn = gr.Button("↺ Retry",  elem_classes=["secondary-btn"], size="sm")

                    gr.Examples(
                        examples=[
                            ["Explain what a neural network is in simple terms"],
                            ["Write a haiku about the cosmos"],
                            ["What makes a good programming language?"],
                            ["Tell me something interesting about mathematics"],
                            ["Write the opening of a sci-fi short story"],
                        ],
                        inputs=msg,
                        label="Try an example",
                    )

                # ── Right: Settings ───────────────────────────────────────────
                with gr.Column(scale=1, min_width=260):
                    with gr.Group(elem_classes=["settings-panel"]):

                        gr.HTML("<p style='color:#a78bfa; font-weight:600; font-size:0.9rem;"
                                "margin-bottom:1rem;'>⚙️ Generation Settings</p>")

                        chat_domain = gr.Dropdown(
                            choices=["Auto", "General", "Code", "Math", "Reasoning"],
                            value="Auto",
                            label="Domain Token",
                            info="Dynamic training format tag"
                        )
                        system_prompt = gr.Textbox(
                            value="You are Nova-1, a helpful and knowledgeable AI assistant. Be concise and clear.",
                            label="System Prompt",
                            lines=3,
                        )
                        max_new_tokens = gr.Slider(
                            minimum=32, maximum=512, value=200, step=32,
                            label="Max New Tokens"
                        )
                        temperature = gr.Slider(
                            minimum=0.1, maximum=2.0, value=0.8, step=0.05,
                            label="Temperature"
                        )
                        top_p = gr.Slider(
                            minimum=0.1, maximum=1.0, value=0.95, step=0.05,
                            label="Top-p"
                        )
                        repetition_penalty = gr.Slider(
                            minimum=1.0, maximum=1.5, value=1.1, step=0.05,
                            label="Repetition Penalty"
                        )

                        gr.HTML("""
                        <div class="model-info-card">
                            <p><strong>Architecture</strong></p>
                            <p>· Grouped Query Attention (GQA)</p>
                            <p>· Mixture of Depths (MoD)</p>
                            <p>· SwiGLU FFN</p>
                            <p>· RoPE Embeddings</p>
                            <p>· RMSNorm</p>
                            <br>
                            <p><strong>Training</strong></p>
                            <p>· Trained from scratch</p>
                            <p>· Phase 2 (2048 ctx from day 1)</p>
                            <p>· ~8B tokens seen</p>
                            <p>· 8-bit AdamW optimizer</p>
                            <br>
                            <p><strong>Hardware</strong></p>
                            <p>· RTX Pro 6000 Blackwell</p>
                            <p>· 95GB VRAM</p>
                            <p>· ZeroGPU inference</p>
                        </div>
                        """)

    gr.HTML("""
    <div class="footer-text">
        Nova-1 by Bc-AI · Built with ❤️ · Powered by HuggingFace ZeroGPU
    </div>
    """)

    # ── Completion tab wiring ─────────────────────────────────────────────────
    raw_inputs  = [raw_prompt, raw_max, raw_temp, raw_topp, raw_rep,
                   raw_seed, raw_stops, raw_eos, raw_domain]
    raw_outputs = [raw_output, raw_completion, raw_stats]

    ev_gen  = raw_gen.click(generate_raw, raw_inputs, raw_outputs)
    ev_sub  = raw_prompt.submit(generate_raw, raw_inputs, raw_outputs)

    def merge_prompt_completion(prompt, completion):
        return (prompt or "") + (completion or "")

    ev_cont = raw_cont.click(
        merge_prompt_completion, [raw_prompt, raw_completion], raw_prompt, queue=False
    ).then(
        generate_raw, raw_inputs, raw_outputs
    )

    raw_stop.click(None, None, None, cancels=[ev_gen, ev_sub, ev_cont], queue=False)

    raw_clr.click(
        lambda: ("", None, "", ""), None,
        [raw_prompt, raw_output, raw_completion, raw_stats], queue=False
    )

    # ── Chat tab wiring (Gradio 6 dict format) ────────────────────────────────
    def user_submit(message, history):
        if not message.strip():
            return "", history
        return "", history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": ""}
        ]

    def bot_respond(history, system_prompt, max_new_tokens, temperature, top_p, rep_penalty, domain_choice):
        if not history or len(history) < 2:
            return history

        user_message = history[-2]["content"]
        
        tuple_history = []
        for i in range(0, len(history) - 2, 2):
            if i + 1 < len(history):
                tuple_history.append((
                    history[i]["content"],
                    history[i + 1]["content"]
                ))

        for response in generate(
            user_message,
            tuple_history,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
            rep_penalty,
            domain_choice
        ):
            history[-1]["content"] = response
            yield history

    def clear_chat():
        return [], ""

    def retry_last(history, system_prompt, max_new_tokens, temperature, top_p, rep_penalty, domain_choice):
        if not history or len(history) < 2:
            return history
        
        user_message = history[-2]["content"]
        history[-1]["content"] = ""
        
        tuple_history = []
        for i in range(0, len(history) - 2, 2):
            if i + 1 < len(history):
                tuple_history.append((
                    history[i]["content"],
                    history[i + 1]["content"]
                ))

        for response in generate(
            user_message,
            tuple_history,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
            rep_penalty,
            domain_choice
        ):
            history[-1]["content"] = response
            yield history

    # Wire up events
    msg.submit(
        user_submit, [msg, chatbot], [msg, chatbot], queue=False
    ).then(
        bot_respond,
        [chatbot, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty, chat_domain],
        chatbot
    )

    send_btn.click(
        user_submit, [msg, chatbot], [msg, chatbot], queue=False
    ).then(
        bot_respond,
        [chatbot, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty, chat_domain],
        chatbot
    )

    clear_btn.click(clear_chat, outputs=[chatbot, msg])

    retry_btn.click(
        retry_last,
        [chatbot, system_prompt, max_new_tokens, temperature, top_p, repetition_penalty, chat_domain],
        chatbot
    )

if __name__ == "__main__":
    demo.queue(max_size=8).launch(
        css=CSS,
        show_error=True,
        share=False
    )