import json
import fitz  # PyMuPDF
import os
import fitz
from collections import Counter
import uuid
from kg_build.Prompt import Prompt_compose_tujie1,Prompt_compose_tujie2
from kg_build.extract_util import safe_json_loads,highlight_pdf_file,get_content_bbox
# from modelsAPI.model_api import OpenaiAPI
import re
from openai import AsyncOpenAI
import asyncio
import json
import re
import copy
import pandas as pd
from io import StringIO
import json
from collections import defaultdict
import traceback


def is_tujie_compose_table_by_first_tr(table_body: str) -> bool:
    if not isinstance(table_body, str) or not table_body.strip():
        return False

    text_without_tags = re.sub(r'<[^>]+>', '', table_body)
    text_lower = text_without_tags.lower()

    keyword_groups = [
        ["组成编码", "名称", "是否为关重件", "是否为寿命件"],
        ["组成编码", "名称", "型号/零件号", "重量(kg)", "数量"]
    ]

    for group in keyword_groups:
        if all(kw in text_lower for kw in group):
            return True

    return False


def is_tujie_compose_mulus_table_by_first_tr(table_body: str) -> bool:
    if not isinstance(table_body, str) or not table_body.strip():
        return False

    text_without_tags = re.sub(r'<[^>]+>', '', table_body)
    text_lower = text_without_tags.lower()

    keyword_groups = [
        ["组成编码", "名称", "型号/零件号", "重量(kg)", "数量"]
    ]

    for group in keyword_groups:
        if all(kw in text_lower for kw in group):
            return True

    return False


def get_tujiecompose_chunk_bbox_and_removed_data(data):
    """
    返回：
      - 符合条件的图解组成表分组（每组 [前前, 前, 当前]）
      - 从原始 data 中移除了所有相关记录后的新 data 列表
    """
    title = "组成"
    record_map = {item["id"]: item for item in data}

    results = []
    ids_to_remove = set()

    for ins in data:
        if ins.get("type") == "table":
            caption_has_title = any(title in caption for caption in ins.get("table_caption", []))

            table_body = ins.get("table_body", "")
            body_matches_keywords = is_tujie_compose_table_by_first_tr(table_body)

            if caption_has_title and body_matches_keywords:
                current_id = ins["id"]
                prev_ids = [current_id - 2, current_id - 1]

                group_ids = [pid for pid in prev_ids if pid >= 1] + [current_id]
                ids_to_remove.update(group_ids)

                group = [record_map.get(pid) for pid in prev_ids if pid >= 1] + [ins]
                results.append(group)

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


def get_tujiecompose_mulu_chunk_bbox_and_removed_data(data):
    title = "组成"
    record_map = {item["id"]: item for item in data}

    results = []
    ids_to_remove = set()

    for ins in data:
        if ins.get("type") == "table":
            caption_has_title = any(title in caption for caption in ins.get("table_caption", []))

            table_body = ins.get("table_body", "")
            body_matches_keywords = is_tujie_compose_mulus_table_by_first_tr(table_body)

            if body_matches_keywords:
                current_id = ins["id"]
                prev_ids = [current_id - 2, current_id - 1]

                group_ids = [pid for pid in prev_ids if pid >= 1] + [current_id]
                ids_to_remove.update(group_ids)

                group = [record_map.get(pid) for pid in prev_ids if pid >= 1] + [ins]
                results.append(group)

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


def format_compose_code(val):
    """将 int/float 型组成编码转换为偶数位字符串，不足偶数位则前补 0"""
    if pd.isna(val):
        return None
    if isinstance(val, (int, float)):
        s = str(int(val))
        if len(s) % 2 != 0:
            s = '0' + s
        return s
    return str(val)


