import json
import fitz  # PyMuPDF
import os
from collections import Counter
from Prompt import Prompt_guzhang
from openai import OpenAI
from extract_util import safe_json_loads, highlight_pdf_file, get_content_bbox,get_filter_content_bbox




def get_compose_chunk_bbox_and_removed_data(data):
    """
    返回：
      - 符合条件的故障表分组（每组 [前前, 前, 当前]）
      - 从原始 data 中移除了所有相关记录后的新 data 列表
    """
    title = "组成"
    
    # 支持多个 label 模式（可扩展）
    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></tr><tr><td rowspan=1 colspan=1>1</td>',
        # 可在此添加更多变体，例如：
        # '<table><tr><th>序号</th><th>故障现象</th><th>故障原因</th><th>维修项目</th></tr>',
        # '<table><tr><td>序号</td><td>故障现象</td><td>故障原因</td><td>维修项目</td></tr>',
    ]

    # 构建 id -> record 映射
    record_map = {item["id"]: item for item in data}

    results = []
    ids_to_remove = set()

    for ins in data:
        if ins.get("type") == "table":
            # 检查 caption 是否包含 "故障表"
            caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
            
            # 检查 body 是否匹配任意一个 label
            body = ins.get("table_body", "")
            body_matches_any_label = any(label in body for label in labels)

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

                # 收集要移除的 ID（前两条 + 当前）
                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)

    # 过滤原始 data
    filtered_data = [item for item in data if item["id"] not in ids_to_remove]

    return results, filtered_data

def get_tujiecompose_chunk_bbox_and_removed_data(data):
    """
    返回：
      - 符合条件的故障表分组（每组 [前前, 前, 当前]）
      - 从原始 data 中移除了所有相关记录后的新 data 列表
    """
    title = "组成"
    
    # 支持多个 label 模式（可扩展）
    labels = [
        '<table><tr><td colspan=\"1\" rowspan=\"1\">组成编码</td><td colspan=\"1\" rowspan=\"1\">名称</td><td colspan=\"1\" rowspan=\"1\">是否为关重件</td><td colspan=\"1\" rowspan=\"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></tr><tr><td rowspan=1 colspan=1>1</td>'
        # 可在此添加更多变体，例如：
        # '<table><tr><th>序号</th><th>故障现象</th><th>故障原因</th><th>维修项目</th></tr>',
        # '<table><tr><td>序号</td><td>故障现象</td><td>故障原因</td><td>维修项目</td></tr>',
    ]

    # 构建 id -> record 映射
    record_map = {item["id"]: item for item in data}

    results = []
    ids_to_remove = set()

    for ins in data:
        if ins.get("type") == "table":
            # 检查 caption 是否包含 "故障表"
            caption_has_title = any(title in caption for caption in ins.get("table_caption", []))
            
            # 检查 body 是否匹配任意一个 label
            body = ins.get("table_body", "")
            body_matches_any_label = any(label in body for label in labels)

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

                # 收集要移除的 ID（前两条 + 当前）
                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)

    # 过滤原始 data
    filtered_data = [item for item in data if item["id"] not in ids_to_remove]

    return results, filtered_data

def get_guzhang_chunk_bbox_and_removed_data(data):
    title = "故障表"
    label = "<table><tr><td colspan=\"1\" rowspan=\"1\">序号</td><td colspan=\"1\" rowspan=\"1\">故障现象</td><td colspan=\"1\" rowspan=\"1\">故障原因</td><td colspan=\"1\" rowspan=\"1\">维修项目</td></tr>"

    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", []))
            body_has_label = label in ins.get("table_body", "")
            if caption_has_title and body_has_label:
                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


START_MARKERS = [
    # 标准格式（带 rowspan/colspan）
    "<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><td rowspan=1 colspan=1>维修间隔期</td>",
    
    # 可在此添加其他格式，例如：
    # "<table><tr><td>维修项目编号</td><td>组成编码</td><td>名称</td><td>维修项目</td><td>维修间隔期</td>",
]

