import os
import spaces
import torch
from threading import Thread
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
    BitsAndBytesConfig,
)
import gradio as gr

# ==============================================================================
# CẤU HÌNH MÔ HÌNH VÀ TOKEN
# ==============================================================================
# Bạn có thể đổi sang 'google/gemma-4-31B-it', 'google/gemma-2-9b-it' hoặc mô hình khác
# qua biến môi trường MODEL_ID trên Hugging Face Space Settings.
MODEL_ID = os.getenv("MODEL_ID", "google/gemma-4-31B-it")
HF_TOKEN = os.getenv("HF_TOKEN", None)

print(f"[*] Khởi tạo cấu hình cho mô hình: {MODEL_ID}")

# Cấu hình lượng tử hóa 4-bit (BitsAndBytes NF4) để tối ưu VRAM trên ZeroGPU
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# Tải Tokenizer và Model
print("[*] Đang tải Tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    token=HF_TOKEN,
    trust_remote_code=True,
)

print("[*] Đang tải Model với 4-bit Quantization...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    token=HF_TOKEN,
    trust_remote_code=True,
)
print("[+] Mô hình đã sẵn sàng hoạt động!")

# ==============================================================================
# HÀM SUY LUẬN & STREAMING KẾT QUẢ VỚI ZERO-GPU
# ==============================================================================
@spaces.GPU(duration=120)
def chat_response(
    message,
    history,
    system_prompt,
    temperature,
    max_new_tokens,
    top_p,
):
    if not message or not message.strip():
        yield ""
        return

    # Xây dựng danh sách tin nhắn theo chuẩn hội thoại
    conversation = []

    # 1. Thêm system prompt nếu có
    if system_prompt and system_prompt.strip():
        conversation.append({"role": "system", "content": system_prompt.strip()})

    # 2. Đọc lại lịch sử chat (tương thích cả dict messages và tuple)
    for item in history:
        if isinstance(item, dict):
            role = item.get("role", "user")
            # Gemma chat template quy ước role là 'user' và 'model'
            if role == "assistant":
                role = "model"
            conversation.append({"role": role, "content": item.get("content", "")})
        elif isinstance(item, (list, tuple)) and len(item) == 2:
            user_text, bot_text = item
            if user_text:
                conversation.append({"role": "user", "content": str(user_text)})
            if bot_text:
                conversation.append({"role": "model", "content": str(bot_text)})

    # 3. Thêm tin nhắn hiện tại của người dùng
    conversation.append({"role": "user", "content": message.strip()})

    # 4. Tạo prompt qua chat template của Gemma
    try:
        prompt = tokenizer.apply_chat_template(
            conversation,
            tokenize=False,
            add_generation_prompt=True,
        )
    except Exception:
        # Fallback nếu tokenizer không hỗ trợ role 'system' riêng rẽ
        fallback_conv = []
        for idx, item in enumerate(conversation):
            if item["role"] == "system":
                continue
            content = item["content"]
            if idx == (1 if system_prompt else 0):
                content = f"[Hướng dẫn hệ thống: {system_prompt}]\n\n{content}"
            fallback_conv.append({"role": item["role"], "content": content})
        prompt = tokenizer.apply_chat_template(
            fallback_conv,
            tokenize=False,
            add_generation_prompt=True,
        )

    # 5. Tokenize và đưa vào thiết bị của model
    model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)

    # 6. Thiết lập Streamer để bắn từng token ra ngoài màn hình
    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    generation_kwargs = dict(
        model_inputs,
        streamer=streamer,
        max_new_tokens=int(max_new_tokens),
        temperature=float(max(temperature, 0.01)),
        top_p=float(top_p),
        do_sample=True if temperature > 0.0 else False,
    )

    # Chạy model.generate trên luồng nền (background thread)
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # 7. Nhận từng token và yield về giao diện theo thời gian thực
    partial_text = ""
    for token in streamer:
        partial_text += token
        yield partial_text


# ==============================================================================
# TÙY BIẾN GIAO DIỆN THEME GOOGLE GEMINI (SÁNG - TINH TẾ - DỄ ĐỌC)
# ==============================================================================
DEFAULT_SYSTEM_PROMPT = (
    "Bạn là Gemma, một trợ lý trí tuệ nhân tạo thông minh, am hiểu và thân thiện. "
    "Hãy trả lời súc tích, chính xác, định dạng câu trả lời bằng Markdown đẹp mắt."
)

# Khởi tạo theme chuẩn Google (màu xanh dương dịu mắt, tương phản cao)
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "sans-serif"],
)

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

/* Toàn bộ nền trang sáng sủa, sạch sẽ phong cách Google */
body, .gradio-container {
    background-color: #f8fafd !important;
    color: #1f2937 !important;
    font-family: 'Google Sans', 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
    max-width: 960px !important;
    margin: 0 auto !important;
}

