from __future__ import annotations

import os
os.environ["KERAS_BACKEND"] = "jax"

import numpy as np
import jax
import keras
import gradio as gr
import time
from pathlib import Path

from veylon_model import create_llm
from tokenizer import TokenizerWrapper
from config import (
    CONTEXT,
    vocab_size,
    D_MODEL,
    numberoflayers,
    numberofheads,
    d_Latent,
    ffn_mult,
    num_kv_heads,
    swa_window,
)

# ============================================================
# Initialize (runs once)
# ============================================================

keras.mixed_precision.set_global_policy("mixed_bfloat16")

print(f"Backend: {keras.backend.backend()}")
print(f"JAX devices: {jax.devices()}")

# Load tokenizer
tokenizer = TokenizerWrapper("tokenizer.model")
assert tokenizer.vocab_size == vocab_size, (
    f"Tokenizer vocab ({tokenizer.vocab_size}) != config vocab ({vocab_size})"
)
print(f"✓ Tokenizer loaded: {tokenizer.vocab_size} vocab")

# Build model
print("Building model...")
model = create_llm(
    vocab_size=vocab_size,
    d_model=D_MODEL,
    n_layers=numberoflayers,
    n_heads=numberofheads,
    d_latent=d_Latent,
    ffn_mult=ffn_mult,
    max_seq_len=CONTEXT,
    use_moe=False,
    num_kv_heads=num_kv_heads,
    swa_window=swa_window,
)

# Warmup
dummy = np.zeros((1, CONTEXT), dtype=np.int32)
_ = model(dummy, training=False)
print("✓ Model built successfully")

# Load weights
WEIGHTS_PATH = "veylon_final.weights.h5"
if Path(WEIGHTS_PATH).exists():
    print(f"Loading weights from: {WEIGHTS_PATH}")
    model.load_weights(WEIGHTS_PATH)
    print("✓ Weights loaded successfully")
else:
    print(f"WARNING: {WEIGHTS_PATH} not found. Using untrained model.")

print(f"✓ Model params: {model.count_params():,}\n")

# ============================================================
# Sampling
# ============================================================

def sample_from_logits(
    logits_row: np.ndarray,
    temperature: float = 0.7,
    min_p: float = 0.05,
) -> int:
    """
    Min-P sampling (replaces top-k). Top-k/top-p keep a FIXED-size or
    fixed-cumulative-mass candidate pool open regardless of how confident
    the model actually is -- on a small/uncertain model, that means when
    it doesn't know what comes next, top-k still hands it 50 candidates
    to choose from, many of them garbage. Min-P instead sets a threshold
    RELATIVE to the top token's probability: confident predictions
    collapse the pool to 1-2 tokens, uncertain predictions keep it wide.
    Expects and returns a 1D array (single vocab-length row), not (1, vocab).
    """
    logits_row = np.array(logits_row, dtype=np.float32, copy=True)

    if temperature > 0:
        logits_row = logits_row / float(max(temperature, 1e-8))

    row = logits_row - np.max(logits_row)  # numerical stability
    probs = np.exp(row)
    probs = probs / probs.sum()

    if min_p > 0:
        threshold = min_p * probs.max()
        mask = probs >= threshold
        if not mask.any():  # degenerate guard -- never zero out everything
            mask[np.argmax(probs)] = True
        probs = np.where(mask, probs, 0.0)
        probs = probs / probs.sum()

    return int(np.random.choice(len(probs), p=probs))


def apply_repetition_controls(
    logits_row: np.ndarray,
    generated_ids: list,
    no_repeat_ngram_size: int = 3,
    repetition_penalty: float = 1.2,
) -> np.ndarray:
    """
    Applied BEFORE sample_from_logits, on the model's own newly-generated
    tokens only (never the user's prompt -- penalizing someone for the
    model reusing words from THEIR prompt would be wrong, this is about
    stopping the model's own loops).

    1. repetition_penalty (soft): divides a positive logit for an
       already-seen token, MULTIPLIES a negative one -- always pushes the
       token down, regardless of logit sign (a naive "divide by penalty"
       for a negative logit would actually push it UP, backwards).
    2. no_repeat_ngram_size (hard): if continuing with a candidate token
       would recreate an n-gram that's already appeared in this
       generation, that candidate is set to -inf. This is what actually
       kills infinite loops -- it's a hard constraint, not a nudge.
    """
    logits_row = logits_row.copy()

    if repetition_penalty and repetition_penalty != 1.0 and generated_ids:
        for tok_id in set(generated_ids):
            if logits_row[tok_id] > 0:
                logits_row[tok_id] /= repetition_penalty
            else:
                logits_row[tok_id] *= repetition_penalty

    n = no_repeat_ngram_size
    if n and n > 0 and len(generated_ids) >= n - 1:
        prefix = tuple(generated_ids[-(n - 1):]) if n > 1 else tuple()
        banned = set()
        for i in range(len(generated_ids) - n + 1):
            if tuple(generated_ids[i:i + n - 1]) == prefix:
                banned.add(generated_ids[i + n - 1])
        for tok_id in banned:
            logits_row[tok_id] = -np.inf

    return logits_row


