"""
Gradio chat interface for nanochat models on Hugging Face Spaces.
Downloads the model from HF Hub and serves a streaming chat UI.
"""

import os
import json
import random
import torch
import gradio as gr
from huggingface_hub import snapshot_download

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_REPO = os.environ.get("MODEL_REPO", "Slothwolf/MiniChat-0.9B")
MODEL_STEP = int(os.environ.get("MODEL_STEP", "486"))
DEFAULT_TEMPERATURE = float(os.environ.get("DEFAULT_TEMPERATURE", "0.8"))
DEFAULT_MAX_TOKENS = int(os.environ.get("DEFAULT_MAX_TOKENS", "1024"))
DEFAULT_TOP_K = int(os.environ.get("DEFAULT_TOP_K", "50"))

# ---------------------------------------------------------------------------
# Download and load model
# ---------------------------------------------------------------------------
print(f"Downloading model from {MODEL_REPO}...")
model_dir = snapshot_download(MODEL_REPO)

# Point nanochat at the downloaded directory so get_tokenizer() finds tokenizer/
os.environ["NANOCHAT_BASE_DIR"] = model_dir

from nanochat.checkpoint_manager import build_model
from nanochat.engine import Engine

print(f"Loading model (step {MODEL_STEP}) on CPU...")
device = torch.device("cpu")
model, tokenizer, meta = build_model(model_dir, MODEL_STEP, device, "eval")
engine = Engine(model, tokenizer)
print("Model loaded and ready!")

# Grab special token ids once
bos_id = tokenizer.get_bos_token_id()
user_start = tokenizer.encode_special("<|user_start|>")
user_end = tokenizer.encode_special("<|user_end|>")
assistant_start = tokenizer.encode_special("<|assistant_start|>")
assistant_end = tokenizer.encode_special("<|assistant_end|>")




def normalize_history(history):
    if not history:
        return []

    if isinstance(history[0], dict):
        pairs = []
        user_msg = None
        for msg in history:
            role = msg.get("role")
            content = msg.get("content", "")

            if not isinstance(content, str):
                content = str(content) if content is not None else ""

            if role == "user":
                if user_msg is not None:
                    pairs.append([user_msg, None])
                user_msg = content
            elif role == "assistant":
                if user_msg is not None:
                    pairs.append([user_msg, content])
                    user_msg = None
                else:  
                    pairs.append(["", content])
        if user_msg is not None:
            pairs.append([user_msg, None])
        return pairs

    normalized = []
    for turn in history:
        if isinstance(turn, (list, tuple)):
            user = turn[0] if len(turn) > 0 else None
            assistant = turn[1] if len(turn) > 1 else None

            user = str(user) if user is not None else ""
            assistant = str(assistant) if assistant is not None else ""
            normalized.append([user, assistant])
        else:
            continue
    return normalized


# ---------------------------------------------------------------------------
# Chat function (streaming)
# ---------------------------------------------------------------------------
def chat(message, history, temperature, max_tokens, top_k):
    """Generate a streaming response for the chat interface."""

    history = normalize_history(history)
    # Build the full token sequence from conversation history
    # history is a list of [user_msg, assistant_msg] tuples
    tokens = [bos_id]


    for user_msg, assistant_msg in history:
        user_text = user_msg if isinstance(user_msg, str) else str(user_msg) if user_msg is not None else ""
        tokens.append(user_start)
        tokens.extend(tokenizer.encode(user_msg))
        tokens.append(user_end)
        if assistant_msg:
            tokens.append(assistant_start)
            tokens.extend(tokenizer.encode(assistant_msg))
            tokens.append(assistant_end)

    # Add the current user message and prime the assistant
    tokens.append(user_start)
    tokens.extend(tokenizer.encode(message))
    tokens.append(user_end)
    tokens.append(assistant_start)

    # Generate tokens with streaming
    response = ""
    accumulated_tokens = []
    last_clean_text = ""

    for token_column, token_masks in engine.generate(
        tokens,
        num_samples=1,
        max_tokens=int(max_tokens),
        temperature=float(temperature),
        top_k=int(top_k) if int(top_k) > 0 else None,
        seed=random.randint(0, 2**31 - 1),
    ):
        token = token_column[0]
        if token == assistant_end or token == bos_id:
            break

        # Accumulate tokens and decode, handling multi-byte UTF-8
        accumulated_tokens.append(token)
        current_text = tokenizer.decode(accumulated_tokens)
        if not current_text.endswith("\ufffd"):
            new_text = current_text[len(last_clean_text):]
            if new_text:
                response += new_text
                last_clean_text = current_text
                yield response

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
demo = gr.ChatInterface(
    chat,
    additional_inputs=[
        gr.Slider(0.0, 2.0, value=DEFAULT_TEMPERATURE, step=0.1, label="Temperature"),
        gr.Slider(1, 2048, value=DEFAULT_MAX_TOKENS, step=1, label="Max Tokens"),
        gr.Slider(0, 200, value=DEFAULT_TOP_K, step=1, label="Top-K (0 = disabled)"),
    ],
    title="MiniChat",
    description=(
        "Chat with a **0.9B parameter** GPT-like model trained with "
        "[nanochat](https://github.com/karpathy/nanochat) by Andrej Karpathy. "
        "Running on CPU — generation will be slow."
    ),
)

if __name__ == "__main__":
    demo.launch()