import json
import fitz  # PyMuPDF
import os
from collections import Counter
import uuid
from kg_build.Prompt import prompt_extract_repair_node_rela
from openai import OpenAI
import re
from kg_build.extract_util import safe_json_loads, highlight_pdf_file, get_content_bbox
from modelsAPI.model_api import OpenaiAPI
import re
from typing import List, Dict, Any, Tuple
import json
import re
import copy
import asyncio
from openai import AsyncOpenAI
import pandas as pd
from io import StringIO  # 1. 导入 StringIO
def optimize_component_type(data):
    """
    根据组成编码的位数优化记录的type字段
    
    规则：
    - 2位数 → 子系统
    - 4位数 → 部件
    - 6位数 → 一级零部件
    - 8位数 → 二级零部件
    """
    
    # type映射关系
    type_mapping = {
        2: '子系统',
        4: '部件',
        6: '一级零部件',
        8: '二级零部件'
    }
    
    optimized_data = []
    stats = {
        'total': 0,
        'modified': 0,
        'unchanged': 0,
        'no_code': 0
    }
    
    for record in data:
        stats['total'] += 1
        new_record = copy.deepcopy(record)
        
        # 检查是否存在properties和组成编码
        if 'properties' in record and '组成编码' in record['properties']:
            code = str(record['properties']['组成编码']).strip()
            
            # 提取纯数字部分（去除可能的特殊字符）
            code_digits = re.sub(r'[^\d]', '', code)
            code_length = len(code_digits)
            
            # 根据位数判断type
            if code_length in type_mapping:
                expected_type = type_mapping[code_length]
                current_type = record.get('type', '')
                
                if current_type != expected_type:
                    new_record['type'] = expected_type
                    stats['modified'] += 1
                    print(f"【修正】{record['properties'].get('名称', '未知')}：{current_type} → {expected_type} (编码：{code})")
                else:
                    stats['unchanged'] += 1
            else:
                print(f"【警告】{record['properties'].get('名称', '未知')}：编码位数{code_length}不在映射范围内 (编码：{code})")
                stats['unchanged'] += 1
        else:
            stats['no_code'] += 1
            print(f"【跳过】{record.get('type', '未知')} - {record.get('properties', {}).get('名称', '未知')}：无组成编码字段")
        
        optimized_data.append(new_record)
    
    
    return optimized_data




def is_fault_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)
    
    # 提取所有 <td> 或 <th> 中的文本（简单去标签）
    # 更健壮的做法是递归去除标签，但这里用正则足够
    text_in_first_tr = re.sub(r'<[^>]+>', '', first_tr_html)  # 去除所有 HTML 标签
    
    # 统一转为字符串并检查关键词
    text_lower = text_in_first_tr.lower()
    keywords = ["故障现象", "故障原因", "维修项目"]
    return all(kw in text_lower for kw in keywords)


