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_anli
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/890DAA发动机装备故障案例_content_list.json'

# try:
#     # 打开并读取 JSON 文件
#     with open(file_path, 'r', encoding='utf-8') as file:
#         data = json.load(file)
    
# except Exception as e:
#     print(f"发生未知错误：{e}")


# for idx, item in enumerate(data, start=1):
#     item["id"] = idx
def extract_chunks_by_start_keywords(data, start_keywords):
    """
    通用版：根据多个起始关键词提取段落块。
    
    参数：
      - data: List[dict]，每项含 "id", "type", "text"（若 type=="text"）
      - start_keywords: List[str]，如 ["一、故障现象", "二、原因分析"]
    
    规则：
      - 起始点：type == "text" 且 text.strip() in start_keywords
      - 每个段落包含：
          a) 起始点的前一条记录（id-1，如果存在）
          b) 起始点本身
          c) 起始点之后的所有连续记录（任意 type），直到（但不包含）下一个起始关键词出现
      - 所有被选中的记录将从原始 data 中移除

    返回：
      - groups: List[List[record]]，每个 group 是一个完整段落（含前一条）
      - filtered_data: 剩余未被移除的记录列表
    """
    if not data or not start_keywords:
        return [], data

    # 构建 id -> record 映射
    record_map = {item["id"]: item for item in data}
    max_id = max(record_map.keys())

    # 第一步：找出所有起始关键词的位置（按 id 顺序）
    start_keyword_set = set(kw.strip() for kw in start_keywords)
    start_records = []

    for item in data:
        if item.get("type") == "text":
            text_content = item.get("text", "").strip()
            if text_content in start_keyword_set:
                start_records.append(item)

    if not start_records:
        return [], data

    # 按 id 排序（虽然 data 通常有序，但保险起见）
    start_records.sort(key=lambda x: x["id"])

    groups = []
    ids_to_remove = set()

    # 第二步：遍历每个起始点，构建段落
    for record in start_records:
        start_id = record["id"]
        group = []

        # 1. 添加前一条记录（id-1）
        prev_id = start_id - 1
        if prev_id >= 1 and prev_id in record_map:
            group.append(record_map[prev_id])
            ids_to_remove.add(prev_id)

        # 2. 添加当前起始记录
        group.append(record)
        ids_to_remove.add(start_id)

        # 3. 向后收集，直到遇到下一个起始关键词（不含）
        next_id = start_id + 1
        while next_id <= max_id:
            next_record = record_map.get(next_id)
            if next_record is None:
                break

            # 检查是否是任意一个起始关键词
            if (next_record.get("type") == "text" and 
                next_record.get("text", "").strip() in start_keyword_set):
                break  # 遇到下一个起始点，停止

            group.append(next_record)
            ids_to_remove.add(next_id)
            next_id += 1

        groups.append(group)

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

    return groups, 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_wxanli(chunk):
    """
    成功时返回 {'entities': [...], 'relationships': [...]}, 
    任何失败（超时、解析错、格式错）都返回 []。
    """
    final_prompt = Prompt_anli.replace("text", chunk)
    raw_response =await OpenaiAPI.openai_chat_aysnc(final_prompt,timeout=600)

    # 如果模型没返回内容（如超时），直接返回 []
    if raw_response is None:
        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)

        # 必须是 dict，且包含 entities 字段（且为 list）
        if not isinstance(json_result, dict):
            return []

        entities = json_result.get("entities", [])
        relationships = json_result.get("relationships", [])

        if not isinstance(entities, list) or not isinstance(relationships, list):
            return []

        # 返回标准结构
        return {
            "entities": entities,
            "relationships": relationships
        }

    except Exception:
        # 任何解析异常都返回 []
        return []

async def get_wxanli_node_relation(input_pdf_path,data,prefix_url = "http://192.168.0.46:59085"):

  filter = [
        "一、故障现象",
        "单位:",
        "故障名称:",
        "一、基本情况",
        "故障现象及原因:",
        "故障名称：",
        "实例故障"
    ]
  results, filtered_data = extract_chunks_by_start_keywords(data=data,start_keywords=filter)
  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[:200])

              result =await  extract_wxanli(content)

              # 关键判断：只有 result 是 dict 才继续
              if not isinstance(result, dict):
                  print(f"  ⚠️ 第 {i+1} 组：提取结果无效，跳过")
                  continue

              entities = result.get("entities", [])
              relationships = result.get("relationships", [])

              # 添加 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
            #       }]
            #       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
              # 安全追加
              all_entity_results.extend(entities)
              all_relation_results.extend(relationships)

          except Exception as e:
              print(f"  ❌ 处理第 {i+1} 组时发生未预期错误: {e}")
              # 不 re-raise，继续下一轮
              continue
  return all_entity_results,all_relation_results 
# if __name__ == "__main__":
#     file_path = '/app/files/890DAA发动机装备故障案例_content_list.json'
#     input_pdf_path = "/app/files/890DAA发动机装备故障案例.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_wxanli_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}个关系")

