#!/usr/bin/env python3
import json, re, requests, gradio as gr

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

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

history = []
current_response = ""

def stream_chat(user_input: str):
    global current_response
    
    user_input = sanitize(user_input)
    history.append({"role": "user", "content": user_input})
    payload = {"messages": history, "stream": True, "temperature": 0.7}
    
    current_response = ""
    full_conversation = history.copy()
    
    try:
        response = requests.post(URL, headers=HEADERS, json=payload,
                                stream=True, timeout=60)
        response.raise_for_status()
        
        for line in response.iter_lines(decode_unicode=True):
            if line and line.startswith("data: "):
                data = line[6:]
                if data == "[DONE]":
                    break
                
                try:
                    json_data = json.loads(data)
                    if "choices" in json_data and len(json_data["choices"]) > 0:
                        delta = json_data["choices"][0].get("delta", {})
                        token = delta.get("content", "")
                        
                        if token:
                            current_response += token
                            # 构建完整的对话显示
                            display_text = ""
                            for msg in full_conversation:
                                role = "用户" if msg["role"] == "user" else "助手"
                                display_text += f"**{role}**: {msg['content']}\n\n"
                            display_text += f"**助手**: {current_response} ▌"
                            yield display_text
                            
                except json.JSONDecodeError:
                    continue
        
        # 完成响应
        history.append({"role": "assistant", "content": sanitize(current_response)})
        display_text = ""
        for msg in history:
            role = "用户" if msg["role"] == "user" else "助手"
            display_text += f"**{role}**: {msg['content']}\n\n"
        yield display_text
        
    except Exception as e:
        error_msg = f"请求失败: {e}"
        history.append({"role": "assistant", "content": error_msg})
        display_text = ""
        for msg in history:
            role = "用户" if msg["role"] == "user" else "助手"
            display_text += f"**{role}**: {msg['content']}\n\n"
        yield display_text

def clear_chat():
    global history, current_response
    history = []
    current_response = ""
    return "对话已清空"

# 创建 Gradio 界面
with gr.Blocks(title="Chat8000 - 流式对话") as demo:
    gr.Markdown("# Chat8000 - 流式对话版")
    
    # 使用 Markdown 显示对话
    chat_display = gr.Markdown(
        label="对话记录",
        value="欢迎使用 Chat8000！请输入您的问题..."
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="输入消息",
            placeholder="请输入您的问题...",
            lines=2,
            scale=4
        )
        submit_btn = gr.Button("发送", variant="primary", scale=1)
    
    with gr.Row():
        clear_btn = gr.Button("清空对话")
    
    # 事件处理
    def handle_submit(message):
        if not message.strip():
            yield chat_display.value
            return
            
        for display_text in stream_chat(message):
            yield display_text
    
    # 绑定事件
    msg.submit(
        handle_submit,
        inputs=[msg],
        outputs=[chat_display]
    )
    
    submit_btn.click(
        handle_submit,
        inputs=[msg],
        outputs=[chat_display]
    )
    
    clear_btn.click(
        clear_chat,
        inputs=[],
        outputs=[chat_display]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)