from langchain.memory import ConversationBufferMemory
from langchain import LLMChain, PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from config import (
    template_memory,template_memory_character, template,
    template_chat_term, template_chat_character,
    template_basic, template_basic_character,
    template_reco, template_reco_character, context_file
)

#### 1. 用户切换，记忆功能
# template_memory = config['Parameters']['template_memory']
human_template_memory="{human_input}"
system_message_prompt_memory = SystemMessagePromptTemplate.from_template(template_memory, input_variables=["context"])
human_message_prompt_memory = HumanMessagePromptTemplate.from_template(human_template_memory)
prompt_memory = ChatPromptTemplate.from_messages([system_message_prompt_memory, human_message_prompt_memory])

#### 1. 用户切换，带人设的记忆功能
# template_memory_character = config['Parameters']['template_memory_character']
human_template_memory="{human_input}"
system_message_prompt_memory = SystemMessagePromptTemplate.from_template(template_memory_character, input_variables=["context"])
human_message_prompt_memory = HumanMessagePromptTemplate.from_template(human_template_memory)
prompt_memory_character = ChatPromptTemplate.from_messages([system_message_prompt_memory, human_message_prompt_memory])


#### 2. 阳光保险文档问答，不带人设，不带memory
# context_file = config['context']['context_file']
with open(context_file, 'r') as file:
    content_term = file.read()
# template_chat_term = config['Parameters']['template_chat_term']
human_template_term="{human_input}"
system_message_prompt_term = SystemMessagePromptTemplate.from_template(template_chat_term, input_variables=["context"])
human_message_prompt_term = HumanMessagePromptTemplate.from_template(human_template_term)
prompt_chat_term = ChatPromptTemplate.from_messages([system_message_prompt_term, human_message_prompt_term])

#### 2. 闲聊，带人设，带memory
# template_chat_character = config['Parameters']['template_chat_character']
human_template=" {human_input}"
system_message_prompt = SystemMessagePromptTemplate.from_template(template_chat_character)
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
prompt_chat_character = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
memory_chat = ConversationBufferMemory(memory_key="chat_history", ai_prefix="")

#### 3. 搜集个人家庭信息，不带人设，有memory
# template_basic = config['Parameters']['template_basic']
human_template="{human_input}"
system_message_prompt = SystemMessagePromptTemplate.from_template(template_basic)
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
prompt_basic = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
memory_basic = ConversationBufferMemory(memory_key="chat_history", ai_prefix="")

#### 3. 搜集个人家庭信息，带人设，有memory
# template_basic_character = config['Parameters']['template_basic_character']
human_template="{human_input}"
system_message_prompt = SystemMessagePromptTemplate.from_template(template_basic_character)
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
prompt_basic_character = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
memory_basic_character = ConversationBufferMemory(memory_key="chat_history", ai_prefix="")


def extract_name(content):
    import re
    # strings = ["我的儿子", "我儿子", "我的妻子", "我妻子", "我的父母", "我的父亲", "我父亲", "父亲", "女儿", "妻子",'我爸']
    pattern = r'儿子|女儿|小孩|父母|父亲|爸爸|老爸|爸|母亲|妈妈|老妈|妈|妻子|老婆|配偶|老公|丈夫'
    family_titles = []
    # for s in strings:
    matches = re.findall(pattern, content)
    if matches:
        return matches[0]
    else:
        return ""
