import gradio as gr
from huggingface_hub import InferenceClient

def respond(
    message,
    history, 
    hf_token: gr.OAuthToken | None = None,
):
    # 1. Khởi tạo Client với Token an toàn
    token = hf_token.token if hf_token else None
    # Lưu ý: Nếu model này vẫn báo lỗi aborted, bạn có thể thử thay bằng "HuggingFaceH4/zephyr-7b-beta"
    client = InferenceClient(token=token, model="openai/gpt-oss-20b")

    # 2. Xây dựng Bộ nhớ (Memory) sạch
    messages = [{"role": "system", "content": "bạn tên là Neso , bạn có thể tạo code và là một AI viết code toàn cục siêu giỏi lập chình."}]
    
    for entry in history:
        if isinstance(entry, dict):
            # Chỉ thêm nếu nội dung không trống
            if entry.get("content"):
                messages.append(entry)
        elif isinstance(entry, (list, tuple)):
            user_m, bot_m = entry
            if user_m: messages.append({"role": "user", "content": user_m})
            if bot_m: messages.append({"role": "assistant", "content": bot_m})
    
    messages.append({"role": "user", "content": message})

    # 3. Xử lý Stream với Error Handling chuyên sâu
    response = ""
    try:
        # Cấu hình stream an toàn hơn
        stream = client.chat_completion(
            messages,
            max_tokens=1024,
            stream=True,
            temperature=0.7,
            top_p=0.9,
        )

        for msg in stream:
            if msg.choices and msg.choices[0].delta.content:
                token_text = msg.choices[0].delta.content
                response += token_text
                yield response

    except Exception as e:
        # Nếu đang stream mà bị "aborted", trả về phần đã nhận được kèm cảnh báo
        error_msg = str(e)
        if "aborted" in error_msg.lower():
            yield response + "\n\n⚠️ [Kết nối bị ngắt quãng, vui lòng gửi lại câu hỏi hoặc kiểm tra Internet]"
        else:
            yield f"⚠️ Lỗi hệ thống: {error_msg}"

# --- GIAO DIỆN PRO (FULL WIDTH & DARK MODE) ---
css = """
body { background-color: #0f172a !important; }
.gradio-container { max-width: 1000px !important; }
.chatbot { border: 1px solid #334155 !important; border-radius: 15px !important; }
footer { visibility: hidden; }
"""

theme = gr.themes.Soft(primary_hue="indigo", neutral_hue="slate").set(
    block_background_fill="#1e293b",
    input_background_fill="#0f172a",
    button_primary_background_fill="#4f46e5",
)

with gr.Blocks(title="Zola AI Pro", theme=theme, css=css) as demo:
    gr.Markdown(
        """
        <div style="text-align: center; padding: 10px;">
            <h1 style="color: white; font-size: 2.5em;">🤖 Zola AI Pro</h1>
            <p style="color: #6366f1;">Phiên bản ổn định - Tối ưu bộ nhớ</p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.LoginButton()
        with gr.Column(scale=4):
            # ChatInterface chính
            gr.ChatInterface(
                fn=respond,
                chatbot=gr.Chatbot(
                    height=600,
                    show_label=False,
                    avatar_images=(
                        "https://cdn-icons-png.flaticon.com/512/149/149071.png", 
                        "https://cdn-icons-png.flaticon.com/512/4712/4712035.png", 
                    ),
                ),
            )

if __name__ == "__main__":
    # Tắt ssr_mode nếu bạn chạy trên môi trường local/dev để ổn định hơn
    demo.launch(ssr_mode=False)