# Few-shot "anchoring": Veylon is a CONTINUATION-only base model, never
# instruction-tuned -- it has no concept of "Input:"/"Output:" turns, so
# priming it with instruction-style examples (as generic small-LM advice
# often suggests) would push it OUT of its training distribution, not
# into it. What anchoring means correctly here is different: prepend a
# short, clean, well-formed piece of narrative prose so the model's
# recent-attention window is dominated by coherent English *before*
# generation starts, nudging it away from degenerating into word salad.
# This consumes some of the context budget, so it's short and toggleable.
ANCHOR_TEXT = (
    "The old lighthouse stood at the edge of the cliff, its light "
    "sweeping slowly across the dark water. Every night, the keeper "
    "climbed the narrow stairs to check the lamp, and every night the "
    "sea answered with the steady sound of waves against the rocks.\n\n"
)

# ============================================================
# Streaming generation with live telemetry
# ============================================================

def respond(
    message: str,
    history: list,
    max_new_tokens: int = 96,
    temperature: float = 0.7,
    min_p: float = 0.05,
    no_repeat_ngram_size: int = 3,
    repetition_penalty: float = 1.2,
    use_anchor: bool = True,
):
    """
    gr.ChatInterface-compatible streaming generator. Yields incrementally
    as tokens are produced, and appends a real (measured, not decorative)
    telemetry line -- token count / time-to-first-token / tokens-per-
    second -- under the reply, matching the "signature element" of this
    UI: honest, live numbers about what the model is actually doing.

    IMPORTANT, honest by design: `history` is accepted (gr.ChatInterface
    requires the signature) and IS shown to the person as an ongoing
    conversation, but is deliberately NOT concatenated into the prompt
    fed to the model. Arya is a continuation-only base model, never
    instruction-tuned on multi-turn User/Assistant-style dialogue at
    scale -- feeding it a growing instruction-formatted transcript would
    push it further out of its training distribution with every turn,
    not closer to coherence. Each message is generated as its own fresh
    continuation (with the anchor primer). This is stated plainly in the
    UI caption rather than left as a surprise.
    """
    try:
        full_prompt = (ANCHOR_TEXT + message) if use_anchor else message

        tokens = tokenizer.encode(full_prompt, add_bos=True, add_eos=False)
        if len(tokens) == 0:
            tokens = [tokenizer.bos_id if hasattr(tokenizer, "bos_id") else 1]
        tokens = tokens[-CONTEXT:]

        generated_ids: list = []
        t_start = time.perf_counter()
        ttft = None

        prompt_ids = np.array([tokens], dtype=np.int32)
        logits, cache_k, cache_v = model.generate_step(
            prompt_ids, cache_k=None, cache_v=None, cache_pos=0,
        )

        logits_row = np.array(logits[0, -1, :], dtype=np.float32, copy=True)
        logits_row = apply_repetition_controls(
            logits_row, generated_ids, no_repeat_ngram_size, repetition_penalty
        )
        next_token = sample_from_logits(logits_row, temperature=temperature, min_p=min_p)
        tokens.append(next_token)
        generated_ids.append(next_token)
        ttft = time.perf_counter() - t_start

        partial_text = tokenizer.decode(generated_ids)
        yield _with_telemetry(partial_text, len(generated_ids), ttft, t_start)

        if next_token != tokenizer.eos_id and len(tokens) < CONTEXT:
            cache_pos = len(prompt_ids[0])

            for _ in range(max_new_tokens - 1):
                next_input = np.array([[next_token]], dtype=np.int32)
                logits, cache_k, cache_v = model.generate_step(
                    next_input, cache_k=cache_k, cache_v=cache_v, cache_pos=cache_pos,
                )
                cache_pos += 1

                logits_row = np.array(logits[0, -1, :], dtype=np.float32, copy=True)
                logits_row = apply_repetition_controls(
                    logits_row, generated_ids, no_repeat_ngram_size, repetition_penalty
                )
                next_token = sample_from_logits(logits_row, temperature=temperature, min_p=min_p)
                tokens.append(next_token)
                generated_ids.append(next_token)

                if next_token == tokenizer.eos_id:
                    break

                partial_text = tokenizer.decode(generated_ids)
                yield _with_telemetry(partial_text, len(generated_ids), ttft, t_start)

                if len(tokens) >= CONTEXT:
                    break

    except Exception as e:
        yield f"Error: {str(e)}"


