import gradio as gr
import requests
import json
import uuid
import logging
import os

# 配置日志
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

class FastGPTChat:
    def __init__(self, system_prompt="我是一名AI用药咨询顾问，请向我提问有关用药的问题"):
        self.api_key = os.environ.get('FASTGPT_API_KEY')
        self.base_url = os.environ.get('FASTGPT_BASE_URL', 'https://api.fastgpt.in/api')
        self.system_prompt = system_prompt
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
    
    def chat(self, message, chat_history):
        if not message.strip():
            return "", chat_history
        
        # 立即添加用户消息到聊天历史
        chat_history.append((message, ""))
        yield "", chat_history
            
        chat_id = str(uuid.uuid4())
        
        # 构建消息历史，包含系统提示词
        messages = []
        if self.system_prompt:
            messages.append({
                "role": "system",
                "content": self.system_prompt
            })
        
        # 添加聊天历史
        for user_msg, bot_msg in chat_history[:-1]:  # 不包含最新消息
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": bot_msg})
        
        # 添加当前用户消息
        messages.append({"role": "user", "content": message})
        
        # 准备请求数据
        data = {
            "chatId": chat_id,
            "stream": True,
            "detail": True,
            "model": "gpt-4o",
            "messages": messages
        }
        
        try:
            # 发送流式请求
            response = requests.post(
                f"{self.base_url}/v1/chat/completions",
                headers=self.headers,
                json=data,
                stream=True,
                timeout=30
            )
            
            if response.status_code != 200:
                error_msg = f"API Error: Status {response.status_code}"
                logger.error(error_msg)
                chat_history[-1] = (message, f"Error: {error_msg}")
                yield "", chat_history
                return
            
            # 处理流式响应
            full_response = ""
            for line in response.iter_lines():
                if line:
                    try:
                        line = line.decode('utf-8')
                        if line.startswith('data: '):
                            data = json.loads(line[6:])
                            if 'choices' in data and len(data['choices']) > 0:
                                content = data['choices'][0]['delta'].get('content', '')
                                if content:
                                    full_response += content
                                    chat_history[-1] = (message, full_response)
                                    yield "", chat_history
                    except Exception as e:
                        logger.error(f"Error processing stream: {str(e)}")
            
            yield "", chat_history
                
        except requests.exceptions.RequestException as e:
            error_msg = f"Request failed: {str(e)}"
            logger.error(error_msg)
            chat_history[-1] = (message, f"Error: {error_msg}")
            yield "", chat_history

def create_chat_interface():
    fastgpt_chat = FastGPTChat()
    
    with gr.Blocks(
        title="AI用药咨询助手",
        css="""
        /* 整体容器样式 */
        .container { 
            width: 80%;
            margin: auto;
            padding: 2rem;
            box-sizing: border-box;
        }

        /* 标题区域样式 */
        .title-container {
            text-align: center;
            margin-bottom: 2rem;
        }

        .title-container h1 {
            color: #2C3E50;
            font-size: 2.5rem;
            margin-bottom: 0.5rem;
        }

        .title-container h3 {
            color: #7F8C8D;
            font-size: 1.2rem;
            font-weight: normal;
        }

        /* 聊天容器样式 */
        .chat-container { 
            background-color: #FFFFFF;
            border-radius: 20px;
            box-shadow: 0 8px 24px rgba(0, 0, 0, 0.1);
            padding: 2rem;
            margin-bottom: 2rem;
        }

        /* 聊天框样式 */
        .chatbot {
            background-color: #F8F9FA;
            border-radius: 15px;
            padding: 1.5rem;
            margin-bottom: 1.5rem;
            height: 600px !important;
            overflow-y: auto;
        }

        /* 消息样式 */
        .chatbot .user-message {
            background: linear-gradient(135deg, #6366F1, #4F46E5);
            color: white !important;
            border-radius: 18px 18px 4px 18px !important;
            padding: 1rem 1.5rem !important;
            margin: 1rem 0 !important;
            max-width: 80% !important;
            float: right !important;
            clear: both !important;
            box-shadow: 0 4px 12px rgba(99, 102, 241, 0.2) !important;
        }

        .chatbot .bot-message {
            background: white !important;
            border: 1px solid #E5E7EB !important;
            color: #1F2937 !important;
            border-radius: 18px 18px 18px 4px !important;
            padding: 1rem 1.5rem !important;
            margin: 1rem 0 !important;
            max-width: 80% !important;
            float: left !important;
            clear: both !important;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.05) !important;
        }

        /* 输入区域样式 */
        .input-container {
            background: white;
            padding: 1.5rem;
            border-radius: 15px;
            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.05);
        }

        .message-box {
            border: 2px solid #E5E7EB !important;
            border-radius: 12px !important;
            padding: 1rem !important;
            font-size: 1rem !important;
            transition: all 0.3s ease !important;
            background: white !important;
        }

        .message-box:focus {
            border-color: #6366F1 !important;
            box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.2) !important;
        }

        /* 按钮样式 */
        .button-container {
            display: flex;
            gap: 1rem;
            margin-top: 1rem;
        }

        .submit-btn {
            background: linear-gradient(135deg, #6366F1, #4F46E5) !important;
            color: white !important;
            padding: 0.75rem 1.5rem !important;
            border-radius: 12px !important;
            font-weight: 600 !important;
            transition: all 0.3s ease !important;
            flex: 1;
        }

        .submit-btn:hover {
            transform: translateY(-2px);
            box-shadow: 0 4px 12px rgba(99, 102, 241, 0.3) !important;
        }

        .clear-btn {
            background: #EF4444 !important;
            color: white !important;
            padding: 0.75rem 1.5rem !important;
            border-radius: 12px !important;
            font-weight: 600 !important;
            transition: all 0.3s ease !important;
        }

        .clear-btn:hover {
            transform: translateY(-2px);
            box-shadow: 0 4px 12px rgba(239, 68, 68, 0.3) !important;
        }
        """
    ) as interface:
        with gr.Column(elem_classes="container"):
            with gr.Column(elem_classes="title-container"):
                gr.Markdown(
                    """
                    # 🏥 AI用药咨询助手
                    ### 您的专业用药咨询顾问，随时为您解答用药相关问题
                    """
                )
            
            with gr.Column(elem_classes="chat-container"):
                chatbot = gr.Chatbot(
                    height=600,
                    container=False,
                    elem_classes="chatbot",
                    show_label=False,
                    bubble_full_width=False,
                    avatar_images=("👤", "🤖")
                )
                
                with gr.Column(elem_classes="input-container"):
                    message = gr.Textbox(
                        show_label=False,
                        placeholder="请输入您的用药问题...",
                        container=False,
                        elem_classes="message-box",
                        lines=3
                    )
                    
                    with gr.Row(elem_classes="button-container"):
                        submit = gr.Button(
                            "发送 💊", 
                            variant="primary",
                            elem_classes="submit-btn"
                        )
                        clear = gr.Button(
                            "清空对话 🗑️", 
                            variant="secondary",
                            elem_classes="clear-btn"
                        )
        
        # 设置事件处理
        submit_click = submit.click(
            fastgpt_chat.chat,
            inputs=[message, chatbot],
            outputs=[message, chatbot],
            api_name="chat"
        )
        
        message.submit(
            fastgpt_chat.chat,
            inputs=[message, chatbot],
            outputs=[message, chatbot]
        )
        
        clear.click(lambda: None, None, chatbot, queue=False)
        
    return interface

if __name__ == "__main__":
    demo = create_chat_interface()
    demo.launch(debug=True)