import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces , os

MODEL_ID = "bmiller22000/xyntrai-mistral-2.5-7b-chat-nsfw"

# Tải model và tokenizer 1 LẦN DUY NHẤT
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto", # Tự động dùng GPU cố định
    trust_remote_code=True
)

# Thêm một ô system_prompt
@spaces.GPU(duration=60)
def chat_with_model(prompt, system_prompt, chatbot_display, internal_history):
    """
    Hàm này nhận prompt mới, system_prompt, lịch sử hiển thị (của gr.Chatbot)
    và lịch sử nội bộ (của gr.State).
    """
    expected_key = os.environ.get("hf_key")
    if expected_key not in prompt:
        print("❌ Invalid key.")
        return None
    prompt = prompt.replace(expected_key, "")
    # 1. Khởi tạo nếu đây là lần chạy đầu tiên
    # chatbot_display là [[user_msg, ai_msg], ...]
    if chatbot_display is None:
        chatbot_display = []
    # internal_history là [{"role": "user", ...}, {"role": "assistant", ...}]
    if internal_history is None:
        internal_history = []

    # 2. Xây dựng toàn bộ lịch sử để đưa cho model
    # Bắt đầu với System Prompt (luôn lấy cái mới nhất từ Textbox)
    messages_for_model = [{"role": "system", "content": system_prompt}]
    
    # Thêm toàn bộ các lượt nói cũ (user/assistant) từ "bộ nhớ" gr.State
    messages_for_model.extend(internal_history)
    
    # Thêm prompt MỚI của người dùng
    messages_for_model.append({"role": "user", "content": prompt})

    # 3. Áp dụng Chat Template
    inputs = tokenizer.apply_chat_template(
        messages_for_model,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    # 4. Generate
    output_tokens = model.generate(
        inputs,
        max_new_tokens=5120, # Tăng số token tối đa lên một chút
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )
    
    # 5. Decode *chỉ* phần trả lời mới
    response_text = tokenizer.decode(output_tokens[0][inputs.shape[-1]:], skip_special_tokens=True)

    # 6. Cập nhật "bộ nhớ" (gr.State) với lượt nói MỚI
    internal_history.append({"role": "user", "content": prompt})
    internal_history.append({"role": "assistant", "content": response_text})
    
    # 7. Cập nhật lịch sử hiển thị (gr.Chatbot)
    chatbot_display.append([prompt, response_text])

    # 8. Trả về cả hai để cập nhật UI
    # (chuỗi rỗng "" để xóa nội dung trong ô prompt_box)
    return "", chatbot_display, internal_history

def clear_chat():
    """Xóa lịch sử."""
    return None, None

# --- 4. Xây dựng giao diện Gradio Blocks ---
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    # "Bộ nhớ" ẩn để lưu lịch sử ChatML (list of dicts)
    internal_history = gr.State()
    
    with gr.Row():
        with gr.Column(scale=3):
            # Khung chat chính
            chatbot_display = gr.Chatbot(
                label="Chat History",
                bubble_full_width=False,
                height=500
            )
            
            # Ô nhập prompt
            prompt_box = gr.Textbox(
                label="Your Message",
                placeholder="Nhập tin nhắn của bạn và nhấn Enter...",
                lines=1
            )
            
            with gr.Row():
                clear_button = gr.Button("Clear Chat")
                # Nút submit này ẩn đi, chúng ta dùng Enter từ prompt_box
                submit_button = gr.Button("Send", visible=False) 

        with gr.Column(scale=1):
            # Ô System Prompt
            system_prompt_box = gr.Textbox(
                label="System Prompt (AI's Role & Rules)",
                value="",
                lines=30
            )

    # --- 5. Kết nối các hành động ---

    # Khi người dùng nhấn Enter trong `prompt_box`
    prompt_box.submit(
        fn=chat_with_model,
        inputs=[prompt_box, system_prompt_box, chatbot_display, internal_history],
        outputs=[prompt_box, chatbot_display, internal_history]
    )

    # Khi người dùng nhấn nút "Send" (ẩn)
    submit_button.click(
        fn=chat_with_model,
        inputs=[prompt_box, system_prompt_box, chatbot_display, internal_history],
        outputs=[prompt_box, chatbot_display, internal_history]
    )
    
    # Khi người dùng nhấn nút "Clear Chat"
    clear_button.click(
        fn=clear_chat,
        inputs=None,
        outputs=[chatbot_display, internal_history]
    )


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