def _with_telemetry(text: str, n_tok: int, ttft: float, t_start: float) -> str:
    elapsed = max(time.perf_counter() - t_start, 1e-6)
    tok_per_s = n_tok / elapsed
    telemetry = (
        f"\n\n<sub style='font-family:monospace;color:var(--body-text-color-subdued)'>"
        f"{n_tok} tok · TTFT {ttft * 1000:.0f} ms · {tok_per_s:.1f} tok/s</sub>"
    )
    return text + telemetry


# ============================================================
# Gradio UI — Arya
# ============================================================
# NOTE on "gr.server()": there is no such function in Gradio. gr.Blocks
# and gr.ChatInterface are ALREADY built on FastAPI/Uvicorn internally --
# demo.launch() below is what actually starts that server. Same
# clarification as the earlier "use FastAPI for speed" question: the web
# framework was never the bottleneck for a custom JAX model like this.

ARYA_CSS = """
:root {
    --arya-bg: #100C1A;
    --arya-bg-2: #17111F;
    --arya-panel: #1D1730;
    --arya-panel-2: #271F42;
    --arya-accent: #E8A33D;
    --arya-accent-soft: #F0BE73;
    --arya-accent-2: #5FBFA8;
    --arya-indigo: #5B67D9;
    --arya-text: #F3EEE3;
    --arya-text-dim: #948BAA;
    --arya-hairline: #2A2440;

    /* Retheme Gradio's OWN documented CSS variables (verified against the
       installed theme's actual variable names) instead of guessing at
       hashed/scoped internal component classes -- this cascades correctly
       into sliders, checkboxes, accordions, buttons, etc. automatically,
       and won't silently stop working on a Gradio version bump the way
       hand-guessed internal class names would. */
    --body-background-fill: var(--arya-bg);
    --background-fill-primary: var(--arya-bg);
    --background-fill-secondary: var(--arya-panel);
    --block-background-fill: var(--arya-panel);
    --block-border-color: var(--arya-hairline);
    --block-title-text-color: var(--arya-text-dim);
    --block-label-text-color: var(--arya-text-dim);
    --body-text-color: var(--arya-text);
    --body-text-color-subdued: var(--arya-text-dim);
    --border-color-primary: var(--arya-hairline);
    --border-color-accent: var(--arya-accent);
    --color-accent: var(--arya-accent);
    --color-accent-soft: color-mix(in srgb, var(--arya-accent) 20%, transparent);
    --input-background-fill: var(--arya-panel-2);
    --input-border-color: var(--arya-hairline);
    --input-border-color-focus: var(--arya-accent);
    --input-placeholder-color: var(--arya-text-dim);
    --slider-color: var(--arya-accent);
    --checkbox-background-color-selected: var(--arya-accent);
    --checkbox-border-color-selected: var(--arya-accent);
    --checkbox-label-background-fill-selected: color-mix(in srgb, var(--arya-accent) 18%, var(--arya-panel));
    --button-primary-background-fill: linear-gradient(135deg, var(--arya-accent), var(--arya-accent-soft));
    --button-primary-background-fill-hover: var(--arya-accent-soft);
    --button-primary-text-color: #201306;
    --button-secondary-background-fill: var(--arya-panel-2);
    --button-secondary-border-color: var(--arya-hairline);
    --button-secondary-background-fill-hover: var(--arya-panel);
    --link-text-color: var(--arya-accent-2);
}

* { scrollbar-width: thin; scrollbar-color: var(--arya-panel-2) transparent; }
::-webkit-scrollbar { width: 8px; height: 8px; }
::-webkit-scrollbar-thumb { background: var(--arya-panel-2); border-radius: 8px; }
::-webkit-scrollbar-thumb:hover { background: var(--arya-hairline); }

.gradio-container {
    background:
        radial-gradient(ellipse 900px 400px at 15% -10%, color-mix(in srgb, var(--arya-accent) 10%, transparent), transparent 60%),
        radial-gradient(ellipse 700px 400px at 100% 0%, color-mix(in srgb, var(--arya-indigo) 12%, transparent), transparent 55%),
        var(--arya-bg) !important;
    font-family: 'Inter', system-ui, sans-serif !important;
    max-width: 880px !important;
}

/* ── Topbar ── */
#arya-topbar {
    display: flex;
    align-items: center;
    gap: 12px;
    padding: 18px 4px 16px 4px;
    border-bottom: 1px solid var(--arya-hairline);
    margin-bottom: 10px;
}
#arya-mark {
    width: 34px; height: 34px; flex-shrink: 0;
    display: flex; align-items: center; justify-content: center;
    border-radius: 10px;
    background: linear-gradient(135deg, var(--arya-accent), var(--arya-indigo));
    box-shadow: 0 2px 14px color-mix(in srgb, var(--arya-accent) 35%, transparent);
}
#arya-mark svg { width: 19px; height: 19px; }
#arya-title-block { display: flex; flex-direction: column; line-height: 1.15; }
#arya-title {
    font-family: 'IBM Plex Mono', monospace;
    font-weight: 700;
    font-size: 1.18rem;
    letter-spacing: 0.14em;
    color: var(--arya-text);
}
#arya-subtitle {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.66rem;
    letter-spacing: 0.06em;
    color: var(--arya-text-dim);
}
#arya-status-chip {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.72rem;
    font-weight: 500;
    color: var(--arya-accent-2);
    display: flex; align-items: center; gap: 6px;
    padding: 5px 10px;
    border-radius: 999px;
    background: color-mix(in srgb, var(--arya-accent-2) 12%, transparent);
    border: 1px solid color-mix(in srgb, var(--arya-accent-2) 30%, transparent);
    margin-left: 14px;
}
#arya-status-chip::before {
    content: "";
    width: 6px; height: 6px; border-radius: 50%;
    background: var(--arya-accent-2);
    box-shadow: 0 0 8px var(--arya-accent-2);
    animation: arya-pulse 2s ease-in-out infinite;
}
@keyframes arya-pulse {
    0%, 100% { opacity: 1; transform: scale(1); }
    50% { opacity: 0.5; transform: scale(0.8); }
}
#arya-specs {
    margin-left: auto;
    display: flex; gap: 6px; flex-wrap: wrap; justify-content: flex-end;
}
.arya-chip {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.66rem;
    letter-spacing: 0.03em;
    color: var(--arya-text-dim);
    padding: 4px 9px;
    border-radius: 6px;
    background: var(--arya-panel);
    border: 1px solid var(--arya-hairline);
    white-space: nowrap;
}

/* ── Chat transcript (verified real Gradio 6 Chatbot classes, not guessed:
     .user-row / .bot-row / .message-bubble-border confirmed present in the
     installed package's compiled component CSS) ── */
.user-row {
    background: linear-gradient(135deg, var(--arya-accent), var(--arya-accent-soft)) !important;
    color: #201306 !important;
}
.user-row * { color: #201306 !important; }
.bot-row {
    background: var(--arya-panel) !important;
    border: 1px solid var(--arya-hairline) !important;
}
.message-bubble-border {
    border-color: var(--arya-hairline) !important;
    box-shadow: 0 4px 18px rgba(0,0,0,0.25);
    animation: arya-rise 0.25s ease-out;
}
@keyframes arya-rise {
    from { opacity: 0; transform: translateY(6px); }
    to { opacity: 1; transform: translateY(0); }
}
.bot-row code, .bot-row pre {
    font-family: 'IBM Plex Mono', monospace !important;
    background: var(--arya-bg-2) !important;
    border: 1px solid var(--arya-hairline) !important;
    border-radius: 6px !important;
}

/* ── Input row ── */
textarea, input[type="text"] {
    transition: border-color 0.15s ease, box-shadow 0.15s ease !important;
}
textarea:focus, input[type="text"]:focus {
    box-shadow: 0 0 0 3px color-mix(in srgb, var(--arya-accent) 22%, transparent) !important;
}
button.primary {
    transition: transform 0.12s ease, box-shadow 0.12s ease !important;
    box-shadow: 0 2px 12px color-mix(in srgb, var(--arya-accent) 30%, transparent) !important;
}
button.primary:hover { transform: translateY(-1px) scale(1.03); }

/* ── Settings accordion & example chips: themed via the CSS variables
   above, no internal class names needed ── */
.example {
    transition: transform 0.12s ease, border-color 0.12s ease !important;
}
.example:hover {
    transform: translateY(-2px);
    border-color: var(--arya-accent) !important;
}

#arya-caption {
    font-family: 'IBM Plex Mono', monospace;
    font-size: 0.68rem;
    color: var(--arya-text-dim);
    text-align: center;
    padding: 16px 20px 6px 20px;
    letter-spacing: 0.03em;
    border-top: 1px solid var(--arya-hairline);
    margin-top: 12px;
}
"""

