import pandas as pd
from io import StringIO  # 1. 导入 StringIO
import json
import requests
import json
import os
from collections import Counter
from extract_util import safe_json_loads,highlight_pdf_file,get_content_bbox
import re
import asyncio
from bs4 import BeautifulSoup
import logging
# 配置 logger
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)
logger = logging.getLogger(__name__)
# ================= 函数定义 =================
def save_to_file(data, filepath, is_json=True):
    """通用保存函数"""
    try:
        with open(filepath, "w", encoding="utf-8") as f:
            if is_json:
                json.dump(data, f, ensure_ascii=False, indent=2)
            else:
                f.write(data)
        print(f"✅ 已保存: {filepath}")
        return True
    except Exception as e:
        print(f"❌ 保存失败 {filepath}: {e}")
        return False
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_fault_table_type(table_body: str):
    """
    返回故障表类型：
      'type1' -> 包含 [故障现象, 故障原因, 维修项目]
      'type2' -> 包含 [故障现象, 可能原因, 排除方法]
      None    -> 非故障表
    """
    if not isinstance(table_body, str):
        return None

    match = re.search(r'<tr[^>]*>(.*?)</tr>', table_body, re.IGNORECASE | re.DOTALL)
    if not match:
        return None

    text_lower = re.sub(r'<[^>]+>', '', match.group(1)).lower()

    if all(kw in text_lower for kw in ["故障现象", "故障原因", "维修项目"]):
        return "type1"
    elif all(kw in text_lower for kw in ["故障现象", "可能原因", "排除方法"]):
        return "type2"
    elif all(kw in text_lower for kw in ["故障外部现象", "故障可能原因", "操作员指导","排除方法"]):
        return "type3"
    return None