def merge_compose(records):
    """
    合并相同组成编码的记录。
    特殊处理：将 knowledge_source 字段的多个值解析并合并到一个新的扁平列表中。
    """
    grouped_data = defaultdict(list)

    for record in records:
        code = record.get('组成编码')
        if code:
            grouped_data[code].append(record)

    merged_data = []

    for key, group in grouped_data.items():
        merged_record = group[0].copy()

        all_sources = []

        for record in group:
            source_val = record.get('knowledge_source')
            if source_val:
                try:
                    if isinstance(source_val, str):
                        parsed_source = json.loads(source_val)
                    else:
                        parsed_source = source_val

                    if isinstance(parsed_source, list):
                        all_sources.extend(parsed_source)
                    else:
                        all_sources.append(parsed_source)
                except (json.JSONDecodeError, TypeError):
                    all_sources.append(source_val)

        if all_sources:
            merged_record['knowledge_source'] = json.dumps(all_sources, ensure_ascii=False)
        else:
            merged_record['knowledge_source'] = "[]"

        for record in group[1:]:
            for k, v in record.items():
                if k == 'knowledge_source':
                    continue

                existing = merged_record.get(k)

                if v and existing and existing != v:
                    if isinstance(existing, list):
                        if v not in existing:
                            existing.append(v)
                    else:
                        merged_record[k] = [existing, v]
                elif not existing and v:
                    merged_record[k] = v

        merged_data.append(merged_record)

    return merged_data


def remove_fields(data_list, fields_to_remove):
    """
    遍历列表中的每个字典，移除指定的字段（键）。
    如果字段不存在，则跳过该字段。
    """
    cleaned_data = []

    for item in data_list:
        new_item = item.copy()

        for field in fields_to_remove:
            if field in new_item:
                del new_item[field]

        cleaned_data.append(new_item)

    return cleaned_data


# def generate_entities_and_relations(data_list):
#     """
#     所有实体统一 type="设备"。
#     关系构建：若父节点在数据中不存在，自动补充占位实体，确保所有关系都成功建立。
#     """
#     entities = []
#     relationships = []

#     code_name_map = {}   # {code: name}
#     clean_nodes = []

#     # --- 第一步：数据清洗与标准化 ---
#     for item in data_list:
#         code = item.get('组成编码', '')

#         if not code or not isinstance(code, str):
#             continue
#         if not code.isdigit():
#             continue
#         if len(code) < 2 or len(code) % 2 != 0:
#             continue

#         cleaned_item = {}
#         for key, value in item.items():
#             if isinstance(value, list):
#                 cleaned_item[key] = str(value[0]) if value else ""
#             elif value is None:
#                 cleaned_item[key] = ""
#             else:
#                 cleaned_item[key] = value

#         code_name_map[code] = cleaned_item.get('名称', f'未知_{code}')
#         clean_nodes.append(cleaned_item)

#     # --- 第二步：找出所有缺失的父节点，自动递归向上补全 ---
#     placeholder_codes = {}  # {code: name}  仅补充的占位节点

#     for node in clean_nodes:
#         code = node['组成编码']
#         length = len(code)

#         if length < 4:
#             continue  # 2位根节点无父节点，跳过

#         current = code[:-2]
#         while current:
#             if current in code_name_map or current in placeholder_codes:
#                 break  # 已存在，无需补全
#             placeholder_name = f'未知_{current}'
#             placeholder_codes[current] = placeholder_name
#             print(f"  ⚠️ 父节点缺失，自动补充占位实体: 编码={current}, 名称={placeholder_name}")
#             if len(current) >= 4:
#                 current = current[:-2]
#             else:
#                 break

#     # 将占位节点合并进 code_name_map，供关系构建使用
#     code_name_map.update(placeholder_codes)

#     # --- 第三步：构建实体列表，所有节点 type="设备" ---
#     # 真实节点
#     for node in clean_nodes:
#         entities.append({
#             "type": "设备",
#             "properties": node
#         })

#     # 占位节点（最简属性，标记 _is_placeholder 便于下游过滤）
#     for code, name in placeholder_codes.items():
#         entities.append({
#             "type": "设备",
#             "properties": {
#                 "组成编码": code,
#                 "名称": name,
#                 "knowledge_source": "[]",
#                 "_is_placeholder": True
#             }
#         })

#     # --- 第四步：构建关系列表，保证全部成功 ---
#     success_count = 0
#     skip_count = 0

#     for node in clean_nodes:
#         code = node['组成编码']
#         length = len(code)

#         if length < 4 or length % 2 != 0:
#             skip_count += 1
#             continue  # 2位根节点无父，正常跳过

#         parent_code = code[:-2]
#         child_name = node.get('名称', f'未知_{code}')
#         parent_name = code_name_map.get(parent_code)

#         if parent_name:
#             relationships.append({
#                 "type": "包含",
#                 "from_entity": parent_name,
#                 "to_entity": child_name
#             })
#             success_count += 1
#         else:
#             # 经过第二步补全后理论上不会进入此分支
#             print(f"  ❌ 关系构建失败（不应出现）: 子编码={code}, 父编码={parent_code} 仍未找到")
#             skip_count += 1

