import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

# ==============================================================================
# 1. CẤU HÌNH MODEL (ĐÃ CẬP NHẬT CHO EVONET 3B BETA)
# ==============================================================================
REPO_ID = "EvoNet/EvoNet-3B-v0.1-Beta-GGUF"
FILENAME = "EvoNet-3B-v0.1-Beta-Q4_K_M.gguf" # File nén Q4 nhẹ & mượt

print(f"⏳ Đang tải file {FILENAME} từ kho {REPO_ID}...")
try:
    model_path = hf_hub_download(
        repo_id=REPO_ID, 
        filename=FILENAME
    )
    print("✅ Tải xong! Đường dẫn:", model_path)
except Exception as e:
    print(f"❌ Lỗi tải file: {e}")
    print("👉 Founder kiểm tra lại tên file hoặc Repo ID nhé!")
    raise e

# ==============================================================================
# 2. KHỞI ĐỘNG ĐỘNG CƠ AI
# ==============================================================================
print("⚙️ Đang khởi động EvoNet 3B Engine...")
llm = Llama(
    model_path=model_path,
    n_ctx=4096,        # Tăng lên 4096 tokens (nhớ dai hơn bản cũ)
    n_threads=2,       # Dùng 2 nhân CPU (Chuẩn Free Tier HF)
    n_gpu_layers=0,    # Chạy hoàn toàn trên CPU
    verbose=False      
)

# ==============================================================================
# 3. HÀM XỬ LÝ CHAT (STREAMING)
# ==============================================================================
def chat_stream(message, history):
    # Cài đặt danh tính (System Prompt) chuẩn Founder
    system_prompt = (
        "Bạn là EvoNet, một trợ lý AI thông minh được phát triển bởi Founder Huỳnh Dương Phong. "
        "Hãy trả lời ngắn gọn, hữu ích, trung thực và ưu tiên sử dụng Tiếng Việt."
    )
    
    # Format chuẩn ChatML (Qwen 2.5 yêu cầu cái này)
    prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
    for user_msg, bot_msg in history:
        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n{bot_msg}<|im_end|>\n"
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    # Gọi Model chạy
    stream = llm(
        prompt,
        max_tokens=1024,     # Cho phép trả lời dài hơn (Model 3B viết code dài được)
        stop=["<|im_end|>"], # Dấu hiệu dừng
        echo=False,          
        stream=True,         # Hiệu ứng gõ chữ
        temperature=0.7,     # Độ sáng tạo vừa phải
        top_p=0.9
    )
    
    partial_message = ""
    for output in stream:
        token = output['choices'][0]['text']
        partial_message += token
        yield partial_message

# ==============================================================================
# 4. GIAO DIỆN CHATBOT (GRADIO)
# ==============================================================================
custom_css = """
#component-0 {max-width: 800px; margin: auto;}
.gradio-container {background-color: #f9f9f9;}
"""

demo = gr.ChatInterface(
    chat_stream,
    title="🔥 EvoNet-3B Beta (GGUF Version)",
    description=f"""
    <div style="text-align: center;">
        <h3>Sản phẩm AI của Founder Huỳnh Dương Phong</h3>
        <p>Model: <b>EvoNet-3B-v0.1-Beta</b> | Format: <b>GGUF Q4_K_M</b></p>
        <p>⚡ Chạy trực tiếp trên CPU Server (Zero-GPU)</p>
    </div>
    """,
    examples=[
        ["EvoNet là gì và ai tạo ra bạn?"],
        ["Viết một hàm Python kiểm tra số nguyên tố và giải thích."],
        ["Giải thích 'Machine Learning' cho học sinh lớp 5."],
        ["Làm một bài thơ về lập trình viên."],
    ],
    cache_examples=False,
    css=custom_css,
    theme="soft" # Giao diện mềm mại, hiện đại
)

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