from __future__ import annotations

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

import time
from pathlib import Path
from functools import partial

import numpy as np
import jax
import jax.numpy as jnp
import keras
import gradio as gr
from fastapi import FastAPI
import uvicorn

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,
)

# ============================================================
# Config
# ============================================================

WEIGHTS_PATH = "veylon_final.weights.h5"
MODEL_NAME = "Arya"
MODEL_TAGLINE = "27M-param transformer, made by Arush Kumar"

# Chat templating for a BASE (non-instruction-tuned) language model. If
# Arya/Veylon was trained on plain prose (TinyStories/WikiText/etc.) rather
# than dialogue-formatted data, this template is best-effort — multi-turn
# coherence depends entirely on what the training corpus looked like, not
# on anything this app can fix. Swap this format if you fine-tune on a
# specific chat template later.
USER_TAG = "User:"
ASSISTANT_TAG = "Assistant:"


# ============================================================
# Initialize (runs once, at import time)
# ============================================================

keras.mixed_precision.set_global_policy("mixed_bfloat16")

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

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")

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,
)

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

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")


# ============================================================
# JIT-compiled generation steps — the real speed lever here.
# ============================================================
# The original app called model.generate_step() as plain, untraced Python
# on every single token — meaning JAX re-dispatches and re-optimizes the
# whole graph every call instead of running a compiled executable. This is
# the single biggest speed cost in the original code, and (unlike swapping
# inference engines) is something we can fix directly: jax.jit traces the
# step function once per unique input shape/dtype, caches the compiled
# XLA executable, and every subsequent call re-uses it. Since prefill
# always sees shape (1, CONTEXT) and decode always sees shape (1, 1),
# there are exactly two shapes to compile — both get cached on first use,
# after which every call is a compiled dispatch instead of a Python/XLA
# re-trace.
#
# NOTE: this assumes model.generate_step is a pure function of its
# arguments (no hidden Python-side mutation of the cache/model state
# outside of what it returns). That's true for essentially all KV-cache
# transformer implementations, but if your model.generate_step does
# something stateful and doesn't work under jit, set JIT_GENERATE=False
# below and it'll fall back to the original eager behavior with zero
# other changes.

JIT_GENERATE = True

if JIT_GENERATE:
    _generate_step_jit = jax.jit(model.generate_step)
else:
    _generate_step_jit = model.generate_step


def _run_generate_step(tokens_arr, cache_k, cache_v, cache_pos):
    return _generate_step_jit(
        tokens_arr,
        cache_k=cache_k,
        cache_v=cache_v,
        cache_pos=cache_pos,
    )


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

def sample_from_logits(
    logits: np.ndarray,
    temperature: float = 0.8,
    top_k: int = 50,
) -> int:
    """NumPy-only sampling."""
    logits = np.array(logits, dtype=np.float32, copy=True)

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

    if top_k > 0:
        k = min(int(top_k), logits.shape[-1])
        row = logits[0]
        top_indices = np.argpartition(row, -k)[-k:]
        filtered = np.full_like(row, -np.inf)
        filtered[top_indices] = row[top_indices]
        logits[0] = filtered

    row = logits[0]
    row = row - np.max(row)
    probs = np.exp(row)
    probs = probs / probs.sum()

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


# ============================================================
# Prompt construction for chat-style multi-turn context
# ============================================================

def build_prompt(message: str, history: list[dict]) -> str:
    """
    Turns Gradio chat history (list of {"role", "content"} dicts, the
    'messages' format) into a single text block for a base LM, truncated
    to fit CONTEXT. Oldest turns are dropped first if it doesn't fit —
    we build from the end backwards.
    """
    turns = []
    for turn in history:
        role = turn.get("role")
        content = turn.get("content", "")
        tag = USER_TAG if role == "user" else ASSISTANT_TAG
        turns.append(f"{tag} {content}")
    turns.append(f"{USER_TAG} {message}")
    turns.append(ASSISTANT_TAG)

    # Build from the most recent turn backwards until we run out of
    # context budget, so we never silently truncate the *current* message.
    budget_chars = CONTEXT * 4  # rough chars-per-token heuristic pre-tokenize
    kept = []
    running_len = 0
    for t in reversed(turns):
        running_len += len(t) + 1
        if running_len > budget_chars and kept:
            break
        kept.append(t)
    kept.reverse()
    return "\n".join(kept)


