from openai import OpenAI
import json
import re
import os

# 设置OpenAI API密钥和基础URL
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")
model_name = os.getenv("OPENAI_MODEL_NAME", "gpt-3.5-turbo")

def extract_boolean(text):
    """从文本中提取布尔值判断"""
    # 查找明确的"是"或"否"的回答
    text_lower = text.lower()
    
    # 更具体地查找否定表达 - 这些应该优先匹配
    negative_patterns = [
        r'不需要', r'没有提及', r'不涉及', r'没有涉及', r'无关', r'没有提到',
        r'不是', r'否', r'不包含', r'未提及', r'未涉及', r'未提到',
        r'不包括', r'并未', r'不包括', r'没有', r'无'
    ]
    
    # 检查是否有明确的否定
    for pattern in negative_patterns:
        if re.search(r'\b' + pattern + r'\b', text_lower):
            return False
    
    # 如果找到"之前疗程"附近有否定词，也认为是否定
    therapy_negation = re.search(r'(没有|不|未|无).*?(之前|以前|上次|过去|先前).*?(疗程|治疗|会话)', text_lower)
    if therapy_negation:
        return False
        
    # 明确的肯定模式 - 只有在没有否定的情况下才考虑
    positive_patterns = [
        r'是的', r'提及了', r'确实', r'有提到', r'涉及到', 
        r'提及', r'确认', r'有关联', r'有联系', r'包含', r'涉及'
    ]
    
    # 检查是否有肯定模式
    for pattern in positive_patterns:
        if re.search(r'\b' + pattern + r'\b', text_lower):
            return True
            
    # 查找含有"之前疗程"的文本，没有否定词的情况下可能是肯定
    therapy_mention = re.search(r'(之前|以前|上次|过去|先前).*?(疗程|治疗|会话)', text_lower)
    if therapy_mention:
        return True
    
    # 默认情况 - 如果没有明确的肯定或否定，我们假设是否定的
    return False

def extract_knowledge(text):
    """从文本中提取知识总结部分"""
    # 尝试匹配总结部分
    summary_patterns = [
        r'总结[:：]\s*([\s\S]+)$',
        r'知识总结[:：]\s*([\s\S]+)$',
        r'相关信息[:：]\s*([\s\S]+)$',
        r'搜索结果[:：]\s*([\s\S]+)$'
    ]
    
    for pattern in summary_patterns:
        match = re.search(pattern, text)
        if match:
            return match.group(1).strip()
    
    # 如果没有找到明确的总结标记，尝试清理文本
    # 移除可能的指令解释部分
    clean_text = re.sub(r'^.*?(根据|基于).*?[，,。]', '', text, flags=re.DOTALL)
    
    # 移除可能的前导分析部分
    clean_text = re.sub(r'^.*?(分析|查看|判断).*?\n\n', '', clean_text, flags=re.DOTALL)
    
    return clean_text.strip()

def is_need(utterance):
    client = OpenAI(
        api_key=api_key,
        base_url=base_url
    )
    
    instruction = """
    ### 任务
    下面这句话是心理咨询师说的话，请判断它是否提及了之前疗程的内容。
    
    请使用以下确切格式回答:
    判断: [是/否]
    解释: [简要解释为什么]
    
    ### 话语
    "{}"
    """.format(utterance)
    
    response = client.chat.completions.create(
        model=model_name,
        messages=[{"role": "user", "content": instruction}],
        temperature=0
    )
    
    response_text = response.choices[0].message.content
    
    # 首先尝试从格式化输出中提取
    judgment_match = re.search(r'判断:\s*(是|否)', response_text)
    if judgment_match:
        return judgment_match.group(1) == "是"
    
    # 如果没有格式化输出，使用更通用的提取
    return extract_boolean(response_text)

def query(utterance, conversations, scales):
    client = OpenAI(
        api_key=api_key,
        base_url=base_url
    )
    
    # 将scales转换为字符串以便传入
    if isinstance(scales, dict):
        scales_str = json.dumps(scales, ensure_ascii=False)
    else:
        scales_str = str(scales)
    
    instruction = """
    ### 任务
    根据对话内容，从知识库中搜索相关的信息并总结。
    
    请使用以下确切格式回答:
    总结: [提供一个清晰、简洁的总结]
    
    ### 对话内容
    {}
    
    ### 知识库
    对话历史: {}
    量表结果: {}
    """.format(utterance, conversations, scales_str)
    
    response = client.chat.completions.create(
        model=model_name,
        messages=[{"role": "user", "content": instruction}],
        temperature=0
    )
    
    response_text = response.choices[0].message.content
    
    # 尝试提取总结部分
    summary_match = re.search(r'总结:\s*([\s\S]+)$', response_text)
    if summary_match:
        return summary_match.group(1).strip()
    
    # 回退到通用提取
    return extract_knowledge(response_text)

# 测试用例
# if __name__ == "__main__":
#     test_utterance = "上次给你说的方法有用吗"
#     # test_utterance = "我觉得你可以多出去走走"
#     print(f"是否提及疗程: {is_need(test_utterance)}")
    
#     test_convs = ["第一次对话内容", "讨论量表结果", "提到睡眠问题"]
#     test_scales = {"BDI": ["A", "B"], "GHQ": ["C", "D"]}
#     print(f"知识检索结果:\n{query(test_utterance, test_convs, test_scales)}")