# ================= 修改：返回时携带 table_type =================
def get_guzhang_chunk_bbox_and_removed_data(data):
    results = []          # 元素格式: {"group": [...], "table_type": "type1"/"type2"}
    ids_to_remove = set()
    record_map = {item["id"]: item for item in data}

    for ins in data:
        if ins.get("type") == "table":
            table_type = get_fault_table_type(ins.get("table_body", ""))

            if table_type is not None:
                current_id = ins["id"]
                prev_ids = [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": group, "table_type": table_type})  # ← 携带类型

    filtered_data = [item for item in data if item["id"] not in ids_to_remove]
    return results, filtered_data

import re

def _expand_rowspans(content: str) -> list[list[str]]:
    """
    将含 rowspan 的 HTML 表格展开为完整二维列表。
    每个被合并的单元格都会被填充为原始内容。
    """
    all_rows = re.findall(r"<tr>(.*?)</tr>", content, flags=re.DOTALL)
    if not all_rows:
        return []

    matrix  = []
    pending = {}  # {col_index: (remaining_rows, cell_text)}

    for row_html in all_rows:
        cells    = re.findall(r"<td([^>]*)>(.*?)</td>", row_html, flags=re.DOTALL)
        row      = []
        col      = 0
        cell_idx = 0

        while cell_idx < len(cells) or any(c >= col for c in pending):
            if col in pending:
                # 该列被上方 rowspan 占用，直接填充
                remaining, text = pending[col]
                row.append(text)
                if remaining - 1 > 0:
                    pending[col] = (remaining - 1, text)
                else:
                    del pending[col]
                col += 1

            elif cell_idx < len(cells):
                attrs_str, text = cells[cell_idx]
                text = text.strip()

                rowspan_m = re.search(r'rowspan\s*=\s*["\']?(\d+)', attrs_str)
                rowspan   = int(rowspan_m.group(1)) if rowspan_m else 1

                row.append(text)
                if rowspan > 1:
                    pending[col] = (rowspan - 1, text)

                col      += 1
                cell_idx += 1
            else:
                break

        matrix.append(row)

    return matrix


def parse_table_content(content: str, chunk_size: int = 10) -> list[dict]:
    """
    先展开 rowspan，再按 chunk_size 切分，分批解析后汇总返回。
    """
    matrix = _expand_rowspans(content)
    if not matrix:
        print("⚠️ 未找到有效表格，返回空列表")
        return []

    header    = matrix[0]   # 表头
    data_rows = matrix[1:]  # 数据行（rowspan 已展开）
    total     = len(data_rows)

    print(f"📋 展开后共 {total} 条数据行，按每批 {chunk_size} 行切分处理")

    def _to_html(rows: list[list[str]]) -> str:
        def row_html(cells):
            return "<tr>" + "".join(f"<td>{c}</td>" for c in cells) + "</tr>"
        return "<table>" + row_html(header) + "".join(row_html(r) for r in rows) + "</table>"

    result_list = []

    for chunk_start in range(0, total, chunk_size):
        chunk      = data_rows[chunk_start: chunk_start + chunk_size]
        chunk_html = _to_html(chunk)

        try:
            chunk_df_list = pd.read_html(StringIO(chunk_html), header=0)
            if not chunk_df_list:
                print(f"  ⚠️ 切片 [{chunk_start}:{chunk_start+chunk_size}] 解析失败，跳过")
                continue

            chunk_result = json.loads(
                chunk_df_list[0].to_json(orient='records', force_ascii=False)
            )
            result_list.extend(chunk_result)
            print(f"  ✅ 切片 [{chunk_start}:{chunk_start+chunk_size}] 解析 {len(chunk_result)} 条")

        except Exception as e:
            print(f"  ❌ 切片 [{chunk_start}:{chunk_start+chunk_size}] 处理异常: {e}")
            continue

    print(f"共汇总 {len(result_list)} 条记录")
    return result_list
# ================= 修改：主函数按类型分支处理 =================
async def get_guzhang_node_relation(input_pdf_path, data, prefix_url="http://192.168.0.46:59085", device="发动机"):
    for idx, item in enumerate(data, start=1):
        item["id"] = idx

    guzhang_groups, cleaned_data = get_guzhang_chunk_bbox_and_removed_data(data)
    print(f"抽取的故障表总数量为：{len(guzhang_groups)}")

    if not guzhang_groups:
        print("未找到符合条件的组成表格，跳过处理。")
        all_entity_results, all_relation_results = [], []

    else:
        all_entity_results, all_relation_results = [], []

        for i, item in enumerate(guzhang_groups):
            compose     = item["group"]
            table_type  = item["table_type"]

            try:
                print(f"开始处理第 {i+1} 个故障表，类型: {table_type}")
                content, highlight_list = get_content_bbox(compose)
                print("故障表切片内容为：", content[:200])

                result_list = parse_table_content(content, chunk_size=10)
                if not result_list:
                    continue

                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}"
                logger.info(f"处理完成，高亮文件：{highlight_filename}，高亮页码：{most_common_page}")

                # ---------- type1: 故障现象 / 故障原因 / 维修项目 ----------
                if table_type == "type1":
                    for ins in result_list:
                        # 安全获取字段，如果是 None 则转为空字符串
                        guzhang_xianxiang = ins.get("故障现象") or ""
                        guzhang_yuanyin = ins.get("故障原因") or ""
                        weixiu_xiangmu = ins.get("维修项目") or ""
                        
                        # 只有当现象或原因不为空时才构建故障名，避免生成 "导致" 这种无意义字符串
                        if guzhang_xianxiang or guzhang_yuanyin:
                            guzhangname = f"{guzhang_yuanyin}导致{guzhang_xianxiang}"
                        else:
                            continue # 如果关键信息缺失，跳过该行

                        source_data_list = [{"filename": filename, "info": content, "url": url, "bbox": highlight_list}]

                        all_entity_results.append({
                            "type": "故障模式",
                            "properties": {
                                "名称":     guzhangname,
                                "故障现象": guzhang_xianxiang,
                                "故障原因": guzhang_yuanyin,
                                "维修方案": weixiu_xiangmu,
                                "knowledge_source": json.dumps(source_data_list, ensure_ascii=False)
                            }
                        })
                        all_relation_results.append({
                            "type":        "包含故障",
                            "from_entity": f"{device}的维修工作",
                            "to_entity":   guzhangname
                        })

                # ---------- type2: 故障现象 / 可能原因 / 排除方法 ----------
                elif table_type == "type2":
                    for ins in result_list:
                        guzhang_xianxiang = ins.get("故障现象") or ""
                        guzhang_yuanyin = ins.get("可能原因") or ""
                        weixiu_fangfa = ins.get("排除方法") or ""

                        if guzhang_xianxiang or guzhang_yuanyin:
                            guzhangname = f"{guzhang_yuanyin}导致{guzhang_xianxiang}"
                        else:
                            continue

                        source_data_list = [{"filename": filename, "info": content, "url": url, "bbox": highlight_list}]

                        all_entity_results.append({
                            "type": "故障模式",
                            "properties": {
                                "名称":     guzhangname,
                                "故障现象": guzhang_xianxiang,
                                "故障原因": guzhang_yuanyin,
                                "维修方案": weixiu_fangfa,
                                "knowledge_source": json.dumps(source_data_list, ensure_ascii=False)
                            }
                        })
                        all_relation_results.append({
                            "type":        "包含故障",
                            "from_entity": f"{device}的维修工作",
                            "to_entity":   guzhangname
                        })

                elif table_type == "type3":
                    for ins in result_list:
                        guzhang_xianxiang = ins.get("故障外部现象") or ""
                        guzhang_yuanyin = ins.get("故障可能原因") or ""
                        caozuoyuan_zhidao = ins.get("操作员指导") or ""
                        weixiu_fangfa = ins.get("排除方法") or ""

                        if guzhang_xianxiang or guzhang_yuanyin:
                            guzhangname = f"{guzhang_yuanyin}导致{guzhang_xianxiang}"
                        else:
                            continue

                        source_data_list = [{"filename": filename, "info": content, "url": url, "bbox": highlight_list}]

                        all_entity_results.append({
                            "type": "故障模式",
                            "properties": {
                                "名称":     guzhangname,
                                "故障现象": guzhang_xianxiang,
                                "故障原因": guzhang_yuanyin,
                                "操作员指导": caozuoyuan_zhidao,
                                "维修方案": weixiu_fangfa,
                                "knowledge_source": json.dumps(source_data_list, ensure_ascii=False)
                            }
                        })
                        all_relation_results.append({
                            "type":        "包含故障",
                            "from_entity": f"{device}的维修工作",
                            "to_entity":   guzhangname
                        })
            except Exception as e:
                print(f"处理第 {i+1} 组时发生未预期错误: {e}")
                continue

        # 设备 → 维修工作指导 节点和关系（两种类型共用）
        all_entity_results.append({
            "type": "维修工作",
            "properties": {"名称": f"{device}的维修工作"}
        })
        all_relation_results.append({
            "type":        "发生故障使用维修工作",
            "from_entity": device,
            "to_entity":   f"{device}的维修工作"
        })

    return all_entity_results, all_relation_results



"""
实体：设备、维修工作、故障、维修项目
关系：设备---发生故障使用维修工作--->维修工作
维修工作---表征--->故障
故障---被修复--->维修项目
"""
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/122-06A0014-B01001_发动机-维修手册.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_guzhang_node_relation(input_pdf_path, data=data)
    ) 
    for ins in all_entity_results:
        print(ins)
    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}")