def extract_maintenance_groups_and_remaining(records, markers=START_MARKERS, max_following_for_last=None):
    """
    普通起始标记：组 = [start, ..., next_start - 1]
    最后一个起始标记：组 = [start, ..., start + max_following_for_last] （若指定）
    """
    groups = []
    current_group = None
    used_ids = set()

    for record in records:
        body = record.get("table_body", "")
        is_start = (record.get("type") == "table" and 
                   any(marker in body for marker in markers))
        
        if is_start:
            if current_group is not None:
                groups.append(current_group)
                used_ids.update(r["id"] for r in current_group)
            current_group = [record]
        else:
            if current_group is not None:
                # 如果指定了 max_following_for_last，且这是最后一个组（暂时不知道是不是最后一个）
                # 我们无法在遍历时知道“这是不是最后一个起始标记”
                # 所以更好的办法是：先按原逻辑分组，再处理最后一个组
                current_group.append(record)

    if current_group is not None:
        groups.append(current_group)
        used_ids.update(r["id"] for r in current_group)

    # —————— 关键：后处理最后一个组 ——————
    if groups and max_following_for_last is not None:
        last_group = groups[-1]
        if len(last_group) > max_following_for_last + 1:  # +1 包含起始记录
            truncated = last_group[:max_following_for_last + 1]
            # 更新 used_ids：移除被截掉的 ID
            removed_ids = {r["id"] for r in last_group[max_following_for_last + 1:]}
            used_ids -= removed_ids
            groups[-1] = truncated

    remaining_records = [r for r in records if r["id"] not in used_ids]
    return groups, remaining_records


START_MARKERS_ope = [
    '<table><tr><td>项目编号</td><td>操作项目</td></tr><tr>',
    '<table><tr><td rowspan=1 colspan=1>项目编号</td><td rowspan=1 colspan=1>操作项目</td></tr><tr>',
    '<table><tr><td colspan="1" rowspan="1">项目编号</td><td colspan="1" rowspan="1">操作项目</td></tr><tr>'
]

def extract_operation_groups_and_remaining(records, markers=START_MARKERS_ope, max_following_for_last=None):
    """
    普通起始标记：组 = [start, ..., next_start - 1]
    最后一个起始标记：组 = [start, ..., start + max_following_for_last] （若指定）
    """
    groups = []
    current_group = None
    used_ids = set()

    for record in records:
        body = record.get("table_body", "")
        is_start = (record.get("type") == "table" and 
                   any(marker in body for marker in markers))
        
        if is_start:
            if current_group is not None:
                groups.append(current_group)
                used_ids.update(r["id"] for r in current_group)
            current_group = [record]
        else:
            if current_group is not None:
                # 如果指定了 max_following_for_last，且这是最后一个组（暂时不知道是不是最后一个）
                # 我们无法在遍历时知道“这是不是最后一个起始标记”
                # 所以更好的办法是：先按原逻辑分组，再处理最后一个组
                current_group.append(record)

    if current_group is not None:
        groups.append(current_group)
        used_ids.update(r["id"] for r in current_group)

    # —————— 关键：后处理最后一个组 ——————
    if groups and max_following_for_last is not None:
        last_group = groups[-1]
        if len(last_group) > max_following_for_last + 1:  # +1 包含起始记录
            truncated = last_group[:max_following_for_last + 1]
            # 更新 used_ids：移除被截掉的 ID
            removed_ids = {r["id"] for r in last_group[max_following_for_last + 1:]}
            used_ids -= removed_ids
            groups[-1] = truncated

    remaining_records = [r for r in records if r["id"] not in used_ids]
    return groups, remaining_records



def get_chunk_filter(data,max_length=7000):
    for idx, item in enumerate(data, start=1):
        item["id"] = idx
    groups, remaining_records_repair = extract_maintenance_groups_and_remaining(data)
    groups, remaining_records_operation = extract_operation_groups_and_remaining(remaining_records_repair)
    groups, remaining_records_compose = get_compose_chunk_bbox_and_removed_data(remaining_records_operation)
    groups, remaining_records_composetujie =get_tujiecompose_chunk_bbox_and_removed_data(remaining_records_compose)
    results, filtered_data=get_guzhang_chunk_bbox_and_removed_data(remaining_records_composetujie)
    res = get_filter_content_bbox(filtered_data,max_length=max_length)
    return res

# if __name__ == "__main__":
#     file_path = '/storage01/home/hdf/project/wxkgrag/kgextract/output/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"
#     try:
#         with open(file_path, 'r', encoding='utf-8') as file:
#             data = json.load(file)
#     except Exception as e:
#         print(f"读取 JSON 文件时出错：{e}")
#         data = []
#     res = get_chunk_filter(data=data)
#     print(res)