import json
import fitz  # PyMuPDF
import os
from collections import Counter
import uuid
# 假设 Prompt 和 extract_util 是你本地已有的模块
from kg_build.Prompt import prompt_extract_operation_node_rela
from kg_build.extract_util import safe_json_loads, highlight_pdf_file, get_content_bbox, get_most_common_page, highlight_pdf_file_batch
import re
import asyncio
from openai import AsyncOpenAI
# from modelsAPI.model_api import OpenaiAPI


async def openai_chat_aysnc(query: str, timeout: int = 180) -> str | None:
    """异步调用 OpenAI 兼容 API（如本地部署的 Qwen、vLLM、Ollama 等）"""
    client = None
    try:
        base_url = "http://192.168.0.46:59800/v1"
        api_key = "none"

        client = AsyncOpenAI(
            api_key=api_key,
            base_url=base_url,
            timeout=timeout
        )

        response = await client.chat.completions.create(
            model="46-qwen3.6-35B",
            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
    finally:
        if client:
            await client.close()

def is_operation_table_by_first_tr(table_body: str) -> bool:
    """
    判断 table_body 是否为操作项目表，依据是第一行 <tr> 中是否包含 '项目编号' 和 '操作项目'
    """
    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_clean = text_in_first_tr.strip()
    
    # 检查是否同时包含两个关键词
    has_code = "项目编号" in text_clean
    has_item = "操作项目" in text_clean
    return has_code and has_item


def extract_operation_groups_and_remaining(records, 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:
        is_start = (
            record.get("type") == "table" 
            and is_operation_table_by_first_tr(record.get("table_body", ""))
        )
        
        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:
                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]
            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



    
async def extract_entity_relation(chunk):
    """
    批量处理一个文档的多个切片（chunks），每个切片调用大模型结构化，并返回全部结果。
    """
    final_prompt = prompt_extract_operation_node_rela.replace("text", chunk)

    raw_response = await openai_chat_aysnc(final_prompt, timeout=600)

    if raw_response is None:
        print(f"  ⚠️ 切片 API 调用失败")
        return []
    
    # 尝试解析 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)
        return json_result
        
    except Exception as e:
        print(f"  ⚠️ 切片 JSON 解析失败: {e}")
        print(f"  原始响应: {repr(raw_response[:500])}")
        return []



async def get_operation_node_relation(input_pdf_path, data, prefix_url="http://192.168.0.46:59085", device="发动机"):
    # 1. 初始化数据 ID (如果输入数据还没有ID)
    for idx, item in enumerate(data, start=1):
        if "id" not in item:
            item["id"] = idx
            
    devicename = device
    
    # 2. 数据预处理：获取故障组 and 清洗数据
    groups1, remaining_records1 = extract_operation_groups_and_remaining(data)
    groups = groups1[1:] if len(groups1) > 1 else []

    all_entity_results = []
    all_relation_results = []
    device_operation = {
        "type" : "操作程序",
        "properties": {
            "名称": f"{devicename}的操作程序",
      }
    }
    
    _filename = os.path.basename(input_pdf_path)
    # _base, _ext = os.path.splitext(_filename)
    # highlight_filename = f"{_base}_{uuid.uuid4()}{_ext}"
    filename = _filename
    # all_highlights = []

    if not groups:
        print("未找到符合条件的操作项目组，跳过后续处理。")
    else:
        # 3. 遍历每个维修项目组进行抽取
        for i, repair_group in enumerate(groups):
            try:            
                print(f"开始处理第 {i+1} 个操作项目表")
                
                # 获取内容和高亮区域
                content, highlight_list = get_content_bbox(repair_group)
                if not content:
                    print(f"第 {i+1} 组：内容为空，跳过")
                    continue
                
                # ================= 重试逻辑开始 =================
                result = None
                max_retries = 20      
                retry_delay = 0.5       
                highlight_filename, most_common_page,filename = highlight_pdf_file(input_pdf_path,highlight_list)
                url = f"http://192.168.0.46:59085/upload/{highlight_filename}#page={most_common_page}"
                for attempt in range(1, max_retries + 2): 
                    # print(content[:500]) # 调试时可打开，避免日志过多
                    # 1. 调用 API
                    result = await extract_entity_relation(content)
                    
                    # 2. 验证结果有效性
                    is_valid = False
                    if result is not None and isinstance(result, dict):
                        # 情况 1: 标准格式 {'entities': [...], 'relationships': [...]}
                        if "entities" in result and isinstance(result["entities"], list) and len(result["entities"]) > 0:
                            is_valid = True
                        # 情况 2: 兼容日志中的格式 {'result': [...]} 
                        elif "result" in result and isinstance(result["result"], list) and len(result["result"]) > 0:
                            is_valid = True
                    
                    # 3. 判断是否成功
                    if is_valid:
                        if attempt > 1:
                            print(f"  ✅ 第 {attempt} 次尝试成功获取有效数据。")
                        break 
                    
                    # 4. 如果失败且还有重试机会
                    if attempt <= max_retries:
                        reason = "API 返回为空" if result is None else "返回数据格式异常或内容为空"
                        print(f"  ⚠️ 第 {attempt} 次尝试失败：{reason}。等待 {retry_delay} 秒后重试...")
                        await asyncio.sleep(retry_delay)
                    else:
                        print(f"  ❌ 经过 {max_retries + 1} 次尝试仍无法获取有效数据，跳过该组。")
                        result = None 
                        break
                # ================= 重试逻辑结束 =================

                if result is None:
                    continue
                
                # 构造知识来源字符串 (此处保持原有逻辑，实际项目中可能需要动态生成真实 URL)
                source_data_list = [{
                    "filename": filename,
                    "info": content,
                    "url": url,
                    "bbox": highlight_list
                }]
                knowledge_source_str = json.dumps(source_data_list, ensure_ascii=False)
                
                # 兼容不同返回格式
                entities = result.get("entities", result.get("result", []))
                relationships = result.get("relationships", [])
                
                for entity in entities:
                    if "properties" not in entity or not isinstance(entity["properties"], dict):
                        entity["properties"] = {}
                    entity["properties"]["knowledge_source"] = knowledge_source_str
                    if entity.get("type") == "操作项目":
                        entity["properties"]["内容"] = content
                    # 构建基础关系：设备 -> 操作项目
                    rela = {
                        "type": "操作时使用",
                        "from_entity": f"{devicename}的操作程序",
                        "to_entity": entity.get("properties", {}).get("名称", "未知实体")
                    }
                    all_relation_results.append(rela)
                
                all_entity_results.extend(entities)
                all_relation_results.extend(relationships)

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



    # ================= 【新增】实体和关系去重逻辑 =================
    print(f"去重前统计：实体数量={len(all_entity_results)}, 关系数量={len(all_relation_results)}")

    # 1. 实体去重
    unique_entities = []
    seen_entity_keys = set()
    
    for entity in all_entity_results:
        props = entity.get("properties", {})
        # 定义唯一键：优先使用 "编号"，如果没有则使用 "名称"
        # 这样可以防止 OT001 出现两次
        entity_id = props.get("编号") or props.get("名称")
        
        # 如果连名称都没有，生成一个临时ID避免报错，但这种情况通常视为无效数据
        if not entity_id:
            entity_id = str(uuid.uuid4())
        
        if entity_id not in seen_entity_keys:
            seen_entity_keys.add(entity_id)
            unique_entities.append(entity)
        else:
            # 可选：打印跳过的重复项
            # print(f"  ⚠️ 跳过重复实体: {entity_id} - {props.get('名称')}")
            pass

    all_entity_results = unique_entities

    # 2. 关系去重
    unique_relations = []
    seen_relation_keys = set()

    for rela in all_relation_results:
        # 定义关系唯一键：类型 + 源 + 目标
        r_key = (rela.get("type"), rela.get("from_entity"), rela.get("to_entity"))
        
        if r_key not in seen_relation_keys:
            seen_relation_keys.add(r_key)
            unique_relations.append(rela)
    
    all_relation_results = unique_relations

    relation = {'type': '使用', 'from_entity': f'{device}的操作使用', 'to_entity': f'{device}的操作程序'}
    all_relation_results.append(relation)
    all_entity_results.append(device_operation)

    print(f"去重后统计：实体数量={len(all_entity_results)}, 关系数量={len(all_relation_results)}")
    # ========================================================


    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/wxkgrag/kgextract/output/output1/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 = asyncio.run(
        get_operation_node_relation(input_pdf_path, data=data, device="发动机")
    ) 
    
    length_entity = len(all_entity_results)
    length_relation = len(all_relation_results)
    for ins in all_entity_results:
        print(ins)
        print(11111111111111111111111)
    for ins in all_relation_results:
        print(ins)
        print(22222222222222222222)
    print(f"最终抽取结果：{length_entity}个实体，{length_relation}个关系")