#     print(f"  📊 关系构建完成：成功={success_count}, 跳过根节点={skip_count}")
#     return entities, relationships
def generate_entities_and_relations(data_list):
    """
    实体分两类：
      - type="设备"：原始组成表中的真实/占位节点
      - type="组成"：每个有子节点的设备自动生成的中间节点，名称为 "{设备名}的组成"

    关系：
      - 父节点（设备）→ 中间节点（组成）：包含
      - 中间节点（组成）→ 子节点（设备）：包含
    """
    entities = []
    relationships = []

    code_name_map = {}   # {code: name}
    clean_nodes = []

    # ── 第一步：数据清洗与标准化 ──────────────────────────────────────
    for item in data_list:
        code = item.get('组成编码', '')
        if not code or not isinstance(code, str):
            continue
        if not code.isdigit():
            continue
        if len(code) < 2 or len(code) % 2 != 0:
            continue

        cleaned_item = {}
        for key, value in item.items():
            if isinstance(value, list):
                cleaned_item[key] = str(value[0]) if value else ""
            elif value is None:
                cleaned_item[key] = ""
            else:
                cleaned_item[key] = value

        code_name_map[code] = cleaned_item.get('名称', f'未知_{code}')
        clean_nodes.append(cleaned_item)

    # ── 第二步：自动补全缺失的父节点占位实体 ─────────────────────────
    placeholder_codes = {}  # {code: name}

    for node in clean_nodes:
        code = node['组成编码']
        current = code[:-2]
        while current:
            if current in code_name_map or current in placeholder_codes:
                break
            placeholder_name = f'未知_{current}'
            placeholder_codes[current] = placeholder_name
            print(f"  ⚠️ 父节点缺失，自动补充占位实体: 编码={current}, 名称={placeholder_name}")
            current = current[:-2] if len(current) >= 4 else ""

    code_name_map.update(placeholder_codes)

    # ── 第三步：构建"设备"实体（真实节点 + 占位节点）────────────────
    for node in clean_nodes:
        entities.append({"type": "设备", "properties": node})

    for code, name in placeholder_codes.items():
        entities.append({
            "type": "设备",
            "properties": {
                "组成编码": code,
                "名称": name,
                "knowledge_source": "[]",
            }
        })

    # ── 第四步：为每个"有子节点的父设备"生成"组成"中间实体 ──────────
    # 收集所有实际存在父子关系的父编码
    parent_codes_needed = set()
    for node in clean_nodes:
        code = node['组成编码']
        if len(code) >= 4 and len(code) % 2 == 0:
            parent_codes_needed.add(code[:-2])

    # {parent_code: "组成中间实体名称"}
    compose_node_map = {}
    for parent_code in parent_codes_needed:
        parent_name = code_name_map.get(parent_code)
        if not parent_name:
            print(f"  ❌ 找不到父编码对应名称（不应出现）: {parent_code}")
            continue
        compose_name = f"{parent_name}的组成"
        compose_node_map[parent_code] = compose_name

        # 生成"组成"中间实体
        entities.append({
            "type": "组成",
            "properties": {
                "名称": compose_name,
                "knowledge_source": "[]",
            }
        })

        # 父设备 → 组成中间节点
        relationships.append({
            "type": "包含",
            "from_entity": parent_name,
            "to_entity": compose_name
        })

    # ── 第五步：构建"组成中间节点 → 子设备"关系 ──────────────────────
    success_count, skip_count = 0, 0

    for node in clean_nodes:
        code = node['组成编码']
        if len(code) < 4 or len(code) % 2 != 0:
            skip_count += 1
            continue  # 根节点无父，正常跳过

        parent_code = code[:-2]
        compose_name = compose_node_map.get(parent_code)
        child_name = node.get('名称', f'未知_{code}')

        if compose_name:
            relationships.append({
                "type": "包含",
                "from_entity": compose_name,
                "to_entity": child_name
            })
            success_count += 1
        else:
            print(f"  ❌ 找不到组成中间节点（不应出现）: 子编码={code}, 父编码={parent_code}")
            skip_count += 1

    print(f"  📊 关系构建完成：成功={success_count}, 跳过根节点={skip_count}")
    return entities, relationships

