# MiniCPM5-1B Demo

from pathlib import Path
import os
import time
import logging
import threading

import gradio as gr
import spaces
import torch
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

from utils_chatbot import organize_messages_from_messages, stream2display_text

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_PATH = "openbmb/MiniCPM5-1B"

hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    logger.info("Logged in to Hugging Face Hub")
else:
    logger.warning("HF_TOKEN not set — private/gated models will be inaccessible")

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
).to("cuda")


@spaces.GPU(duration=60)
def gpu_generate_stream(inputs, history, temperature, top_p, enable_think):
    prompt_text = tokenizer.apply_chat_template(
        inputs,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=enable_think,
    )
    model_inputs = tokenizer([prompt_text], return_tensors="pt").to("cuda")

    history.append({"role": "assistant", "content": ""})
    yield "", history

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=False,
    )

    gen_kwargs = dict(
        **model_inputs,
        streamer=streamer,
        max_new_tokens=4096,
    )
    if temperature > 0:
        gen_kwargs.update(temperature=temperature, top_p=top_p, do_sample=True)
    else:
        gen_kwargs.update(do_sample=False)

    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    stream_text = ""
    gen_tk_count = 0
    start_time = time.time()

    for new_token_text in streamer:
        if not new_token_text:
            continue
        stream_text += new_token_text
        gen_tk_count += 1
        elapsed = time.time() - start_time
        token_per_sec = gen_tk_count / elapsed if elapsed > 0 else 0
        display_text = stream2display_text(stream_text, token_per_sec)
        history[-1]["content"] = display_text
        yield "", history

    thread.join()
    history[-1]["content"] = stream_text.replace("<|im_end|>", "")
    yield "", history


def gen_response_stream(message, history, temperature, top_p, enable_think):
    chat_msg_ls = organize_messages_from_messages(message, history)
    history.append({"role": "user", "content": message})
    yield from gpu_generate_stream(
        chat_msg_ls, history,
        temperature=temperature,
        top_p=top_p,
        enable_think=enable_think,
    )


def create_app():
    assets_path = Path.cwd().absolute() / "assets"
    gr.set_static_paths(paths=[assets_path])

    with gr.Blocks() as demo:
        # Header
        gr.HTML(
            '<div class="app-header">'
            '  <div class="header-left">'
            '    <img src="/gradio_api/file=assets/OpenBMB-MiniCPM.png" alt="Logo" class="header-logo">'
            '    <div class="header-title">'
            '      <span class="title-main">MiniCPM5-1B</span>'
            '      <span class="title-sub">By OpenBMB</span>'
            '    </div>'
            '  </div>'
            '  <div class="header-badge">'
            '    <span class="status-dot"></span>'
            '    MiniCPM5-1B'
            '  </div>'
            '</div>'
        )

        with gr.Row(elem_classes=["main-row"]):
            with gr.Column(scale=4, elem_classes=["chat-col"]):
                chatbot = gr.Chatbot(
                    show_label=False,
                    placeholder="Send a message to start chatting...",
                    height="70vh",
                    elem_classes=["dark-chatbot"],
                )
                prompt = gr.Textbox(
                    show_label=False,
                    placeholder="Ask MiniCPM5...",
                    lines=1,
                    elem_classes=["input-pill"],
                )

            with gr.Column(scale=1, min_width=220, elem_classes=["settings-col"]):
                gr.HTML('<div class="settings-title">Settings</div>')
                temperature = gr.Slider(
                    minimum=0, maximum=1, value=0.9, step=0.05,
                    label="Temperature",
                    elem_classes=["dark-slider"],
                )
                top_p = gr.Slider(
                    minimum=0, maximum=1, value=0.95, step=0.01,
                    label="Top-p",
                    elem_classes=["dark-slider"],
                )
                enable_think = gr.Checkbox(
                    label="Enable Thinking",
                    value=True,
                    elem_classes=["dark-checkbox"],
                )
                clear = gr.Button(
                    "Clear History",
                    variant="secondary",
                    elem_classes=["clear-btn"],
                )

        prompt.submit(
            gen_response_stream,
            inputs=[prompt, chatbot, temperature, top_p, enable_think],
            outputs=[prompt, chatbot],
        )
        clear.click(lambda: None, None, chatbot, queue=False)

    return demo