# ============================================================
# Generation (streaming generator for gr.ChatInterface)
# ============================================================

def generate_stream(
    message: str,
    history: list[dict],
    max_new_tokens: int,
    temperature: float,
    top_k: int,
):
    """
    Streaming generator: yields the partial assistant response as tokens
    are produced, so the chat UI fills in live instead of blocking on the
    whole generation.
    """
    try:
        prompt = build_prompt(message, history)

        tokens = tokenizer.encode(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:]

        prompt_len = len(tokens)
        prompt_ids = np.array([tokens], dtype=np.int32)

        t0 = time.time()
        logits, cache_k, cache_v = _run_generate_step(
            prompt_ids, cache_k=None, cache_v=None, cache_pos=0,
        )

        next_token = sample_from_logits(
            np.array(logits[:, -1, :], dtype=np.float32, copy=True),
            temperature=temperature,
            top_k=top_k,
        )

        generated_ids = [next_token]
        partial_text = tokenizer.decode(generated_ids)
        yield partial_text

        if next_token != tokenizer.eos_id and prompt_len < CONTEXT:
            cache_pos = prompt_len

            for _ in range(int(max_new_tokens) - 1):
                next_input = np.array([[next_token]], dtype=np.int32)

                logits, cache_k, cache_v = _run_generate_step(
                    next_input, cache_k=cache_k, cache_v=cache_v, cache_pos=cache_pos,
                )
                cache_pos += 1

                next_token = sample_from_logits(
                    np.array(logits[:, -1, :], dtype=np.float32, copy=True),
                    temperature=temperature,
                    top_k=top_k,
                )

                if next_token == tokenizer.eos_id:
                    break

                generated_ids.append(next_token)
                partial_text = tokenizer.decode(generated_ids)
                yield partial_text

                if prompt_len + len(generated_ids) >= CONTEXT:
                    break

        elapsed = time.time() - t0
        n_tok = len(generated_ids)
        tok_per_sec = n_tok / elapsed if elapsed > 0 else 0.0
        print(f"Generated {n_tok} tokens in {elapsed:.2f}s ({tok_per_sec:.1f} tok/s)")

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


# ============================================================
# Styling — dark terminal / glassmorphism, JetBrains Mono, neon accent
# ============================================================

CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;700&display=swap');

:root {
    --veylon-bg: #0a0e14;
    --veylon-panel: rgba(20, 26, 38, 0.65);
    --veylon-border: rgba(0, 255, 200, 0.18);
    --veylon-accent: #00ffc8;
    --veylon-accent-dim: #00b894;
    --veylon-text: #e6f1ff;
    --veylon-text-dim: #7d8ba1;
}

* {
    font-family: 'JetBrains Mono', ui-monospace, monospace !important;
}

body, .gradio-container {
    background: radial-gradient(circle at 20% 10%, #101826 0%, #0a0e14 60%) !important;
    color: var(--veylon-text) !important;
}

#veylon-header {
    text-align: center;
    padding: 18px 12px 6px 12px;
}

#veylon-header h1 {
    font-size: 1.8rem;
    font-weight: 700;
    letter-spacing: 0.5px;
    background: linear-gradient(90deg, var(--veylon-accent), #4dd0ff);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 2px;
}

#veylon-header p {
    color: var(--veylon-text-dim);
    font-size: 0.9rem;
    margin: 0;
}

.gr-panel, .block, .form {
    background: var(--veylon-panel) !important;
    border: 1px solid var(--veylon-border) !important;
    border-radius: 14px !important;
    backdrop-filter: blur(10px);
}