def process_Info(content):
    """['人物1*称谓: 我的儿子  *年龄: 10岁  *职业: 小学生  *健康状况: 不大好，总是感冒  *生活习惯: 未知']"""
    result = []
    for item in content:
        if len(item.strip()) != 0:
            data = {}
            tmp = item.split('*')
            if len(tmp) >= 1:
                for it in tmp[1:]:  # tmp = ['人物1', '称谓: 先生  ', '年龄: 42  ', '职业: 技术高管  ', '健康状况: 三高，腰椎不太好 ', '生活习惯: 每周会去健身房锻炼一到两次，但饮食不太规律，经常加班，每天的睡觉时间也很少  ', '风险提示: 加强身体管理，科学饮食及减少加班，增加睡眠时间。']
                    if len(it.split(':')) == 2:
                        if it.split(':')[0].strip() == "称谓":
                            data.update({"name"  : extract_name(it.split(':')[1].strip()) })
                        elif it.split(':')[0].strip() == "年龄":
                            data.update({"age"  : it.split(':')[1].strip()})
                        elif it.split(':')[0].strip() == "职业":
                            data.update({"career"  : it.split(':')[1].strip()})
                        elif it.split(':')[0].strip() == "健康状况":
                            data.update({"health"  : it.split(':')[1].strip()})
                        elif it.split(':')[0].strip() == "生活习惯":
                            data.update({"live"  : it.split(':')[1].strip()})
            result.append(data)
    return result

def get_detailInfo(text):
    # template = config['Parameters']['template']
    prompt = PromptTemplate(
        input_variables=["context"], 
        template=template
    )
    llm = ChatOpenAI(request_timeout = 8*60, model_name="gpt-4")
    llm_chain = LLMChain(llm=llm, prompt=prompt, verbose=True)
    res = llm_chain.predict(context=text)
    return res



#### 4. 保险推荐，不带人设，有memory
# template_reco = config['Parameters']['template_reco']
prompt_reco = PromptTemplate(
    input_variables=["chat_history", "human_input", "context", "product"], 
    template=template_reco
)
memory_reco = ConversationBufferMemory(memory_key="chat_history", input_key = "human_input")

#### 4. 保险推荐，带人设，有memory
# template_reco_character = config['Parameters']['template_reco_character']
prompt_reco_character = PromptTemplate(
    input_variables=["chat_history", "human_input", "context", "product"], 
    template = template_reco_character
)
memory_reco_character = ConversationBufferMemory(memory_key="chat_history", input_key = "human_input")

product = { \
    "医疗保险和定期寿险":
    [
        {"k1": "产品一名称", "k2": "神农父母防癌医疗险", "k3": "亮点", "k4": "对于恶性肿瘤（含原位癌）的保障比较全面，保额较高", "k5": "适合人群", "k6": "出生满30日-70周岁", "k7": "保障额度", "k8": "最高保额300万元"},
        {"k1": "产品一名称", "k2": "擎天柱6号定期寿险", "k3": "亮点", "k4": "0等待期，提供额外赔付的猝死保障以及可选的交通意外身故保障，性价比高", "k5": "适合人群", "k6": "18-60岁", "k7": "保障额度", "k8": "最高保额350万元"}
    ],
    "医疗保险和重疾险":
    [
        {"k1": "产品一名称", "k2": "神农父母防癌医疗险", "k3": "亮点", "k4": "对于恶性肿瘤（含原位癌）的保障比较全面，保额较高", "k5": "适合人群", "k6": "出生满30日-70周岁", "k7": "保障额度", "k8": "最高保额300万元"},
        {"k1": "产品一名称", "k2": "达尔文易核版重疾险2021", "k3": "亮点", "k4": "缴费期限宽松，对7种常见疾病患者更为宽松", "k5": "适合人群", "k6": "出生满30日-50周岁", "k7": "保障额度", "k8": "最高保额50万元"}
    ],
    "重疾险和定期寿险":
    [
        {"k1": "产品一名称", "k2": "达尔文易核版重疾险2021", "k3": "亮点", "k4": "缴费期限宽松，对7种常见疾病患者更为宽松", "k5": "适合人群", "k6": "出生满30日-50周岁", "k7": "保障额度", "k8": "最高保额50万元"},
        {"k1": "产品一名称", "k2": "擎天柱6号定期寿险", "k3": "亮点", "k4": "0等待期，提供额外赔付的猝死保障以及可选的交通意外身故保障，性价比高", "k5": "适合人群", "k6": "18-60岁", "k7": "保障额度", "k8": "最高保额350万元"}
    ]
}