ARYA_HEAD = """
<link rel="preconnect" href="https://fonts.googleapis.com">
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=IBM+Plex+Mono:wght@400;500;700&display=swap" rel="stylesheet">
"""

# Simple geometric monogram, not a decorative flourish -- the gradient fill
# (accent -> indigo) matches #arya-mark's CSS background, single mark
# rather than a busier logo, keeps it legible at 19px.
ARYA_MARK_SVG = """
<svg viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
  <path d="M12 3L21 20H16.5L12 11L7.5 20H3L12 3Z" fill="#100C1A"/>
</svg>
"""


def main():
    spec_chips = [
        f"{model.count_params() / 1e6:.1f}M PARAMS",
        f"{numberoflayers}L",
        f"CTX {CONTEXT}",
        f"GQA {numberofheads}/{num_kv_heads}",
    ]
    spec_html = "".join(f'<span class="arya-chip">{c}</span>' for c in spec_chips)

    with gr.Blocks(title="Arya") as demo:

        gr.HTML(f"""
        <div id="arya-topbar">
            <div id="arya-mark">{ARYA_MARK_SVG}</div>
            <div id="arya-title-block">
                <span id="arya-title">ARYA</span>
                <span id="arya-subtitle">FROM-SCRATCH SMALL LM</span>
            </div>
            <span id="arya-status-chip">READY</span>
            <div id="arya-specs">{spec_html}</div>
        </div>
        """)

        gr.ChatInterface(
            fn=respond,
            additional_inputs=[
                gr.Slider(10, 256, value=96, step=10, label="Max tokens"),
                gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
                gr.Slider(0.0, 0.3, value=0.05, step=0.01, label="Min-P",
                          info="Replaces top-k. Higher = stricter."),
                gr.Slider(0, 6, value=3, step=1, label="No-repeat n-gram size",
                          info="Hard-blocks repeating any phrase this long. 0 disables."),
                gr.Slider(1.0, 1.5, value=1.2, step=0.05, label="Repetition penalty"),
                gr.Checkbox(value=True, label="Narrative anchor",
                            info="Primes the model with clean prose before your prompt."),
            ],
            additional_inputs_accordion=gr.Accordion("Settings", open=False),
            examples=[
                ["Once upon a time, in a village by the mountains,"],
                ["The scientist looked at the data and realized"],
                ["Explain photosynthesis simply."],
            ],
            chatbot=gr.Chatbot(height=460, show_label=False),
            textbox=gr.Textbox(placeholder="Ask Arya anything...", show_label=False),
        )

        gr.HTML("""
        <div id="arya-caption">
            CONTINUATION MODEL, NOT INSTRUCTION-TUNED — EACH REPLY IS A FRESH
            GENERATION, NOT A REMEMBERED CONVERSATION · RUNS FULLY ON YOUR HARDWARE
        </div>
        """)

    # css/head/theme live on launch(), not Blocks(), as of Gradio 6 --
    # verified against the actually-installed version rather than assumed.
    demo.launch(
        share=True, server_name="0.0.0.0", server_port=7860,
        css=ARYA_CSS, head=ARYA_HEAD,
        theme=gr.themes.Base(primary_hue=gr.themes.colors.orange,
                              neutral_hue=gr.themes.colors.slate),
    )

if __name__ == "__main__":
    main()