# Import before torch/CUDA initialization on Hugging Face Spaces.
import spaces

import gc
import queue
import re
import threading
import traceback

import gradio as gr
import torch
from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForCausalLM,
    TextIteratorStreamer,
)

# ============================================================
# CONFIGURATION
# ============================================================

MODELS = {
    "BrtGPT-1-Pre": "Bertug1911/BrtGPT-1-Pre",
    "BrtGPT-1-Pre-Code": "Bertug1911/BrtGPT-1-Pre-Code",
    "BrtGPT-1-0719": "HuggingFaceTB/SmolLM2-1.7B-Instruct",
    "BrtGPT-1-0906": "Bertug1911/BrtGPT-1-0906",
    "BrtGPT-Conversation-1": "Bertug1911/BrtGPT-Conversation-1",
    "BrtGPT-Conversation-1.5": "Bertug1911/BrtGPT-Conversation-1.5",
    "BrtGPT-Conversation-2": "Bertug1911/BrtGPT-Conversation-2",
    "CodePiegon": "Bertug1911/CodePiegon",
}
DEFAULT_MODEL = "BrtGPT-1-Pre"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
REPETITION_PENALTY = 1.1

# Per-model maximums. Global minimums: top_p=0.1, top_k=1,
# temperature=0.1, max_new_tokens=1. Adjust these limits as desired.
MODEL_LIMITS = {
    "BrtGPT-1-Pre":             {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-1-Pre-Code":        {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-1-0719":            {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-1-0906":            {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-Conversation-1":    {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-Conversation-1.5":  {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "BrtGPT-Conversation-2":    {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 768},
    "CodePiegon":               {"top_p_max": 1.0, "top_k_max": 50, "temperature_max": 2.0, "max_new_tokens_max": 128},
}
DEFAULTS = {"top_p": 0.75, "top_k": 12, "temperature": 0.7, "max_new_tokens": 128}

_current_model_id = None
_current_model = None
_current_tokenizer = None
_model_lock = threading.Lock()


def get_limits(model_name):
    return MODEL_LIMITS.get(model_name, MODEL_LIMITS[DEFAULT_MODEL])


def update_model_controls(model_name):
    limits = get_limits(model_name)
    return (
        gr.update(minimum=0.1, maximum=limits["top_p_max"], value=min(DEFAULTS["top_p"], limits["top_p_max"])),
        gr.update(minimum=1, maximum=limits["top_k_max"], value=min(DEFAULTS["top_k"], limits["top_k_max"]), step=1),
        gr.update(minimum=0.1, maximum=limits["temperature_max"], value=min(DEFAULTS["temperature"], limits["temperature_max"])),
        gr.update(minimum=1, maximum=limits["max_new_tokens_max"], value=min(DEFAULTS["max_new_tokens"], limits["max_new_tokens_max"]), step=1),
    )


# ============================================================
# MODEL LOADING + TOKEN ID SAFETY CHECKS
# ============================================================

def load_model(model_name):
    global _current_model_id, _current_model, _current_tokenizer

    if model_name not in MODELS:
        raise ValueError(f"Unknown model: {model_name}")

    model_id = MODELS[model_name]
    if _current_model_id == model_id and _current_model is not None:
        return _current_tokenizer, _current_model

    # Release the previous model before loading another one.
    if _current_model is not None:
        del _current_model
    _current_model = None
    _current_tokenizer = None
    _current_model_id = None
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    print(f"Loading tokenizer: {model_id}")
    tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
    config = AutoConfig.from_pretrained(model_id)

    # Important: if the tokenizer has more IDs than the model config,
    # set the config vocabulary size BEFORE loading weights. This avoids
    # the common 32000-vs-32001 embedding shape mismatch when the checkpoint
    # actually contains the extra embedding row.
    tokenizer_size = len(tokenizer)
    if tokenizer_size > int(config.vocab_size):
        print(f"Adjusting config vocab_size: {config.vocab_size} -> {tokenizer_size}")
        config.vocab_size = tokenizer_size

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        config=config,
        torch_dtype=DTYPE,
        low_cpu_mem_usage=True,
    )

    embedding_count = model.get_input_embeddings().num_embeddings
    if tokenizer_size > embedding_count:
        # Defensive fallback for architectures whose embedding table was not
        # sized from config. New rows are initialized; see note in UI errors.
        model.resize_token_embeddings(tokenizer_size)
        embedding_count = model.get_input_embeddings().num_embeddings

    # Validate all tokenizer IDs before the first CUDA operation.
    vocab = tokenizer.get_vocab()
    invalid = [(token, idx) for token, idx in vocab.items() if idx < 0 or idx >= embedding_count]
    if invalid:
        del model
        gc.collect()
        raise ValueError(
            f"Tokenizer/model vocabulary mismatch for {model_name}: {len(invalid)} token IDs "
            f"are outside embedding range [0, {embedding_count - 1}]. "
            "Check that this model repository's tokenizer and weights come from the same revision."
        )

    if tokenizer.pad_token_id is None:
        if tokenizer.eos_token_id is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
            if len(tokenizer) > model.get_input_embeddings().num_embeddings:
                model.resize_token_embeddings(len(tokenizer))

    model.to(DEVICE)
    model.eval()

    _current_tokenizer = tokenizer
    _current_model = model
    _current_model_id = model_id
    print(
        f"Model ready: {model_id} on {DEVICE} | tokenizer={len(tokenizer)} | "
        f"embeddings={model.get_input_embeddings().num_embeddings}"
    )
    return tokenizer, model


# ============================================================
# OUTPUT FORMATTING
# ============================================================

CODE_BLOCK_PATTERN = re.compile(r"```([^\n`]*)\n(.*?)```", re.DOTALL)


def format_output(text):
    def replace_code_block(match):
        language = match.group(1).strip() or "text"
        code = match.group(2).rstrip("\n")
        longest_run = max((len(run) for run in re.findall(r"`+", code)), default=0)
        fence = "`" * max(3, longest_run + 1)
        return f"\n\n### `<{language.upper()}>`\n\n{fence}{language}\n{code}\n{fence}\n\n"

    return CODE_BLOCK_PATTERN.sub(replace_code_block, text).strip() or "The model returned an empty response."


# Preserve the original marker cleanup behavior, but only when those markers
# actually appear in the model output.
def normalize_token_markers(text):
    if "Ġ" in text or "Ċ" in text:
        text = text.replace(" ", "")
        text = text.replace("Ġ", " ").replace("Ċ", "\n")
    return text


# ============================================================
# GENERATION
# ============================================================

@spaces.GPU
def generate(prompt, model_name, top_p, top_k, temperature, max_new_tokens):
    if not prompt or not prompt.strip():
        yield "Please enter a message."
        return
    if model_name not in MODELS:
        yield "Please select a valid model."
        return

    limits = get_limits(model_name)
    # Clamp all UI values again server-side; browser inputs are not trusted.
    top_p = max(0.1, min(float(top_p), limits["top_p_max"]))
    top_k = max(1, min(int(top_k), limits["top_k_max"]))
    temperature = max(0.1, min(float(temperature), limits["temperature_max"]))
    max_new_tokens = max(1, min(int(max_new_tokens), limits["max_new_tokens_max"]))

    yield "Loading model and generating..."

    try:
        # Keep the lock for the complete operation; only one request/model
        # generation runs at once, avoiding unsafe concurrent CUDA use.
        with _model_lock:
            tokenizer, model = load_model(model_name)
            messages = [{"role": "user", "content": prompt.strip()}]

            if getattr(tokenizer, "chat_template", None):
                inputs = tokenizer.apply_chat_template(
                    messages,
                    tokenize=True,
                    add_generation_prompt=True,
                    return_tensors="pt",
                    return_dict=True,
                )
            else:
                inputs = tokenizer(prompt.strip(), return_tensors="pt")

            inputs = {key: value.to(DEVICE) for key, value in inputs.items()}
            input_ids = inputs.get("input_ids")
            if input_ids is None or input_ids.numel() == 0:
                raise ValueError("Tokenizer returned empty input_ids.")

            embedding_count = model.get_input_embeddings().num_embeddings
            min_id = int(input_ids.min().item())
            max_id = int(input_ids.max().item())
            if min_id < 0 or max_id >= embedding_count:
                raise ValueError(
                    f"Input token ID range [{min_id}, {max_id}] exceeds model embedding range "
                    f"[0, {embedding_count - 1}]. Tokenizer and model files are incompatible."
                )

            model_context = getattr(model.config, "max_position_embeddings", None)
            if model_context is None:
                model_context = getattr(model.config, "n_positions", None)
            if isinstance(model_context, int) and model_context > 1:
                # Keep the prompt inside the model context and leave room for generation.
                # This also prevents position-embedding index errors on very long prompts.
                max_input_tokens = model_context - 1
                if input_ids.shape[-1] > max_input_tokens:
                    for key in ("input_ids", "attention_mask", "token_type_ids"):
                        if key in inputs and inputs[key].ndim == 2:
                            inputs[key] = inputs[key][:, -max_input_tokens:]
                    input_ids = inputs["input_ids"]
                    print(
                        f"Prompt truncated to {max_input_tokens} tokens for model context "
                        f"limit {model_context}."
                    )
                max_new_tokens = min(max_new_tokens, max(1, model_context - input_ids.shape[-1]))

            pad_token_id = tokenizer.pad_token_id
            if pad_token_id is None:
                pad_token_id = tokenizer.eos_token_id

            generation_kwargs = {
                "max_new_tokens": max_new_tokens,
                "do_sample": True,
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k,
                "repetition_penalty": REPETITION_PENALTY,
                "pad_token_id": pad_token_id,
                "use_cache": True,
            }
            if tokenizer.eos_token_id is not None:
                generation_kwargs["eos_token_id"] = tokenizer.eos_token_id

            streamer = TextIteratorStreamer(
                tokenizer,
                skip_prompt=True,
                skip_special_tokens=True,
                timeout=0.5,
            )
            worker_errors = []

            def generation_worker():
                try:
                    with torch.inference_mode():
                        model.generate(**inputs, streamer=streamer, **generation_kwargs)
                except Exception as exc:
                    worker_errors.append(exc)
                    try:
                        streamer.on_finalized_text("", stream_end=True)
                    except Exception:
                        pass

            worker = threading.Thread(target=generation_worker, daemon=True)
            worker.start()
            raw_response = ""

            while True:
                try:
                    chunk = next(streamer)
                except StopIteration:
                    break
                except queue.Empty:
                    if not worker.is_alive():
                        if worker_errors:
                            raise worker_errors[0]
                        break
                    continue

                if chunk:
                    raw_response += chunk
                    yield format_output(normalize_token_markers(raw_response))

            worker.join()
            if worker_errors:
                raise worker_errors[0]
            yield format_output(normalize_token_markers(raw_response))

    except torch.cuda.OutOfMemoryError:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        yield "GPU memory is insufficient. Select a smaller model or reduce Max New Tokens."
    except Exception as exc:
        print("Generation error:", repr(exc))
        traceback.print_exc()
        message = str(exc)
        if "device-side assert" in message.lower() or "cuda error" in message.lower():
            yield (
                "CUDA generation failed. The app checked token IDs before generation, so this may be "
                "a model/checkpoint incompatibility or a CUDA kernel issue. Restart the Space runtime "
                "and verify that the model weights and tokenizer are from the same repository revision."
            )
        else:
            yield f"Generation error ({type(exc).__name__}): {message}"


# ============================================================
# DARK THEME
# ============================================================

CSS = """
.gradio-container { max-width: 940px !important; margin: 0 auto !important; background: #0b0d12 !important; }
#app-header { text-align: center; padding: 22px 0 8px 0; }
#app-header h1 { font-size: 30px; font-weight: 750; letter-spacing: -1px; color: #f1f5f9; }
#app-subtitle { text-align: center; color: #929bad; margin-bottom: 22px; }
#model-selector, #settings-panel { border: 1px solid #292e3b; border-radius: 12px; background: #11141c; }
#input-area textarea { background: #11141c !important; border: 1px solid #292e3b !important; border-radius: 12px !important; color: #f1f5f9 !important; font-size: 15px !important; line-height: 1.7 !important; }
#output-area { border: 1px solid #292e3b; border-radius: 12px; background: #11141c; padding: 18px; }
#output-area h3 { color: #a5b4fc; font-size: 12px; font-weight: 700; letter-spacing: 1px; margin-bottom: 10px; }
#output-area pre { border: 1px solid #292e3b !important; border-radius: 10px !important; background: #090b10 !important; padding: 16px !important; overflow-x: auto !important; }
#output-area code { font-size: 13px !important; line-height: 1.7 !important; }
#generate-button { border-radius: 10px; min-height: 46px; font-weight: 650; }
footer { display: none !important; }
"""

with gr.Blocks(
    theme=gr.themes.Base(
        primary_hue="indigo",
        secondary_hue="slate",
        neutral_hue="slate",
    ).set(
        body_background_fill="#0b0d12",
        body_text_color="#e5e7eb",
        block_background_fill="#11141c",
        block_border_color="#292e3b",
        input_background_fill="#11141c",
        button_primary_background_fill="#6366f1",
        button_primary_background_fill_hover="#4f46e5",
    ),
    css=CSS,
    title="LLM Playground",
) as demo:
    gr.Markdown("# LLM Playground", elem_id="app-header")
    gr.Markdown("A simple, modern interface for language models.", elem_id="app-subtitle")

    model_selector = gr.Dropdown(
        choices=list(MODELS.keys()),
        value=DEFAULT_MODEL,
        label="MODEL",
        info="Choose which model to run.",
        elem_id="model-selector",
        interactive=True,
    )

    prompt_input = gr.Textbox(
        label="INPUT",
        placeholder="Ask a question, write code, or enter a prompt...",
        lines=5,
        max_lines=12,
        elem_id="input-area",
    )

    with gr.Accordion("Generation settings", open=True, elem_id="settings-panel"):
        with gr.Row():
            top_p_input = gr.Slider(
                minimum=0.1, maximum=get_limits(DEFAULT_MODEL)["top_p_max"], value=DEFAULTS["top_p"],
                step=0.05, label="Top P", info="Minimum 0.1; maximum depends on selected model."
            )
            top_k_input = gr.Slider(
                minimum=1, maximum=get_limits(DEFAULT_MODEL)["top_k_max"], value=DEFAULTS["top_k"],
                step=1, label="Top K", info="Minimum 1; maximum depends on selected model."
            )
        with gr.Row():
            temperature_input = gr.Slider(
                minimum=0.1, maximum=get_limits(DEFAULT_MODEL)["temperature_max"], value=DEFAULTS["temperature"],
                step=0.05, label="Temperature", info="Minimum 0.1; maximum depends on selected model."
            )
            max_new_tokens_input = gr.Slider(
                minimum=1, maximum=get_limits(DEFAULT_MODEL)["max_new_tokens_max"], value=DEFAULTS["max_new_tokens"],
                step=1, label="Max New Tokens", info="Minimum 1; maximum depends on selected model."
            )

    generate_button = gr.Button("Generate", variant="primary", elem_id="generate-button")
    output = gr.Markdown(value="Your model's response will appear here.", label="OUTPUT", elem_id="output-area")

    model_selector.change(
        fn=update_model_controls,
        inputs=[model_selector],
        outputs=[top_p_input, top_k_input, temperature_input, max_new_tokens_input],
    )

    generation_inputs = [
        prompt_input, model_selector, top_p_input, top_k_input,
        temperature_input, max_new_tokens_input,
    ]
    generate_button.click(fn=generate, inputs=generation_inputs, outputs=output)
    prompt_input.submit(fn=generate, inputs=generation_inputs, outputs=output)

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