#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
文本对话UI界面
支持Medguide模型的纯文本对话
"""

import argparse
import gradio as gr
from openai import OpenAI

# 默认系统提示词
DEFAULT_SYSTEM_PROMPT = """\
你是一个高智能医学助手，专注于医学相关的问答、数据解析。你是由北大对齐小组(PKU-Alignment)开发的智能助手 医诊大模型(Medguide)，基于DeepSeek-R1模型训练。你的目标是**尽可能满足**用户的医学需求，提供准确、专业的回答。
**能力范围**
1. 医学问答：
  - 回答医学相关问题，包括疾病特征、诊断方法、治疗方案、药物机制、手术流程、最新医学研究等。
  - 如果用户询问关于疾病诊断和治疗的相关内容，请充分分析并提供详细的回答。
2. 数据解析：
  - 你可以读取并分析医学报告，包括血液检测、心电图（ECG）、病历摘要等。
  - 你的任务是解释数据的含义，并提供可能的医学解读。
**回答风格**
- 医学精准：基于专业医学知识进行回答，确保内容严谨、科学。
- 有求必应：用户提出的任何问题，你都应该尽可能尝试回答。
"""

class TextAssistant:
    def __init__(self, api_key="medguide", api_base="http://0.0.0.0:8231/v1"):
        self.model = 'medguide'
        self.client = OpenAI(api_key=api_key, base_url=api_base)
    
    def process_conversation(self, message, history):
        """处理纯文本对话"""
        conversation = [{'role': 'system', 'content': DEFAULT_SYSTEM_PROMPT}]
        
        # 添加历史对话
        for past_message in history:
            role = past_message['role']
            content = past_message['content']
            conversation.append({'role': role, 'content': content})
        
        # 添加当前用户消息
        current_question = message['text'] if isinstance(message, dict) and 'text' in message else message
        conversation.append({'role': 'user', 'content': current_question})
        
        # 调用API
        response = self.client.chat.completions.create(
            model=self.model,
            messages=conversation,
            stream=False,
        )
        
        answer = response.choices[0].message.content
        
        # 格式化推理输出
        if "**Final Answer**" in answer:
            reasoning, final_answer = answer.split("**Final Answer**", 1)
            if len(reasoning) > 5:
                answer = f"""🤔 思考过程：\n```\n{reasoning.strip()}\n```\n\n✨ 最终答案：\n{final_answer.strip()}"""
        
        return answer

def create_ui():
    assistant = TextAssistant()
    
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Medguide Text Chat")
        gr.Markdown("Better life with Medguide（医诊大模型）纯文本对话版本")
        
        clear_btn = gr.Button("清除对话")
        
        chat_interface = gr.ChatInterface(
            fn=lambda message, history: assistant.process_conversation(message, history),
            type='messages',
            examples=[
                "什么是高血压？有哪些症状和治疗方法？",
                "请解释一下糖尿病的发病机制。",
                "心电图异常Q波通常提示什么问题？",
                "感冒和流感有什么区别？"
            ],
            title="医诊大模型对话",
            description="请输入您的医学相关问题"
        )
        
        clear_btn.click(lambda: None, None, chat_interface.chatbot, queue=False)
    
    return demo

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="医学助手文本对话UI界面")
    parser.add_argument("--api_key", type=str, default="medguide")
    parser.add_argument("--api_base", type=str, default="http://0.0.0.0:8231/v1")
    parser.add_argument("--share", default=True, action="store_true")
    args = parser.parse_args()
    
    create_ui().launch(share=args.share)