from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts.prompt import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser

from extraction_info import neo4j_url, neo4j_username, neo4j_password

qwen_llm = ChatOpenAI(
    openai_api_base="https://api.siliconflow.cn/v1",
    model_name="Qwen/Qwen2.5-72B-Instruct",
    # sk-muuqautpcyuowjtgfecbnivqodlhzydtfslqkmwbknawejsx
    openai_api_key="sk-bbeamiumkouptsrueilgufqqyuumelcsivxwjbdugqwsqhwj",
    temperature=0.1
)

deep_v3 = ChatOpenAI(
    openai_api_base="https://api.siliconflow.cn/v1",
    model_name="deepseek-ai/DeepSeek-V3",
    # sk-muuqautpcyuowjtgfecbnivqodlhzydtfslqkmwbknawejsx
    openai_api_key="sk-bbeamiumkouptsrueilgufqqyuumelcsivxwjbdugqwsqhwj",
    temperature=0.1
)

def Problem_rewrite():
    PromptTemplate1 = """
    请根据用户的输入内容，替换其中内容：  

    【用户输入】：
    {query}

    【检索内容】：
    {retriever}

    【举例】：
    用户输入：塔材装材费是多少？
    检索内容：角钢塔_塔材装材费_元
    得到的结果：【角钢塔_塔材装材费_元】是多少？

    【要求】：
    - 不允许输出任何解释、标点符号或额外内容，仅在原输入上进行替换
    - 替换后的内容要加上【】
    """

    Prompt = ChatPromptTemplate.from_template(PromptTemplate1)

    Chain = Prompt | deep_v3 | StrOutputParser()
    # Chain = Prompt | llm | StrOutputParser()

    return Chain 

def question_answer():
    PromptTemplate1 = """
    请根据用户的输入内容，和检索到的信息，回答问题：  

    【用户输入】：
    {query}

    【实际检索目标】：
    {retriever_keywords}

    【检索内容】：
    {retriever_info}

    【要求】：
    - 不允许输出任何解释、标点符号或额外内容
    - 要选找到和实际检索目标最接近的检索内容，直接读取相关信息进行解答，不做任何运算
    - 检索到的信息不一定有用，如果和问题毫无相关性，则回答不会
    """

    Prompt = ChatPromptTemplate.from_template(PromptTemplate1)

    Chain = Prompt | deep_v3 | StrOutputParser()
    # Chain = Prompt | llm | StrOutputParser()

    return Chain 


def question_answer_calculation():
    PromptTemplate1 = """
    请根据用户的实际检索目标查到的内容，基于用户输入和计算公式回答问题：  

    【用户输入】：
    {query}

    【实际检索目标】：
    {retriever_keywords}

    【计算公式】：
    {calculation}

    【检索内容】：
    {retriever_info}

    【要求】：
    - 不允许输出任何解释、标点符号或额外内容
    - 要选找到和实际检索目标最接近的检索内容，直接读取相关信息进行解答，不做任何运算
    - 检索到的信息不一定有用，如果和问题毫无相关性，则回答不会
    """

    Prompt = ChatPromptTemplate.from_template(PromptTemplate1)

    Chain = Prompt | deep_v3 | StrOutputParser()
    # Chain = Prompt | llm | StrOutputParser()

    return Chain 


"""neo4j"""

from langchain_community.graphs import Neo4jGraph



graph = Neo4jGraph(
    url = neo4j_url,
    username = neo4j_username,
    password = neo4j_password,
)

graph.refresh_schema()


from langchain.prompts import (
    PromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    ChatPromptTemplate,
)