#chatbot {
    border-radius: 16px !important;
    border: 1px solid var(--veylon-border) !important;
    background: var(--veylon-panel) !important;
    box-shadow: 0 0 24px rgba(0, 255, 200, 0.05);
}

.message.user {
    background: rgba(0, 255, 200, 0.08) !important;
    border: 1px solid var(--veylon-border) !important;
}

.message.bot {
    background: rgba(255, 255, 255, 0.03) !important;
    border: 1px solid rgba(255, 255, 255, 0.06) !important;
}

button.primary, .gr-button-primary {
    background: linear-gradient(90deg, var(--veylon-accent-dim), var(--veylon-accent)) !important;
    color: #04120e !important;
    border: none !important;
    font-weight: 700 !important;
    box-shadow: 0 0 16px rgba(0, 255, 200, 0.25);
    transition: transform 0.12s ease, box-shadow 0.12s ease;
}

button.primary:hover, .gr-button-primary:hover {
    transform: translateY(-1px);
    box-shadow: 0 0 24px rgba(0, 255, 200, 0.4);
}

input, textarea, .gr-box {
    background: rgba(255, 255, 255, 0.03) !important;
    color: var(--veylon-text) !important;
    border: 1px solid var(--veylon-border) !important;
    border-radius: 10px !important;
}

label span {
    color: var(--veylon-text-dim) !important;
    font-size: 0.78rem !important;
    text-transform: uppercase;
    letter-spacing: 0.06em;
}

#stats-badge {
    text-align: center;
    color: var(--veylon-text-dim);
    font-size: 0.78rem;
    padding-top: 4px;
}

footer { display: none !important; }
"""

THEME = gr.themes.Base(
    primary_hue="teal",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
).set(
    body_background_fill="*neutral_950",
    block_background_fill="*neutral_900",
    border_color_primary="*primary_700",
)


# ============================================================
# Build the Gradio Blocks app
# ============================================================

def build_demo() -> gr.Blocks:
    with gr.Blocks(title=MODEL_NAME, theme=THEME, css=CUSTOM_CSS) as demo:
        gr.HTML(f"""
        <div id="veylon-header">
            <h1>◆ {MODEL_NAME}</h1>
            <p>{MODEL_TAGLINE} · {model.count_params():,} params · {CONTEXT} ctx</p>
        </div>
        """)

        with gr.Row():
            max_tokens = gr.Slider(
                label="Max new tokens", minimum=10, maximum=768, value=64, step=10,
            )
            temperature = gr.Slider(
                label="Temperature", minimum=0.1, maximum=2.0, value=0.8, step=0.1,
            )
            top_k = gr.Slider(
                label="Top-K", minimum=1, maximum=100, value=50, step=1,
            )

        gr.ChatInterface(
            fn=generate_stream,
            type="messages",
            additional_inputs=[max_tokens, temperature, top_k],
            chatbot=gr.Chatbot(elem_id="chatbot", type="messages", height=480),
            examples=[
                ["Once upon a time,"],
                ["The scientist looked at the data and"],
                ["In a small village, there lived"],
            ],
            title=None,
        )

        gr.HTML(f"""
        <div id="stats-badge">
            base language model · not instruction-tuned · coherence across
            turns depends on training data format
        </div>
        """)

    return demo


demo = build_demo()
demo.queue()  # required for streaming generators to work over the wire


# ============================================================
# FastAPI mount — real production entrypoint instead of demo.launch()
# ============================================================

app = FastAPI(title=f"{MODEL_NAME} API")


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model": MODEL_NAME,
        "params": model.count_params(),
        "context": CONTEXT,
        "backend": keras.backend.backend(),
        "jax_devices": [str(d) for d in jax.devices()],
    }


app = gr.mount_gradio_app(app, demo, path="/")


if __name__ == "__main__":
    # Run via uvicorn directly rather than demo.launch(): gives you
    # multi-worker / reload / proper ASGI server config for actual
    # deployment, plus the /health route above for load-balancer checks.
    uvicorn.run(app, host="0.0.0.0", port=7860)