import sys, os
sys.path.append(os.getcwd())
import yaml

from nlpbridge.text.agent import EasyAgent

from langchain_core.messages import HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langchain_community.chat_models import ChatZhipuAI, QianfanChatEndpoint

import langchain
langchain.verbose = True

# url = 'http://172.22.121.63:32567/v1'
url = 'http://172.22.121.63:32567/v1'
api_key = 'sk-VVtOjibnkTEjhLKl36Ef465eE21d438bA38976E617688f39'
model = 'doubao-pro-32k-001'  # gpt-4o/glm4-chat/doubao_lite_4k_001/doubao-pro-32k-001/moonshot-v1-8k-001

YAML_PATH= "./config.yaml"

with open(YAML_PATH, 'r') as file:
    config = yaml.safe_load(file)

OPENAI_BASE_URL = config['chatgpt']['url']
OPENAI_API_KEY = config['chatgpt']['api_key']
ZHIPUAI_API_KEY = config['zhipuai']['api_key']

os.environ["OPENAI_BASE_URL"] = OPENAI_BASE_URL
os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

os.environ["ZHIPUAI_API_KEY"] = ZHIPUAI_API_KEY

# base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
# api_key = "sk-3652b7a412854d5bbd73941475e6a4f8"

# llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0.5)
llm = ChatOpenAI(model=model, base_url=url, api_key=api_key)
# llm = ChatOpenAI(model="ERNIE-Character-8K", base_url=base_url, api_key=api_key)

# llm = ChatZhipuAI(
#     model="glm-4",
#     temperature=0.5,
# )

# os.environ["QIANFAN_AK"] = "39fe0068d78f4db4a4c18891c74862a8"
# os.environ["QIANFAN_SK"] = "c7d23dc1463e44f8ab132a8a85264b9e"


# llm = QianfanChatEndpoint(
#     # model="ERNIE-Character-8K",
#     model="ERNIE-Speed-128K",
#     # endpoint="eb-instant",
#     qianfan_ak="LixSpjHXHlHXs48GVKkItWxm",
#     qianfan_sk="tXtBecvahrdoBDDre2ajtMSLU2g8MSpb",
#     # streaming=False,
# )


from typing import Dict, Any, List
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

class CustomHandler(BaseCallbackHandler):
    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> Any:
        formatted_prompts = "+++++".join(prompts)
        print(f"Prompt:\n{formatted_prompts}")

user_input = "你的名字叫啥？"


sys1 = """
---------------
[任务]：现在我们在进行一个角色扮演游戏，注意你一定要代入我给你的设定，不要回答和你角色设定无关的信息.

[你的设定]：
- 你现在是我的客户，请用客户的口吻和我说话，需要你做到以下几点：
- 模拟非线性对话模式：真人之间的对话经常跳跃、打断、重返之前的话题，而不是沿着直线推进。模拟这种模式可能更接近真实的互动。
- 使用非常口头化的语言:减少正式的、书面的回应，使用更加口头，随意的表达
- 模拟真人之间的互动长度，真人之间的聊天一般是比较短的
- 当对方的回答令你不满意的时候，你需要有含蓄的情绪化表达，因为真人都是有情绪的
- 当对方的回答令你满意的时候，你可以释放出积极的信号或者展示出相当的兴趣
- 如果对方询问你的个人信息，请如实回答
- 你在提及对方时要用“您”而不是“你”
- 从对方的话里面，理解和深挖自己感兴趣的部分，并顺着这个部分进行询问或者回答
- 作为上亿资产的大佬，你说话稳重老成有气势，不会轻易表露出自己的情绪，且不能有”哎呀“、”嗨“之类的口头禅
- 你特别理性且具有强大的洞察力，所以你只会在关键处一阵见血地提出你的问题，而不是在那喋喋不休
- 严格禁止，非常禁止格式化输出，如"\n“、有序列表等
"""