/* Tiêu đề phong cách Gemini Sparkle Gradient */
h1 {
    text-align: center !important;
    font-size: 2.2rem !important;
    font-weight: 700 !important;
    background: linear-gradient(135deg, #1a73e8 0%, #8b5cf6 50%, #ec4899 100%) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    letter-spacing: -0.5px !important;
    margin-bottom: 0.25rem !important;
}

/* Mô tả phụ bên dưới tiêu đề */
.description, p.description {
    text-align: center !important;
    color: #5f6368 !important;
    font-size: 0.95rem !important;
    margin-bottom: 1.25rem !important;
}

/* Khung Chatbot chính: Nền trắng, bo tròn mềm mại, đổ bóng nhẹ */
[data-testid="chatbot"], .chatbot {
    background-color: #ffffff !important;
    border: 1px solid #e2e8f0 !important;
    border-radius: 20px !important;
    box-shadow: 0 4px 20px -2px rgba(0, 0, 0, 0.05) !important;
    min-height: 520px !important;
}

/* Tin nhắn Người dùng: Nền xanh nhạt dịu mắt (Google Blue tint), chữ xanh đậm rõ nét */
[data-testid="user"], .message.user, .user-row .message {
    background-color: #e8f0fe !important;
    border: 1px solid #d2e3fc !important;
    color: #174ea6 !important;
    border-radius: 20px 20px 4px 20px !important;
    padding: 12px 18px !important;
    font-size: 15px !important;
    font-weight: 500 !important;
    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.04) !important;
}

[data-testid="user"] *, .message.user * {
    color: #174ea6 !important;
}

/* Tin nhắn Trợ lý Bot: Chữ tối rõ nét trên nền trắng, dễ đọc */
[data-testid="bot"], .message.bot, .bot-row .message {
    background-color: transparent !important;
    color: #1f2937 !important;
    border: none !important;
    font-size: 15px !important;
    line-height: 1.7 !important;
    padding: 12px 18px !important;
}

[data-testid="bot"] *, .message.bot * {
    color: #1f2937 !important;
}

/* Code block hiển thị đẹp mắt và tương phản chuẩn */
pre {
    background-color: #1e293b !important;
    border-radius: 10px !important;
    padding: 14px 16px !important;
    color: #f8fafc !important;
}

pre code, pre span {
    color: #f8fafc !important;
}

/* Ô nhập câu hỏi dạng viên thuốc (Pill shape) chuẩn Google */
input[type="text"], textarea {
    background-color: #ffffff !important;
    border: 1.5px solid #dadce0 !important;
    border-radius: 26px !important;
    color: #1f2937 !important;
    padding: 12px 20px !important;
    font-size: 15px !important;
    box-shadow: 0 2px 6px rgba(0, 0, 0, 0.04) !important;
}

input[type="text"]:focus, textarea:focus {
    border-color: #1a73e8 !important;
    box-shadow: 0 0 0 3px rgba(26, 115, 232, 0.15) !important;
}

/* Khối Accordion cài đặt: Thu gọn, nền trắng, tinh gọn */
details, .gr-accordion {
    background-color: #ffffff !important;
    border: 1px solid #e2e8f0 !important;
    border-radius: 14px !important;
    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.03) !important;
    margin-top: 0.75rem !important;
}

label, span.label-text {
    color: #374151 !important;
    font-weight: 500 !important;
}

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

# ==============================================================================
# KHỞI TẠO GIAO DIỆN CHAT INTERFACE
# ==============================================================================
demo = gr.ChatInterface(
    fn=chat_response,
    title="✨ Google Gemma Assistant",
    description="Trải nghiệm đối thoại AI phong cách Google Gemini • Tải trực tiếp mô hình Gemma 4 31B (4-bit ZeroGPU)",
    textbox=gr.Textbox(
        placeholder="Hỏi Gemma bất kỳ điều gì (nhấn Enter để gửi)...",
        lines=1,
        max_lines=6,
        container=False,
        scale=7,
    ),
    additional_inputs=[
        gr.Textbox(
            label="Chỉ dẫn hệ thống (System Prompt)",
            value=DEFAULT_SYSTEM_PROMPT,
            lines=2,
        ),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.7,
            step=0.05,
            label="Độ sáng tạo (Temperature)",
            info="Giá trị thấp trả lời chuẩn xác hơn; giá trị cao trả lời phong phú và sáng tạo hơn.",
        ),
        gr.Slider(
            minimum=128,
            maximum=4096,
            value=1536,
            step=128,
            label="Độ dài tối đa (Max New Tokens)",
            info="Số lượng token tối đa mô hình tạo ra trong mỗi lượt trả lời.",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.9,
            step=0.05,
            label="Top-P (Nucleus Sampling)",
        ),
    ],
    additional_inputs_accordion=gr.Accordion(
        label="⚙️ Cài đặt tham số mô hình & System Prompt",
        open=False,
    ),
    examples=[
        [
            "Giải thích ngắn gọn cách hoạt động của cơ chế Multi-Head Attention trong mô hình Transformer.",
            DEFAULT_SYSTEM_PROMPT,
            0.7,
            1536,
            0.9,
        ],
        [
            "Viết một kịch bản ngắn giới thiệu vẻ đẹp thiên nhiên và văn hóa cồng chiêng của vùng đất Gia Lai.",
            DEFAULT_SYSTEM_PROMPT,
            0.7,
            1536,
            0.9,
        ],
        [
            "Tạo một hàm Python đọc tệp CSV, làm sạch dữ liệu thiếu và vẽ biểu đồ xu hướng bằng Matplotlib.",
            DEFAULT_SYSTEM_PROMPT,
            0.7,
            1536,
            0.9,
        ],
        [
            "So sánh ưu điểm và nhược điểm giữa cấu trúc Dense Model và Mixture-of-Experts (MoE).",
            DEFAULT_SYSTEM_PROMPT,
            0.7,
            1536,
            0.9,
        ],
    ],
    cache_examples=False,
)

if __name__ == "__main__":
    demo.queue().launch(theme=custom_theme, css=CUSTOM_CSS)