import json
from extraction_info import info_data_json
from langchain_neo4j import nl_query_to_function_call
from vector_lab import intersection_of_three_lists
from chains_lab import Problem_rewrite
from utils import find_target_item, find_target_items, pre_mapping

problem_rewrite = Problem_rewrite()

# 读取JSON文件
with open(info_data_json, "r", encoding="utf-8") as f:
    json_data = json.load(f)

while True:
    input_str = input("请输入查询内容（输入 'exit' 退出）：")
    if input_str.lower() == 'exit':
        print("程序结束。")
        break

    try:
        results = intersection_of_three_lists(input_str)
        if not results:
            print("未找到相关指标，请重新输入。")
            continue

        retriever = results[0]
        keywords = problem_rewrite.invoke({"query": input_str, "retriever": retriever})
        input_neo4j = pre_mapping(keywords, json_data)

        print(f"Neo4j 问题理解: {input_neo4j}")
        print(f"Neo4j 查询输入: {retriever}")

        temp_result = None
        for item in json_data:
            if item["指标名称"] == retriever:
                temp_result = item
                break

        if temp_result:
            if temp_result["code"] == "":
                question = {
                    "type": "query",
                    "value": input_neo4j,
                }
                result = nl_query_to_function_call(question)

                if result["message"] == "成功":
                    code_result = result["data"]["code"]
                    temp_result["code"] = code_result

                    # 更新JSON文件
                    with open(info_data_json, "w", encoding="utf-8") as f:
                        json.dump(json_data, f, ensure_ascii=False, indent=4)

                    print(f"已更新 code 字段: {code_result}")
                else:
                    print(f"查询失败：{result}")
            else:
                question = {
                    "type": "code",
                    "value": temp_result["code"],
                }
                result = nl_query_to_function_call(question)
                print(f"查询结果：{result}")

        else:
            print("未找到对应的指标。")

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