import streamlit as st
import requests
import json
import time
import os
from openai import OpenAI
import psutil
import platform

# 页面配置
st.set_page_config(
    page_title="Qwen3.5 27B Chat",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# 自定义 CSS
st.markdown("""
<style>
    .stApp {
        background-color: #f5f5f5;
    }
    .main-header {
        color: #1f77b4;
        font-size: 2.5rem;
        font-weight: 600;
        margin-bottom: 1rem;
    }
    .info-box {
        background-color: #e8f4f8;
        padding: 1rem;
        border-radius: 0.5rem;
        margin-bottom: 1rem;
    }
</style>
""", unsafe_allow_html=True)

# 初始化 OpenAI 客户端
@st.cache_resource
def init_client():
    """初始化 vLLM 客户端"""
    # 等待 vLLM 服务就绪
    max_retries = 30
    for i in range(max_retries):
        try:
            response = requests.get("http://localhost:8000/health")
            if response.status_code == 200:
                st.success("✅ 已连接到推理服务器")
                break
        except:
            if i % 5 == 0:
                st.info(f"⏳ 正在连接推理服务器... ({i+1}/{max_retries})")
            time.sleep(2)
    
    return OpenAI(
        base_url="http://localhost:8000/v1",
        api_key="EMPTY",
    )

# 侧边栏
with st.sidebar:
    st.image("https://qianwen-res.oss-cn-beijing.aliyuncs.com/assets/blog/2025/qwen3.5-logo.png", 
             use_column_width=True)
    
    st.markdown("## ⚙️ 配置参数")
    
    # 模型状态
    with st.expander("🖥️ 系统状态", expanded=False):
        col1, col2 = st.columns(2)
        with col1:
            st.metric("CPU 使用率", f"{psutil.cpu_percent()}%")
        with col2:
            mem = psutil.virtual_memory()
            st.metric("内存使用", f"{mem.percent}%")
        
        # GPU 信息（如果可用）
        try:
            import torch
            if torch.cuda.is_available():
                gpu_memory = torch.cuda.memory_allocated() / 1024**3
                gpu_total = torch.cuda.get_device_properties(0).total_memory / 1024**3
                st.metric("GPU 显存", f"{gpu_memory:.1f}GB / {gpu_total:.1f}GB")
        except:
            pass
    
    # 采样参数
    st.markdown("### 🎛️ 采样参数")
    temperature = st.slider("Temperature", 0.0, 2.0, 0.7, 0.1,
                           help="控制输出的随机性，越高越随机")
    top_p = st.slider("Top P", 0.0, 1.0, 0.8, 0.05,
                     help="核采样，控制词汇选择的累积概率")
    max_tokens = st.number_input("Max Tokens", 512, 32768, 4096, 512,
                                help="最大生成token数")
    
    # 高级选项
    with st.expander("🔧 高级选项"):
        presence_penalty = st.slider("Presence Penalty", -2.0, 2.0, 0.0, 0.1,
                                    help="话题重复惩罚")
        frequency_penalty = st.slider("Frequency Penalty", -2.0, 2.0, 0.0, 0.1,
                                     help="词频惩罚")
        enable_thinking = st.checkbox("🧠 启用思考模式", True,
                                     help="启用模型的内部思考过程")
    
    st.divider()
    st.markdown("### 📝 预设配置")
    preset = st.selectbox("快速选择", ["通用对话", "创意写作", "代码生成", "精确回答"])
    if preset == "通用对话":
        temperature, top_p = 0.7, 0.8
    elif preset == "创意写作":
        temperature, top_p = 1.2, 0.95
    elif preset == "代码生成":
        temperature, top_p = 0.2, 0.1
    elif preset == "精确回答":
        temperature, top_p = 0.3, 0.5

# 主界面
st.markdown('<p class="main-header">🤖 Qwen3.5 27B Chat</p>', 
            unsafe_allow_html=True)

# 信息框
with st.container():
    st.markdown("""
    <div class="info-box">
    📌 基于 Qwen3.5-27B-GPTQ-Int4 模型，4-bit 量化，可在 T4 16GB 显存高效运行。
    支持多轮对话、代码生成、文档理解等功能。
    </div>
    """, unsafe_allow_html=True)

# 初始化客户端
client = init_client()

# 初始化聊天历史
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "system", "content": "You are Qwen3.5, a helpful assistant."}
    ]

# 显示聊天历史
for message in st.session_state.messages[1:]:  # 跳过系统消息
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 聊天输入
if prompt := st.chat_input("输入你的问题..."):
    # 添加用户消息
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # 生成回复
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        full_response = ""
        
        try:
            # 准备消息
            messages_for_api = st.session_state.messages.copy()
            
            # 调用 API
            stream = client.chat.completions.create(
                model="Qwen/Qwen3.5-27B-GPTQ-Int4",
                messages=messages_for_api,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                presence_penalty=presence_penalty,
                frequency_penalty=frequency_penalty,
                stream=True,
                extra_body={
                    "chat_template_kwargs": {
                        "enable_thinking": enable_thinking
                    }
                } if not enable_thinking else None
            )
            
            # 处理流式响应
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    full_response += chunk.choices[0].delta.content
                    message_placeholder.markdown(full_response + "▌")
            
            message_placeholder.markdown(full_response)
            
            # 添加到历史
            st.session_state.messages.append({
                "role": "assistant",
                "content": full_response
            })
            
        except Exception as e:
            st.error(f"生成失败: {str(e)}")
            if "connection" in str(e).lower():
                st.info("正在重新连接推理服务器，请稍候...")
                st.cache_resource.clear()
                time.sleep(2)
                st.rerun()

# 底部工具栏
col1, col2, col3 = st.columns([1, 1, 4])
with col1:
    if st.button("🗑️ 清空对话"):
        st.session_state.messages = [st.session_state.messages[0]]
        st.rerun()
with col2:
    if st.button("📋 复制对话"):
        import pyperclip
        texts = [f"{m['role']}: {m['content']}" 
                for m in st.session_state.messages[1:]]
        pyperclip.copy("\n\n".join(texts))
        st.toast("已复制到剪贴板")

# 显示 token 使用统计（如果有）
if "last_response" in st.session_state:
    st.caption(f"📊 最后响应: ~{len(st.session_state.last_response)} tokens")

# 心跳机制（防止休眠）
st.markdown("---")
st.caption("💡 提示：此页面保持打开可防止Space自动休眠")