cypher_generation_template = """
# 任务：
为 Neo4j 图数据库生成 Cypher 查询。

# 说明：
仅使用架构中提供的关系类型和属性。 不得使用架构中未提供的关系类型或属性。

# 架构：
{schema}

# 注意：
不得在回答中包含任何解释或道歉。 
不得回答任何不要求构造 Cypher 查询的问题。
回答中不得包含除生成的 Cypher 查询以外的任何文本。
请确保查询中关系的方向正确，并正确为实体和关系设置别名。
在查询中使用 WITH 为后续语句设置别名（例如，WITH v as visit, c.billing_amount as billing_amount）。
如果需要进行数字除法运算，请确保对分母进行非零过滤。

# 注意
在查询关键字时，不能使用MATCH (n:PropertyNode)，要用最原始的MATCH (n)
在查询中分析缺失属性时，必须使用 IS NULL 或 IS NOT NULL。 
不得在查询中返回嵌套属性。
不得在查询中包含 "GROUP BY" 语句。
请确保为所有后续语句使用 WITH 设置别名（例如，WITH v as visit, c.billing_amount as billing_amount）。
如果需要进行数字除法运算，请确保对分母进行非零过滤。

# 示例：

1.
输入：
查找一下【工程数据/安装工程/安装/架空输电线路本体工程/基础工程】的【定额】下的【YX2-1/YX2-2/YX2-3/YX2-4/YX2-5/YX2-6/YX2-7】
输出：
MATCH (root:ProjectDivisionItem [name: '架空输电线路本体工程'])-[:HAS_CHILD]->(base:ProjectDivisionItem [name: '基础工程'])
MATCH path = (base)-[:HAS_CHILD]->(child)-[:HAS_COMPONENT*1..]->(component)
WHERE ANY(prop IN keys(component) WHERE toString(component[prop]) =~ 'YX2-[1-7]')
RETURN component

2.
输入：
查找一下【工程数据/安装工程/安装/架空输电线路本体工程/杆塔工程/杆塔组立/铁塔、钢管杆组立】的【主材】下的【角钢】
输出：
MATCH (root:ProjectDivisionItem [name: '架空输电线路本体工程'])-[:HAS_CHILD]->(base1:ProjectDivisionItem [name: '杆塔工程'])-[:HAS_CHILD]->(base2:ProjectDivisionItem [name: '杆塔组立'])-[:HAS_CHILD]->(base3:ProjectDivisionItem [name: '铁塔、钢管杆组立'])
MATCH path = (base3)-[:HAS_COMPONENT*1..]->(component)
WHERE ANY(prop IN keys(component) WHERE toString(component[prop]) CONTAINS '角钢') 
RETURN component

3.
输入：
查找一下【架空输电线路本体工程/架线工程/导地线架设】和【架空输电线路本体工程/架线工程/导地线跨越架设】的【主材】下的【高导电率】
输出：
MATCH (root:ProjectDivisionItem [name: '架空输电线路本体工程'])
-[:HAS_CHILD]->(base1:ProjectDivisionItem [name: '架线工程'])-[:HAS_CHILD]->(base:ProjectDivisionItem)
WHERE base.name IN ['导地线架设', '导地线跨越架设']
MATCH path = (base)-[:HAS_COMPONENT*1..]->(component)
WHERE ANY(prop IN keys(component) WHERE toString(component[prop]) CONTAINS '高导电率') 
RETURN component


问题：
{question}
"""

cypher_generation_prompt = PromptTemplate(
    input_variables=["schema", "question"], template=cypher_generation_template
)

qa_generation_template = """你是一个助手，根据 Neo4j Cypher 查询的结果生成可读的回答。查询结果部分包含根据用户的自然语言问题生成的 Cypher 查询结果。提供的信息是权威的；你必须始终使用这些信息来构建回答，不得使用内部知识来质疑或更正这些信息。确保回答听起来像是对问题的直接回应。

用户提出了以下问题：
{question}

运行了一个 Cypher 查询，生成了以下结果：
{context}

如果提供的信息是空的，就说明你不知道答案。
空的信息看起来是这样的：[]

如果查询结果不为空，你必须提供一个回答。

如果有查询结果数据，绝不能说你没有正确的信息。如果用户提问时需要显示所有相关查询结果，确保你显示所有相关结果。你必须始终假设提供的查询结果与问题相关。回答时只能基于提供的查询结果构建答案。

"""

qa_generation_prompt = PromptTemplate(
    input_variables=["context", "question"], template=qa_generation_template
)

from langchain.chains import GraphCypherQAChain

booway_cypher_chain = GraphCypherQAChain.from_llm(
    allow_dangerous_requests=True,
    cypher_llm=deep_v3,                       # 用于生成Cypher查询的LLM
    qa_llm=deep_v3,                           # 用于根据Cypher查询结果生成答案的LLM
    graph=graph,
    verbose=True,
    qa_prompt=qa_generation_prompt,
    cypher_prompt=cypher_generation_prompt,
    validate_cypher=True,
    top_k=100,
    return_intermediate_steps=True
)