def get_guzhang_chunk_bbox_and_removed_data(data):
    title_keywords = {"故障"}  # 可选：仍保留标题中有“故障”的约束
    results = []
    ids_to_remove = set()
    record_map = {item["id"]: item for item in data}

    for ins in data:
        if ins.get("type") == "table":
            table_body = ins.get("table_body", "")
            
            # 核心判断：是否包含三个关键字段
            if is_fault_table_by_first_tr(table_body):
                current_id = ins["id"]
                
                # 可选：向上找最近的标题（更智能），但若坚持用前3条也可保留
                # 这里先按你的原逻辑：取前3条 + 当前
                prev_ids = [current_id - 3, 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

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

            client = AsyncOpenAI(
                api_key=api_key,
                base_url=base_url,
                timeout=timeout  # 整体请求超时（秒）
            )

            response = await client.chat.completions.create(
                model="Qwen3.5-35B-A3B",
                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:
            # 可选：显式关闭客户端（通常不需要，但可避免警告）
            await client.close()

async def extract_wxindextable(chunk):
    final_prompt = prompt_extract_repair_node_rela.replace("text",chunk)
    raw_response =await OpenaiAPI.openai_chat_aysnc(final_prompt,timeout=420)
    
    if raw_response is None:
        print("⚠️ 切片 API 调用失败")
        return None

    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 None


def extract_text_from_first_tr(html: str) -> str:
    """提取 HTML 表格中第一个 <tr> 内的文本（去除标签）"""
    # if not isinstance(html, str):
    #     return ""
    # match = re.search(r'<tr[^>]*>(.*?)</tr>', html, re.IGNORECASE | re.DOTALL)
    # if not match:
    #     return ""
    # first_tr_content = match.group(1)
    # 去除所有 HTML 标签
    text = re.sub(r'<[^>]+>', '', html)
    return text.strip()


def matches_any_pattern(text: str, patterns: List[List[str]]) -> bool:
    """检查文本是否包含任意一组关键词（全部命中）"""
    if not text:
        return False
    text_lower = text.lower()
    for pattern in patterns:
        if all(keyword in text_lower for keyword in pattern):
            return True
    return False


def is_maintenance_table(table_body: str, patterns: List[List[str]]) -> bool:
    """判断是否为维修项目表格：仅检查第一个 <tr> 是否匹配任一模式"""
    first_tr_text = extract_text_from_first_tr(table_body)
    return matches_any_pattern(first_tr_text, patterns)


# ==============================
# 维修项目表格的关键词模式（可按需扩展）
# ==============================

MAINTENANCE_TABLE_PATTERNS: List[List[str]] = [
    ["维修项目编号", "组成编码", "名称", "维修项目", "维修间隔期"],
    ["故障现象", "编码", "维修内容", "周期"],
    ["序号", "维修项", "标准", "频次"],  # 示例：可继续添加业务变体
    ["设备名称", "维修级别", "维修项目编号", "维修间隔期"],  # 示例：可继续添加业务变体
    ["设备名称", "设备型号"],  # 示例：可继续添加业务变体
    ["维修项", "组成编码", "名称", "维修项目", "维修间隔期"],  # 示例：可继续添加业务变体

]


# ==============================
# 主函数：提取维修项目表格及其后续内容组
# ==============================

def extract_maintenance_groups_and_remaining(
    records: List[Dict[str, Any]],
    patterns: List[List[str]] = MAINTENANCE_TABLE_PATTERNS,
    max_following_for_last: int | None = None
) -> Tuple[List[List[Dict]], List[Dict]]:
    groups = []
    current_group = None
    used_ids = set()
    
    # 标记是否已经遇到第一个维修表格（用于跳过）
    skipped_first = False

    for record in records:
        is_start = (
            record.get("type") == "table" and
            is_maintenance_table(record.get("table_body", ""), patterns)
        )

        if is_start:
            # 如果是第一个维修表，跳过（不加入 groups，也不标记 used_ids）
            if not skipped_first:
                skipped_first = True
                # 注意：这里不清空 current_group，因为前面不可能有 group（第一个 start）
                current_group = None  # 确保不会把之前的非 start 内容误加
                continue  # 跳过这个 record，不加入任何 group

            # 从第二个维修表开始，才正常分组
            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)

    # 处理最后一个 group（仅当有有效 group 时）
    if current_group is not None:
        groups.append(current_group)
        used_ids.update(r["id"] for r in current_group)

    # 限制最后一个 group 长度
    if groups and max_following_for_last is not None:
        last_group = groups[-1]
        max_len = max_following_for_last + 1
        if len(last_group) > max_len:
            truncated = last_group[:max_len]
            removed_ids = {r["id"] for r in last_group[max_len:]}
            used_ids -= removed_ids
            groups[-1] = truncated

    # remaining_records 包含所有未被 used_ids 标记的记录
    # 由于第一个 group 没有加入 used_ids，所以会保留在 remaining_records 中
    remaining_records = [r for r in records if r["id"] not in used_ids]
    return groups, remaining_records


FAULT_TABLE_KEYWORDS_wxsm = ["故障现象", "故障原因", "维修人数", "维修工具", "排除步骤"]

def is_fault_table_wxsm(table_body: str) -> bool:
    if not isinstance(table_body, str):
        return False
    return all(kw in table_body for kw in FAULT_TABLE_KEYWORDS_wxsm)


def extract_fault_groups_and_remaining_wxsm(
    records: List[Dict[str, Any]],
) -> Tuple[List[List[Dict]], List[Dict]]:
    groups = []
    current_group = None
    used_ids = set()

    for record in records:
        is_start = (
            record.get("type") == "table" and
            is_fault_table_wxsm(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)

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

def format_compose_code(val):
    """将 int/float 型组成编码转换为偶数位字符串，不足偶数位则前补 0"""
    if pd.isna(val):
        return None
    # 统一处理 int 和 float
    if isinstance(val, (int, float)):
        s = str(int(val))           # 101 / 101.0 → '101'
        if len(s) % 2 != 0:        # 奇数位 → 前补 0
            s = '0' + s             # '101' → '0101'
        return s
    return str(val)   
              # 已经是字符串则直接返回


async def get_repair_node_relation(input_pdf_path, data, prefix_url="http://192.168.0.111:9085"):
    # 1. 初始化数据 ID
    for idx, item in enumerate(data, start=1):
        item["id"] = idx

    # 2. 数据预处理：获取故障组 and 清洗数据
    guzhang_groups, cleaned_data = get_guzhang_chunk_bbox_and_removed_data(data)
    groups_wxxm, remaining_records = extract_maintenance_groups_and_remaining(cleaned_data)
    groups_wxsm, remaining_records_wxsm =extract_fault_groups_and_remaining_wxsm(remaining_records)  # 可选：先抽取故障表，进一步清洗数据
    print(f"维修项目表格总数量: {len(groups_wxxm)}")
    print(f"维修说明表格总数量: {len(groups_wxsm)}")
    groups = groups_wxxm + groups_wxsm  # 合并维修项目表和维修说明表的结果，后续一起处理
    all_entity_results = []
    all_relation_results = []
    
    final_filename = None
    final_url = None

    if not groups_wxxm:
        print("未找到符合条件的维修项目组，跳过后续处理。")
    else:
        for i, repair_group in enumerate(groups_wxxm):
            try:
                
                print(f"开始处理第 {i+1} 个维修项目表")
                # 获取内容和高亮区域
                content, highlight_list = get_content_bbox(repair_group)
                if not content:
                    print(f"第 {i+1} 组：内容为空，跳过")
                    continue
                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)

                wx_bianhao = result_list[0].get("维修项目编号")
                wx_mingcheng = result_list[0].get("维修项目")
                device_bianma = result_list[0].get("组成编码")
                device_name = result_list[0].get("名称")
                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 not wx_bianhao or not wx_mingcheng or not device_bianma or not device_name:
                    print("html解析不成功，需要大模型進行抽取")
                    result = await extract_wxindextable(content)
                    if result is None:
                        continue
                    entities = result.get("entities", [])
                    
                    if not isinstance(entities, list):
                        print(f"第 {i+1} 组：entities 字段格式异常，跳过")
                        continue
                    
                    relationships = result.get("relationships", [])
                    # 更新全局文件信息 (用于最后创建手册节点)

                    for entity in entities:
                        if entity['type'] == '维修项目':
                            entity["properties"]["内容"] = content
                    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)
                else:
                    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)
                    wxentity = {
                        "type" : "维修项目",
                        "properties":{
                            "名称" : wx_mingcheng,
                            "编号" : wx_bianhao,
                            "内容" : content,
                            "knowledge_source" :knowledge_source_str
                        }
                    }
                    deviceentity ={
                        "type" : "设备",
                        "properties":{
                            "名称" : device_name,
                            "组成编码" : device_bianma,
                            "维修项目名称" : wx_mingcheng,
                            "维修项目编号" : wx_bianhao,
                            "knowledge_source" :knowledge_source_str
                        }
                    }
                    all_entity_results.append(wxentity)
                    all_entity_results.append(deviceentity)

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

    if not groups_wxsm:
        print("未找到符合条件的维修说明表，跳过后续处理。")
    else:
        for i, repair_group in enumerate(groups_wxsm):
            try:
                
                print(f"开始处理第 {i+1} 个维修说明表")
                # 获取内容和高亮区域
                content, highlight_list = get_content_bbox(repair_group)
                if not content:
                    print(f"第 {i+1} 组：内容为空，跳过")
                    continue
                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)
                
                # 转换为 JSON 列表
                json_data = df.to_json(orient='records', force_ascii=False)
                result_list = json.loads(json_data)
                wx_bianhao = result_list[0].get("维修项目编号")
                wx_mingcheng = result_list[0].get("维修项目")
                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 not wx_bianhao or not wx_mingcheng:
                    result = await extract_wxindextable(content)
                    if result is None:
                        continue
                    entities = result.get("entities", [])
                    
                    if not isinstance(entities, list):
                        print(f"第 {i+1} 组：entities 字段格式异常，跳过")
                        continue
                    
                    relationships = result.get("relationships", [])
                    # 更新全局文件信息 (用于最后创建手册节点)

                    for entity in entities:
                        if entity['type'] == '维修项目':
                            entity["properties"]["内容"] = 'content'
                    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)
                else:
                    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)
                    wxentity = {
                        "type" : "维修项目",
                        "properties":{
                            "名称" : wx_mingcheng,
                            "编号" : wx_bianhao,
                            "内容" : content,
                            "knowledge_source" :knowledge_source_str
                        }
                    }
                    all_entity_results.append(wxentity)


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


    # 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 = await extract_wxindextable(content)
    #             # 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 is None:
    #                 continue
    #             entities = result.get("entities", [])
                
    #             if not isinstance(entities, list):
    #                 print(f"第 {i+1} 组：entities 字段格式异常，跳过")
    #                 continue
                
    #             relationships = result.get("relationships", [])
    #             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:
    #                 if entity['type'] == '维修项目':
    #                     entity["properties"]["内容"] = content
    #             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}")
    #             import traceback
    #             traceback.print_exc()
    #             continue
    # 5. 最终优化
    # all_entity_results = optimize_component_type(all_entity_results)

    return all_entity_results, all_relation_results

"""
实体：设备、子设备、维修项目
关系：设备-维修时使用-维修项目
"""

if __name__ == "__main__":
    file_path = '/app/wxceshi/files/雷达/output/101-0440011-B04001_HLJP-349型跟踪雷达维修说明书-秘密_content_list.json'
    input_pdf_path = "/app/wxceshi/files/雷达/101-0440011-B04001_HLJP-349型跟踪雷达维修说明书-秘密.pdf"
    output_json_path = '/app/wxceshi/files/雷达/output/维修细目表抽取结果.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_repair_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 all_entity_results:
        print(ins)
        print(11111111111111111)

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