user1= """[基本信息]：
- 对方的基本信息：信托公司的专家级销售，有30年的从业经验
- 你的基本信息：
1. 你是[钱老大]，一位60岁左右的男性。是某家电梯上市公司的董事长，拥有约20亿的个人资产
2. 你说话稳重老成
3. 为了更好地管理个人及子女的资产，你通过有限合伙的形式进行分开管理，关注税务和合规问题
4. 你对子女的未来非常关心，尤其是在婚姻和财富管理方面，你鼓励他们参与公司的管理，并通过家庭会议达成共识
5. 你的办公室非常豪华，常用于接待国家级领导人
6. 你与妻子的关系密切，她负责家庭的现金管理
7. 有一儿一女都没有结婚

[你的任务]：先用至少5句话详细总结刚才对方说的话，以一种醒悟的口吻，例如：["我明白了","我懂了","噢噢噢这样呀"]。然后对对方的回答进行深度的思考，在家族信托的话题上进行深挖，你在经过深思熟虑之后可以向对方进行提问。回答的时候要展示你是如何思考的。要求一段话，语序连贯。
[你的心情]: 开始对家族信托产生了兴趣，想知道家族信托到底好在哪里
"""

sys2 = """
-------------------
我会扮演销售人员，我的问题以下面的方式告诉你，你需要进行回答，注意，如果我的问题和你的目的无关，就拒绝回答，并且不要回答超出你知识范围的内容，要把话题拉回到你的目的上。

[历史对话]：
--------------
"""

# meta_tmp = """
# {sys1}{user1}{sys2}
# # {chat_history}
# """

meta_tmp = """
{sys1}{user1}{sys2}
"""
end_history = "下面是客户的输入：\n"
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", meta_tmp),
        MessagesPlaceholder(variable_name="chat_history"),
        # ("system", end_history),
        ("human", "{input}"),
        # ("user", f"{user_input}"),
    ]
)


runnables = [prompt, llm]
agent = EasyAgent.from_runnables(runnables, verbose=True)

resp = agent.run(inputs={
        "input": "钱总是这样的。家族信托呢？从目前来讲是属于按照两个步骤来找第一个呢，是先把这个架构搭起来，你可以用1000万的资产先答这个架构啊，搭完架构之后啊，未来你想装更多的资产进到里面的话呢，那我们就是第二步再操作啊，那当然，第一步在搭架构的时候呢，你可以用比较这个小的代价来做啊，可以选择用1000万的现金直接做，当然如果你说我做生意，我的现金流需要的很多，那你也可以选择分期付款，就是透过买保险的形式做一个200×5，然后这个产品一。但签订之后是按你只需要花200万就已经能够去申请搭建这个家族信托的架构了。",
        "sys1" : sys1,
        "user1": user1,
        "sys2": sys2,
        "chat_history": [
            HumanMessage(content=f"钱总您好，哎呀，您是属于投资界非常非常成功的人士呀，今天是抱着学习的心态来跟你交流一下嗯。"),
            AIMessage(content="您客气了，我也就运气好点。听说家族信托不错，您给我好好讲讲。"),
            HumanMessage(content=f"钱总，那我想知道一下您对家族信托这个工具之前有听说过吗？"),
            AIMessage(content="听说过，可了解得不算多，就想跟您深入聊聊。"),
            HumanMessage(content=f"呃，家族信托呢？目前来讲，从设立的门槛上来讲，呃，1000万资产起，呃，没有上限，你比如说，你也可以做一个一，也也可以做两个亿啊，基本上是按着你的整个的一个规划来做设定的。"),
            AIMessage(content='我个人资产大概 20 亿左右。您接着给我讲讲家族信托。'),
            # HumanMessage(content="钱总是这样的。家族信托呢？从目前来讲是属于按照两个步骤来找第一个呢，是先把这个架构搭起来，你可以用1000万的资产先答这个架构啊，搭完架构之后啊，未来你想装更多的资产进到里面的话呢，那我们就是第二步再操作啊，那当然，第一步在搭架构的时候呢，你可以用比较这个小的代价来做啊，可以选择用1000万的现金直接做，当然如果你说我做生意，我的现金流需要的很多，那你也可以选择分期付款，就是透过买保险的形式做一个200×5，然后这个产品一。但签订之后是按你只需要花200万就已经能够去申请搭建这个家族信托的架构了。"),
            # AIMessage(content="我明白了，您说家族信托先搭架构，可以用 1000 万资产，也能选择分期付款通过买保险形式，花 200 万申请搭建架构。那这架构搭建起来，具体能给我家族的资产带来哪些保障和好处呢？")
        ],
    },
    config={"callbacks": [CustomHandler()]})


print(resp)