from chains_lab import Problem_rewrite
from vector_lab import intersection_of_three_lists
from utils import find_target_item, find_target_items, pre_mapping, pre_mapping2
import json

# 初始化
problem_rewrite = Problem_rewrite()

from utils import extract_concrete_info, extract_query_prefix_list

from chains_lab import question_answer, question_answer_calculation

qa_chains = question_answer()

calculation_chains = question_answer_calculation()

from chains_lab import  booway_cypher_chain

# 加载数据
with open('./data/data.json', 'r', encoding='utf-8') as file:
    data = json.load(file)

print("📥 请输入查询内容，输入 'exit' 可退出程序。\n")

while True:
    input_str = input("🔍 输入问题：")

    if input_str.lower() == 'exit':
        print("👋 已退出。")
        break

    try:
        results = intersection_of_three_lists(input_str)
        if not results:
            print("⚠️ 无法从向量中获取候选项。")
            continue

        retriever = results[0]

        print(f"➡️ 匹配向量检索结果：{retriever}")

        # 重写问题，提取关键词
        keywords = problem_rewrite.invoke({
            "query": input_str,
            "retriever": retriever
        })

        print(f"🧠 提取关键词：{keywords}")

        # 预映射为图数据库结构
        input_neo4j = pre_mapping2(keywords, data)

        print(f"📊 图谱相关知识输出：\n{input_neo4j}\n")

        question = input_neo4j
        input_str = input_str

        if isinstance(question, str):
    
            response = booway_cypher_chain.invoke(question)
    
            generated_cypher = response.get("intermediate_steps")[0]
    
            # print(str(generated_cypher))

            temp = response.get("result")

            # print(temp)
    
            finally_result = qa_chains.invoke({"query":input_str, 
                                               "retriever_keywords":keywords, 
                                               "retriever_info":temp})

            print(f"📊 图谱最终检索输出：\n{finally_result}\n")


        elif isinstance(question, list):
            ques = extract_query_prefix_list(question)
            ques_info = extract_concrete_info(question)

            retriever_info = []
            for i in ques:
                response = booway_cypher_chain.invoke(ques)    
                # generated_cypher = response.get("intermediate_steps")[0]
                temp = response.get("result")
                retriever_info.append(temp)

            retriever_keywords = ques_info[0]
            calculation = ques_info[-1]
    
            finally_result = calculation_chains.invoke({"query":question, 
                                                        "retriever_keywords":retriever_keywords, 
                                                        "calculation":calculation, 
                                                        "retriever_info":retriever_info})
            print(f"📊 图谱最终检索输出：\n{finally_result}\n")

    except Exception as e:
        print(f"❌ 发生错误：{e}")