#!/usr/bin/env python3
# qwen_style_gui.py ── 相容 Gradio 4.x+，Qwen 風格，重整清空

import gradio as gr
import requests, json, re, time
from threading import Thread

URL = "http://localhost:8000/v1/chat/completions"
HEADERS = {"Content-Type": "application/json"}

def sanitize(text: str):
    return re.sub(r"[\ud800-\udfff]", "", text)

# ======================
# 全域狀態（每次重啟或重整頁面會重設）
# ======================
history_openai = []  # OpenAI 格式 [{"role": "user", "content": "..."}, ...]
is_streaming = False
assistant_buffer = ""

# ======================
# Streaming 主執行緒
# ======================
def stream_to_buffer(user_input: str):
    global assistant_buffer, is_streaming, history_openai
    user_input = sanitize(user_input)
    history_openai.append({"role": "user", "content": user_input})
    payload = {"messages": history_openai, "stream": True, "temperature": 0.7}
    assistant_buffer = ""
    is_streaming = True

    try:
        with requests.post(URL, headers=HEADERS, json=payload, stream=True, timeout=60) as r:
            r.raise_for_status()
            byte_buf = b""
            for chunk in r.iter_content(chunk_size=1024):
                if not chunk:
                    break
                byte_buf += chunk
                while b"\n" in byte_buf:
                    line_bytes, byte_buf = byte_buf.split(b"\n", 1)
                    line = line_bytes.decode("utf-8", errors="replace").strip()
                    if not line.startswith("data: "):
                        continue
                    data = line[6:]
                    if data == "[DONE]":
                        break
                    try:
                        tok = json.loads(data)["choices"][0]["delta"].get("content")
                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue
                    if tok:
                        assistant_buffer += tok
            # 結束後更新歷史
            final_content = sanitize(assistant_buffer)
            history_openai.append({"role": "assistant", "content": final_content})
    except Exception as e:
        assistant_buffer = f"请求失败: {e}"
    finally:
        is_streaming = False

# ======================
# 使用者送出訊息
# ======================
def user_submit(user_msg: str):
    global assistant_buffer, is_streaming
    if not user_msg.strip():
        return [], "", ""
    # 清空 buffer，啟動 streaming
    assistant_buffer = ""
    is_streaming = True
    Thread(target=stream_to_buffer, args=(user_msg,), daemon=True).start()
    # 回傳初始對話（user + 空 assistant）
    return [{"role": "user", "content": user_msg}], "", user_msg

# ======================
# Timer 定時刷新
# ======================
def flush_buffer(current_messages):
    global assistant_buffer, is_streaming
    if not current_messages:
        return current_messages

    # 如果最後一條是 user，就加上 streaming 中的 assistant 回覆
    if current_messages[-1]["role"] == "user":
        return current_messages + [{"role": "assistant", "content": assistant_buffer}]
    elif current_messages[-1]["role"] == "assistant":
        # 更新最後一條 assistant 的內容
        updated = current_messages.copy()
        updated[-1]["content"] = assistant_buffer
        return updated
    return current_messages

# ======================
# 清除對話
# ======================
def clear_conversation():
    global history_openai, assistant_buffer, is_streaming
    history_openai = []
    assistant_buffer = ""
    is_streaming = False
    return [], "", ""

# ======================
# Qwen 風格 CSS
# ======================
QWEN_CSS = """
#chatbot {
    height: 600px;
    overflow-y: auto;
}
#chatbot .user {
    background-color: #e6f0ff;
    border-radius: 12px;
    padding: 12px;
    margin: 8px 0;
    text-align: right;
}
#chatbot .bot {
    background-color: #f0f0f0;
    border-radius: 12px;
    padding: 12px;
    margin: 8px 0;
    text-align: left;
}
.gradio-container {
    font-family: 'Segoe UI', 'PingFang SC', 'Microsoft YaHei', sans-serif;
}
"""

# ======================
# Gradio 介面（使用 type='messages'）
# ======================
with gr.Blocks(css=QWEN_CSS, title="Qwen Chat") as demo:
    gr.Markdown("## 💬 Qwen Chat Service")
    chatbot = gr.Chatbot(elem_id="chatbot", type="messages")  # ✅ 關鍵：type="messages"

    with gr.Row():
        msg = gr.Textbox(
            show_label=False,
            placeholder="輸入訊息後按 Enter",
            container=False,
            scale=8
        )
        clear_btn = gr.Button("🗑️ 清除", scale=1)

    # 提交時：更新 chatbot 為 [{"role": "user", "content": "..."}]
    msg.submit(
        fn=user_submit,
        inputs=msg,
        outputs=[chatbot, msg, gr.State()],  # 第三個 output 是為了觸發 timer（可選）
        queue=False
    )

    clear_btn.click(clear_conversation, outputs=[chatbot, msg, gr.State()])

    # Timer 每 0.3 秒刷新
    timer = gr.Timer(value=0.3)
    timer.tick(flush_buffer, inputs=chatbot, outputs=chatbot)

if __name__ == "__main__":
    # ⚠️ 不要加 reload=...，這不是 launch() 的參數
    demo.launch(server_name="0.0.0.0", server_port=7860)