THEME = gr.themes.Base(
    primary_hue=gr.themes.Color(
        c50="#eef2ff", c100="#e0e7ff", c200="#c7d2fe", c300="#a5b4fc",
        c400="#818cf8", c500="#6366f1", c600="#4f46e5", c700="#4338ca",
        c800="#3730a3", c900="#312e81", c950="#1e1b4b",
    ),
    secondary_hue="slate",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
).set(
    body_background_fill="#05070A",
    body_background_fill_dark="#05070A",
    body_text_color="#e2e8f0",
    body_text_color_dark="#e2e8f0",
    block_background_fill="rgba(255,255,255,0.03)",
    block_background_fill_dark="rgba(255,255,255,0.03)",
    block_border_color="rgba(255,255,255,0.08)",
    block_border_color_dark="rgba(255,255,255,0.08)",
    block_label_text_color="#94a3b8",
    block_label_text_color_dark="#94a3b8",
    block_title_text_color="#e2e8f0",
    block_title_text_color_dark="#e2e8f0",
    input_background_fill="rgba(255,255,255,0.04)",
    input_background_fill_dark="rgba(255,255,255,0.04)",
    input_border_color="rgba(255,255,255,0.08)",
    input_border_color_dark="rgba(255,255,255,0.08)",
    input_placeholder_color="#64748b",
    input_placeholder_color_dark="#64748b",
    button_primary_background_fill="linear-gradient(135deg, #3B5BFF, #27D4EA)",
    button_primary_background_fill_dark="linear-gradient(135deg, #3B5BFF, #27D4EA)",
    button_primary_text_color="#ffffff",
    button_secondary_background_fill="rgba(255,255,255,0.06)",
    button_secondary_background_fill_dark="rgba(255,255,255,0.06)",
    button_secondary_text_color="#94a3b8",
    button_secondary_border_color="rgba(255,255,255,0.1)",
    button_secondary_border_color_dark="rgba(255,255,255,0.1)",
    slider_color="#6366f1",
    slider_color_dark="#6366f1",
    checkbox_background_color="rgba(255,255,255,0.06)",
    checkbox_background_color_dark="rgba(255,255,255,0.06)",
    checkbox_border_color="rgba(255,255,255,0.15)",
    checkbox_border_color_dark="rgba(255,255,255,0.15)",
    checkbox_label_text_color="#cbd5e1",
    checkbox_label_text_color_dark="#cbd5e1",
    shadow_drop="none",
    shadow_drop_lg="none",
)

