import json
import fitz  # PyMuPDF
import os
from collections import Counter
from kg_build.Prompt import prompt_extract_beipinbeijian_node_rela
from openai import OpenAI
from extract_util import safe_json_loads, highlight_pdf_file, get_content_bbox
from modelsAPI.model_api import OpenaiAPI

# 指定 JSON 文件路径
# file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/122-06A0014-B01001_发动机-维修手册_content_list.json'

# try:
#     with open(file_path, 'r', encoding='utf-8') as file:
#         data = json.load(file)
# except Exception as e:
#     print(f"读取 JSON 文件时出错：{e}")
#     data = []

# 为每条记录添加 id（从1开始）


def get_beipinbeijian_chunk_bbox_and_removed_data(data):
    # 定义三种可能的表头格式
    labels = [
        '<table><tr><td rowspan=1 colspan=1>序</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型专</td><td rowspan=1 colspan=1>规格</td>',
        '<table><tr><td rowspan=1 colspan=1>序号</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型号</td><td rowspan=1 colspan=1>规格</td>',
        '<<table><tr><td rowspan=2 colspan=1>序号</td><td rowspan=2 colspan=1>名称</td><td rowspan=2 colspan=1>型号</td><td rowspan=2 colspan=1>规格</td>',
        '<table><tr><td rowspan=2 colspan=1>序号</td><td rowspan=2 colspan=1>名称</td><td rowspan=2 colspan=1>型号</td><td rowspan=2 colspan=1>规格</td>',
        '<table><tr><td rowspan=1 colspan=1>房</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型专</td><td rowspan=1 colspan=1>规格</td>',
        '<table><tr><td rowspan=1 colspan=23>备品备件、附件和工具清单</td></tr><tr><td rowspan=1 colspan=9>设备备品备件清单</td><td rowspan=1 colspan=5>设备型号</td><td rowspan=1 colspan=3></td><td rowspan=1 colspan=4>设备装舰台套数</td>',
        '<table><tr><td rowspan=2 colspan=1>序号</td><td rowspan=2 colspan=1>名      称</td><td rowspan=2 colspan=1>型号规格</td><td rowspan=2 colspan=1>图号或标准号</td>',
        '<table><tr><td rowspan=1 colspan=1>序号</td><td rowspan=1 colspan=1>名称</td><td rowspan=1 colspan=1>型专</td><td rowspan=1 colspan=1>规格</td>'
    ]

    record_map = {item["id"]: item for item in data}
    results = []
    ids_to_remove = set()

    for ins in data:
        if ins.get("type") == "table":
            table_body = ins.get("table_body", "")
            if any(label in table_body for label in labels):
                current_id = ins["id"]
                # 获取当前及前4个记录的 id（确保 >=1）
                group_ids = [current_id - i for i in range(4, -1, -1) if current_id - i >= 1]
                ids_to_remove.update(group_ids)
                # 按顺序构建 group（从最早到当前）
                group = [record_map[pid] for pid in group_ids if pid in record_map]
                results.append(group)

    filtered_data = [item for item in data if item["id"] not in ids_to_remove]
    return results, filtered_data


async def extract_beipinbeijian_table_entity(chunk):

    print("提取备品备件类实体和关系")

    final_prompt = prompt_extract_beipinbeijian_node_rela.format(text=chunk)
    raw_response = await OpenaiAPI.openai_chat_aysnc(final_prompt,timeout=600)

    if raw_response is None:
        print("API 调用失败，跳过当前 chunk")
        return []

    try:
        cleaned = raw_response.strip()
        if cleaned.startswith("```json"):
            cleaned = cleaned[7:].lstrip()
        if cleaned.endswith("```"):
            cleaned = cleaned[:-3].rstrip()

        json_result = safe_json_loads(cleaned)
        if json_result is None:
            print("JSON 解析失败: 返回 None")
            return []

        return json_result

    except Exception as e:
        print(f"⚠️ JSON 解析异常: {e}")
        print(f"原始响应片段: {repr(raw_response[:500])}")
        return []




async def get_beipinbeijian_node_relation(input_pdf_path,data,prefix_url = "http://192.168.0.46:59085"):
    for idx, item in enumerate(data, start=1):
        item["id"] = idx
    # 主流程
    guzhang_groups, cleaned_data = get_beipinbeijian_chunk_bbox_and_removed_data(data)

    if not guzhang_groups:
        print("未找到符合条件的备品备件表格，跳过处理。")
        all_entity_results,all_relation_results = [],[]
    else:
        all_entity_results,all_relation_results = [],[] # 用于收集所有成功提取的结果

        for i, compose in enumerate(guzhang_groups):
            try:
                print(f"开始处理第{i+1}个备品备件表")
                # 获取内容和高亮信息
                content, highlight_list = get_content_bbox(compose)
                print("备品备件表切片内容为：",content)

                # 提取实体
                result =await  extract_beipinbeijian_table_entity(content)
                if result is []:
                    print(f"第 {i+1} 组：实体提取失败，跳过")
                    continue

                # 安全访问 entities 字段
                entities = result.get("entities", [])
                if not isinstance(entities, list):
                    print(f"第 {i+1} 组：entities 字段格式异常，跳过")
                    continue
                relationships = result.get("relationships", [])

                # 为每个 entity 添加 knowledge_source
                
                highlight_filename, most_common_page,filename = highlight_pdf_file(input_pdf_path,highlight_list)
                url = f"{prefix_url}/upload/{highlight_filename}#page={most_common_page}"
                # for entity in entities:
                #     knowledge_source = [{
                #         "filename": filename,
                #         "info": content,
                #         "url": url
                #     }]
                #     # 确保 properties 存在
                #     if "properties" not in entity or not isinstance(entity["properties"], dict):
                #         entity["properties"] = {}
                #     entity["properties"]["knowledge_source"] = knowledge_source
                for entity in entities:
                        # 构造原始数据对象 (List of Dict)
                        source_data_list = [{
                            "filename": filename,
                            "info": content,
                            "url": url
                        }]
                        
                        # 2. 【关键修改】将列表序列化为 JSON 字符串
                        # ensure_ascii=False 确保中文内容不会被转义成 \uXXXX
                        knowledge_source_str = json.dumps(source_data_list, ensure_ascii=False)

                        # 确保 properties 存在且是字典
                        if "properties" not in entity or not isinstance(entity["properties"], dict):
                            entity["properties"] = {}
                        
                        # 3. 赋值字符串而不是对象列表
                        entity["properties"]["knowledge_source"] = knowledge_source_str
                # 更新 result 中的 entities
                all_entity_results.extend(entities)
                all_relation_results.extend(relationships)

            except Exception as e:
                print(f"处理第 {i+1} 组时发生未预期错误: {e}")
                continue  # 跳过当前组，继续下一轮
    return all_entity_results,all_relation_results 

# if __name__ == "__main__":
#     file_path = '/app/files/122-06A0014-B01001_发动机-维修手册_content_list.json'
#     input_pdf_path = "/app/files/122-06A0014-B01001_发动机-维修手册.pdf"

#     try:
#         with open(file_path, 'r', encoding='utf-8') as file:
#             data = json.load(file)
#     except Exception as e:
#         print(f"读取 JSON 文件时出错：{e}")
#         data = []
#     all_entity_results,all_relation_results = get_guzhang_node_relation(input_pdf_path,data=data)
#     length_entity = len(all_entity_results)
#     length_relation = len(all_relation_results)

#     print(f"抽取{length_entity}个实体，抽取{length_relation}个关系")