async def get_tujiecompose_node_relation(input_pdf_path, data, prefix_url="http://192.168.0.46:59085", device="发动机", lower_entity_type="设备"):
    # 初始化默认返回值（空列表）
    all_entity_results = []
    all_relation_results = []
    devicename = device
    lower_entity_type = lower_entity_type

    try:
        # 1. 初始化 ID（确保从 1 开始连续）
        for idx, item in enumerate(data, start=1):
            item["id"] = idx

        # 2. 获取分组数据
        results, filtered_data = get_tujiecompose_chunk_bbox_and_removed_data(data=data)
        compose_mulu_result, filtered_data_mulu = get_tujiecompose_mulu_chunk_bbox_and_removed_data(data=data)

        # 3. 合并所有需要处理的组
        all_results = results + compose_mulu_result

        # 初始化中间结果容器
        all_htmltojson_results = []

        print(f"共发现 {len(all_results)} 个图解组成表组，开始处理...")

        for i, compose_group in enumerate(all_results):
            try:
                print(f"开始处理第 {i+1} 个图解组成表")

                content, highlight_list = get_content_bbox(compose_group)

                if not content:
                    print(f"  ⚠️ 第 {i+1} 组未获取到有效内容，跳过")
                    continue

                # 文本清洗
                content = content.replace("香", "否")
                content = content.replace("房号", "序号")

                # 解析 HTML 表格
                html_obj = StringIO(content)
                df_list = pd.read_html(html_obj, header=0)

                if not df_list:
                    print(f"  ⚠️ 第 {i+1} 组未解析到表格，跳过")
                    continue

                df = df_list[0]

                # 清洗数据：删除全为空的行
                df.dropna(how='all', inplace=True)

                # 格式化组成编码
                if '组成编码' in df.columns and pd.api.types.is_numeric_dtype(df['组成编码']):
                    df['组成编码'] = df['组成编码'].apply(format_compose_code)

                # 转换为 JSON 列表
                json_data = df.to_json(orient='records', force_ascii=False)
                result_list = json.loads(json_data)
                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}"
                if result_list:
                    for ins in result_list:
                        source_data_list = [{
                            "filename": filename,
                            "info": content,
                            "url": url
                        }]
                        knowledge_source_str = json.dumps(source_data_list, ensure_ascii=False)
                        ins["knowledge_source"] = knowledge_source_str

                    all_htmltojson_results.extend(result_list)
                    print(f"  ✅ 第 {i+1} 组提取成功，包含 {len(result_list)} 条记录")
                else:
                    print(f"  ⚠️ 第 {i+1} 组提取结果为空")

            except Exception as e:
                print(f"❌ 处理第 {i+1} 组时发生错误: {e}")
                traceback.print_exc()
                continue

        # 后处理步骤
        if all_htmltojson_results:
            all_htmltojson_results = merge_compose(all_htmltojson_results)
            all_htmltojson_results = remove_fields(all_htmltojson_results, ["序号"])

            # ✅ 修复：先生成实体和关系
            all_entity_results, all_relation_results = generate_entities_and_relations(all_htmltojson_results)

            # ✅ 修复：循环移到 generate 之后，此时 all_entity_results 才有数据
            for ins in all_entity_results:
                if ins["properties"].get("名称") == devicename:
                    ins["type"] = lower_entity_type
        else:
            print("⚠️ 没有提取到任何有效数据，跳过实体关系生成。")

        return all_entity_results, all_relation_results

    except Exception as e:
        print(f"❌ get_tujiecompose_node_relation 主流程发生严重错误: {e}")
        traceback.print_exc()
        return [], []


if __name__ == "__main__":
    file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/output1/122-06A0014-B01001_发动机-图解目录_content_list.json'
    input_pdf_path = "/storage01/home/hdf/project/wxkgrag/kgextract/know/122-06A0014-B01001_发动机-维修手册_highlighted_fdb9e9db-98d2-4ef6-bc72-52f875453f85.pdf"
    output_json_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/output2/图解组成表抽取结果.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 = []

    all_entity_results, all_relation_results = asyncio.run(
        get_tujiecompose_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} 个关系")

    result = {
        "entities": all_entity_results,
        "relations": all_relation_results
    }
    for ins in result["entities"]:
        print(ins)
        print(111111111111111111111)

    try:
        with open(output_json_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=4)
        print(f"结果已保存至：{output_json_path}")
    except Exception as e:
        print(f"保存 JSON 文件时出错：{e}")