CSS = """
/* Force dark background everywhere */
html, body, .gradio-container, .main, footer { background: #05070A !important; }
footer { display: none !important; }

/* Header */
.app-header {
    display: flex;
    align-items: center;
    justify-content: space-between;
    padding: 16px 24px;
    border-bottom: 1px solid rgba(255,255,255,0.05);
}
.header-left {
    display: flex;
    align-items: center;
    gap: 12px;
}
.header-logo {
    width: 40px;
    height: 40px;
    border-radius: 10px;
    filter: drop-shadow(0 0 12px rgba(39,212,234,0.35));
}
.header-title {
    display: flex;
    flex-direction: column;
}
.title-main {
    font-size: 18px;
    font-weight: 700;
    letter-spacing: -0.02em;
    background: linear-gradient(to right, #fff, rgba(255,255,255,0.6));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
.title-sub {
    font-size: 10px;
    text-transform: uppercase;
    letter-spacing: 0.2em;
    color: rgba(255,255,255,0.35);
    font-weight: 600;
}
.header-badge {
    display: flex;
    align-items: center;
    gap: 6px;
    font-size: 10px;
    font-weight: 700;
    letter-spacing: 0.12em;
    text-transform: uppercase;
    color: rgba(255,255,255,0.5);
    background: rgba(255,255,255,0.04);
    padding: 6px 14px;
    border-radius: 100px;
    border: 1px solid rgba(255,255,255,0.06);
}
.status-dot {
    width: 6px;
    height: 6px;
    border-radius: 50%;
    background: #27D4EA;
    animation: pulse-dot 2s infinite;
}
@keyframes pulse-dot {
    0%, 100% { opacity: 1; }
    50% { opacity: 0.4; }
}

/* Layout */
.main-row { gap: 0 !important; }
.chat-col { padding-right: 8px !important; }
.settings-col {
    border-left: 1px solid rgba(255,255,255,0.06) !important;
    padding: 20px !important;
}
.settings-title {
    font-size: 14px;
    font-weight: 700;
    color: rgba(255,255,255,0.7);
    margin-bottom: 12px;
    letter-spacing: 0.05em;
    text-transform: uppercase;
}

/* Chatbot dark styling */
.dark-chatbot {
    border: none !important;
    background: transparent !important;
}
.dark-chatbot .wrapper { background: transparent !important; }
.dark-chatbot .bubble-wrap { background: transparent !important; }
.dark-chatbot .message-row {
    animation: msg-fade-in 0.35s cubic-bezier(0.16, 1, 0.3, 1) forwards;
}
@keyframes msg-fade-in {
    from { opacity: 0; transform: translateY(12px); }
    to { opacity: 1; transform: translateY(0); }
}

/* User message bubble */
.dark-chatbot .role-user .message-bubble-border {
    background: linear-gradient(135deg, #3B5BFF, #27D4EA) !important;
    border: none !important;
    border-radius: 20px 20px 4px 20px !important;
    box-shadow: 0 6px 20px rgba(59,91,255,0.15);
}
.dark-chatbot .role-user .message-bubble-border .message-content {
    color: #ffffff !important;
}

/* Bot message bubble */
.dark-chatbot .role-assistant .message-bubble-border {
    background: rgba(255,255,255,0.03) !important;
    border: 1px solid rgba(255,255,255,0.08) !important;
    border-radius: 20px 20px 20px 4px !important;
    backdrop-filter: blur(8px);
}
.dark-chatbot .role-assistant .message-bubble-border .message-content {
    color: #e2e8f0 !important;
}

/* Thinking block (blockquote in bot messages) */
.dark-chatbot .role-assistant blockquote {
    background: rgba(59,91,255,0.06) !important;
    border-left: 3px solid #3B5BFF !important;
    border-radius: 4px 10px 10px 4px;
    padding: 10px 14px !important;
    color: #8B949E !important;
    font-style: italic;
    margin-bottom: 10px !important;
}

/* Input pill */
.input-pill {
    border: 1px solid rgba(255,255,255,0.08) !important;
    border-radius: 24px !important;
    background: rgba(255,255,255,0.04) !important;
    backdrop-filter: blur(20px);
    transition: all 0.3s ease;
}
.input-pill:focus-within {
    border-color: rgba(99,102,241,0.4) !important;
    background: rgba(255,255,255,0.06) !important;
    box-shadow: 0 0 30px rgba(99,102,241,0.06);
}
.input-pill textarea, .input-pill input {
    background: transparent !important;
    border: none !important;
    color: #e2e8f0 !important;
}

/* Settings controls */
.dark-slider input[type="range"] {
    accent-color: #6366f1;
}
.clear-btn button {
    border-radius: 12px !important;
    transition: all 0.2s ease;
}
.clear-btn button:hover {
    background: rgba(255,255,255,0.1) !important;
    color: #fff !important;
}

/* Code blocks in chat */
.dark-chatbot pre {
    background: rgba(0,0,0,0.4) !important;
    border: 1px solid rgba(255,255,255,0.06) !important;
    border-radius: 10px !important;
    color: #e2e8f0 !important;
}
.dark-chatbot code {
    color: #a5b4fc !important;
}
"""


demo = create_app()

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS)