import json
import fitz  # PyMuPDF
import os
import fitz
from collections import Counter
import uuid
from openai import OpenAI
from kg_build.Prompt import Prompt_compose
from extract_util import safe_json_loads,highlight_pdf_file,get_content_bbox
import re
from typing import List, Dict, Any, Tuple
from modelsAPI.model_api import OpenaiAPI
import asyncio  # <-- 新增导入

import re

def is_compose_table_by_first_tr(table_body: str) -> bool:
    if not isinstance(table_body, str):
        return False

    # 非贪婪匹配第一个 <tr> ... </tr>
    match = re.search(r'<tr[^>]*>(.*?)</tr>', table_body, re.IGNORECASE | re.DOTALL)
    if not match:
        return False

    first_tr_html = match.group(1)

    # 去除所有 HTML 标签，只保留纯文本
    text_in_first_tr = re.sub(r'<[^>]+>', '', first_tr_html)

    # 统一转为小写便于匹配
    text_lower = text_in_first_tr.lower()

    # 定义多组关键词组合（任意一组满足即可）
    keyword_groups = [
        ["序号", "组成部分", "功能", "数量"],
        ["序号", "组成部分", "功能"]
    ]

    # 检查是否存在至少一组关键词全部出现在首行中
    for group in keyword_groups:
        if all(kw in text_lower for kw in group):
            return True

    return False


async def get_compose_chunk_bbox_and_removed_data(data):
    """
    返回：
      - 符合条件的组成表分组（每组 [前前, 前, 当前]）
      - 从原始 data 中移除了所有相关记录后的新 data 列表
    """
    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 is_compose_table_by_first_tr(table_body):
                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 openai_chat(query: str, model="Qwen3-32B"):
#     """调用OpenAI API"""
#     try:
#         client = OpenAI(
#             api_key="none",
#             base_url="http://192.168.0.46:59800/v1"
#         )

#         response = client.chat.completions.create(
#             model=model,
#             messages=[
#                 {"role": "user", "content": query},
#             ],
#             temperature=0.1,
#             stream=False,
#             response_format={"type": "json_object"},
#             extra_body={"chat_template_kwargs": {"enable_thinking": False}}
#         )
#         return response.choices[0].message.content
#     except Exception as e:
#         print(f"调用OpenAI API时出错: {e}")
#         return None
    
async def extract_compose_table_entity(chunk):
    """
    提取组成表的设备、零部件实体以及关系
    """
    print("提取组成表实体和关系")

    # 建议：将 all_results 改为复数形式以符合习惯
    

    final_prompt = Prompt_compose.format(text=chunk)
    raw_response = await OpenaiAPI.openai_chat_aysnc(final_prompt,timeout=480)
    
    if raw_response is None:
        print(f"  ⚠️ 切片 API 调用失败")
    
    # 尝试解析 JSON
    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)

        # 打印美化版
    #     print(json.dumps(json_result, ensure_ascii=False, indent=2))
        
    except Exception as e:
        print(f"  ⚠️切片 JSON 解析失败: {e}")
        print(f"  原始响应: {repr(raw_response[:500])}")
        json_result = []


    # 返回按切片分组的结果列表
    return json_result


async def get_compose_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
    results, filtered_data =await get_compose_chunk_bbox_and_removed_data(data=data)
    print(f"抽取的组成表总数量为：{len(results)}")


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

        for i, compose in enumerate(results):
            try:
                # 获取内容和高亮信息
                print(f"开始处理第{i+1}个组成表")

                content, highlight_list = get_content_bbox(compose)
                print("组成表切片内容为：",content)
                print(len(content))
                # 提取实体
                result =await extract_compose_table_entity(content)
                if result is None:
                    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 = '/storage01/home/hdf/project/wxkgrag/kgextract/output/output1/122-06A0014-B01001_发动机-维修手册_content_list.json'
#     input_pdf_path = "/storage01/home/hdf/project/agent/测试.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_compose_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
#     }

#     # 保存为 JSON 文件
#     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}")

