"""
extract_filtertext_serial.py
对外接口：get_filtertext_node_relation(input_pdf_path, data, prefix_url, device)

全程不调用大模型，纯规则处理。

逻辑分支：
  - 文件名含 备品备件/图解目录            → 直接跳过，不做任何解析
  - 文件名含 接线图/原理图/总图/连接图     → 直接生成"图录"实体 + 关系
  - 文件名含 芯线表                       → 直接生成"芯线表"实体 + 关系
  - 文件名含 随机文件目录/随机文件          → 直接生成"文件目录"实体 + 关系
  - 其他文件（如维修手册等）               → 按章节结构提取实体 + 关系

关系统一格式：
  {"type": "包含", "from_entity": device, "to_entity": 实体type}

章节切分逻辑已迁移至 chapter_segmenter.py，本文件直接调用。
"""

import json
import os
import re
from typing import List
from kg_build.extract_util import highlight_pdf_file
from entity_registry import register as _register_entity_type
from kg_build.chapter_segmenter import segment_chapters, Chapter


# ==============================
# 配置
# ==============================

DEVICE_NAME = "发动机"

# 一级标题映射：去编号后的标题文本 → 替换为指定名称（精确匹配）
L1_TITLE_MAPPING = {
    "概述": "基本情况",
    "使用操作": "操作使用",
    "系统结构组成" : "组成",
    "设备组成" : "组成",
    "设备功能" : "功能",
    "安装与调试" : "调试",
    "调试细则" : "调试",
    "系统功能" : "功能",
    "系统组成" : "组成",
    "维修工作指导" : "维修工作"
}

# 二级标题（特殊章节独立实体）类型映射
L2_TYPE_MAPPING = {
    "技术指标": "技术指标参数",
    "工作条件": "环境条件",
    "组成简介" : "组成",
    "主要用途及适用范围" :"功能",
    "主要功能" : "功能",
    "接口关系" : "接口",
    "主要性能参数" : "性能",
}

# 实体 type → 关系 type（未命中默认"包含"）
ENTITY_RELATION_MAPPING = {
    "基本情况":     "实现",
    "功能":         "实现",
    "环境条件":     "适用",
    "组成":         "包括",
    "技术指标参数": "具有",
    "接口":         "配备",
    "安全警告":     "安全警告",
    "操作使用":     "支持操作",
    "维修工作":     "维修",
    "调试":         "调试",
    "器材保障":     "依赖",
    "图册":         "配套",
    "设计单位":     "设计",
    "性能"  :  "具有",
}

# 设计单位识别关键词（出现在封面 page_idx=0 的文本中）
DESIGN_UNIT_KEYWORDS = ["研究所", "集团公司", "有限公司", "工厂", "研究院", "研究中心"]

# ---------- 文件名关键词配置 ----------
DIAGRAM_KEYWORDS = ["接线图", "原理图", "总图", "连接图"]
TABLE_KEYWORDS = ["芯线表"]
CATALOG_KEYWORDS = ["随机文件目录", "随机文件"]
SPARE_KEYWORDS = ["备品备件"]
SKIP_KEYWORDS = ["图解目录"]


# ==============================
# 文件名判断
# ==============================

def _should_skip(filename: str) -> bool:
    """文件名包含跳过关键词时返回True，直接不处理。"""
    basename = os.path.basename(filename)
    return any(kw in basename for kw in SKIP_KEYWORDS)


def _detect_nomodel_category(filename: str):
    """
    根据文件名判断属于哪种特殊类别。
    返回: (category, keyword) 或 (None, None) 表示走通用章节提取。
    """
    basename = os.path.basename(filename)
    for kw in DIAGRAM_KEYWORDS:
        if kw in basename:
            return "diagram", kw
    for kw in TABLE_KEYWORDS:
        if kw in basename:
            return "table", kw
    for kw in CATALOG_KEYWORDS:
        if kw in basename:
            return "catalog", kw
    for kw in SPARE_KEYWORDS:
        if kw in basename:
            return "spare", kw
    return None, None


# ==============================
# 域映射 / 关系类型
# ==============================

def _map_l1_title(title_clean: str) -> str:
    """对一级标题做映射，如 '概述' → '基本情况'，未命中则原样返回。"""
    return L1_TITLE_MAPPING.get(title_clean, title_clean)


def _get_relation_type(entity_type: str) -> str:
    """根据实体 type 返回对应的关系 type，未命中默认'包含'。"""
    return ENTITY_RELATION_MAPPING.get(entity_type, "包含")


# ==============================
# 设计单位提取
# ==============================

def _extract_design_unit(data: list) -> str:
    """
    从封面前两页（page_idx=0 或 1）提取设计单位名称。
    规则：文本含机构关键词（不限 type，兼容 discarded 记录）。
    返回第一个命中的文本，找不到返回空字符串。
    """
    for ins in data:
        if ins.get("page_idx", -1) not in (0, 1):
            continue
        text = ins.get("text", "").strip()
        if text and any(kw in text for kw in DESIGN_UNIT_KEYWORDS):
            return text
    return ""


# ==============================
# 单条 record 转文本
# ==============================

def _record_to_text(ins: dict) -> str:
    """将单条 record 转为纯文本字符串。"""
    if ins.get("type") in ("text", "discarded"):
        return ins.get("text", "").strip()

    elif ins.get("type") == "table":
        table_caption = ins.get("table_caption", "")
        table_body = ins.get("table_body", "")
        if isinstance(table_caption, list):
            table_caption = "\n".join(str(c) for c in table_caption)
        table_text = re.sub(r'<[^>]+>', ' ', table_body or "")
        table_text = re.sub(r'\s+', ' ', table_text).strip()
        parts = [str(table_caption), table_text]
        return "\n".join(p for p in parts if p).strip()

    elif ins.get("type") == "image":
        image_caption = ins.get("image_caption", "")
        img_path = ins.get("img_path", "")
        if isinstance(image_caption, list):
            image_caption = "\n".join(str(i) for i in image_caption)
        elif not isinstance(image_caption, str):
            image_caption = str(image_caption)
        if isinstance(img_path, list):
            img_path = "\n".join(str(p) for p in img_path)
        elif not isinstance(img_path, str):
            img_path = str(img_path)
        parts = [image_caption, img_path]
        return "\n".join(p for p in parts if p).strip()

    return ""


# ==============================
# 高亮相关
# ==============================

def _collect_highlights(records: list) -> list:
    """从 records 中收集高亮信息。"""
    highlight_list = []
    for ins in records:
        bbox = ins.get("bbox", [])
        page_idx = ins.get("page_idx", -1)
        if bbox and page_idx >= 0:
            highlight_list.append({"page_idx": page_idx, "bbox": bbox})
    return highlight_list



def _build_highlight_url(input_pdf_path, highlight_list, prefix_url, include_page=True):
    """调用 highlight_pdf_file 生成高亮PDF并返回URL。"""
    if not highlight_list:
        return ""
    try:
        highlight_filename, most_common_page, filename = highlight_pdf_file(
            input_pdf_path, highlight_list
        )
        if include_page:
            return f"{prefix_url}/upload/{highlight_filename}#page={most_common_page}"
        return f"{prefix_url}/upload/{highlight_filename}"
    except Exception as e:
        print(f"  ⚠️ highlight_pdf_file 失败: {e}")
        return ""


# ==============================
# 图/表/目录类：直接生成实体 + 关系
# ==============================

def _special_file_generate_entities(input_pdf_path, prefix_url, device, category, keyword):
    """
    对于图/表/目录类文件，直接根据类别生成实体和关系。
    返回: (entities, relations)
    """
    filename = os.path.basename(input_pdf_path)
    #url = f"{prefix_url}/upload/{filename}"
   # url = f"{prefix_url}/files/uploads_v1/{filename}"
    url = f"{prefix_url}/files/{filename}"
    entities = []
    relations = []

    if category == "diagram":
        entity_type = f"{device}的{keyword}"
        entity = {
            "type": "图册",
            "properties": {
                "名称": entity_type,
                keyword: url,
                "knowledge_source": json.dumps([{"filename": filename, "info": "", "url": url}], ensure_ascii=False)
            }
        }
        _register_entity_type("图册", properties=["名称", keyword, "knowledge_source"])
        entities.append(entity)
        relations.append({"type": "配套", "from_entity": device, "to_entity": entity_type})
        print(f"  [图/表/目录] 生成图册实体: {entity_type}")

    elif category == "table":
        entity_type = f"{device}的{keyword}"
        entity = {
            "type": keyword,
            "properties": {
                "名称": entity_type,
                keyword: url,
                "knowledge_source": json.dumps([{"filename": filename, "info": "", "url": url}], ensure_ascii=False)
            }
        }
        _register_entity_type(keyword, properties=["名称", keyword, "knowledge_source"])
        entities.append(entity)
        relations.append({"type": "包含", "from_entity": device, "to_entity": entity_type})
        print(f"  [图/表/目录] 生成表实体: {entity_type}")

    elif category == "catalog":
        entity_type = f"{device}的文件目录"
        entity = {
            "type": "文件目录",
            "properties": {
                "名称": entity_type,
                "文件目录": url,
                "knowledge_source": json.dumps([{"filename": filename, "info": "", "url": url}], ensure_ascii=False)
            }
        }
        _register_entity_type("文件目录", properties=["名称", "文件目录", "knowledge_source"])
        entities.append(entity)
        relations.append({"type": "包含", "from_entity": device, "to_entity": entity_type})
        print(f"  [图/表/目录] 生成文件目录实体: {entity_type}")

    elif category == "spare":
        entity_name = f"{device}的备品备件及清单"
        entity = {
            "type": "备品备件及清单",
            "properties": {
                "名称": entity_name,
                "备品备件及清单": url,
                "knowledge_source": json.dumps([{"filename": filename, "info": "", "url": url}], ensure_ascii=False)
            }
        }
        _register_entity_type("备品备件及清单", properties=["名称", "备品备件及清单", "knowledge_source"])
        #entities.append(entity)
        relations.append({
            "type": "包含",
            "from_entity": f"{device}的保障资源",
            "to_entity": entity_name
        })
        print(f"  [图/表/目录] 生成备品备件及清单实体: {entity_name}")

    return entities, relations


# ==============================
# 通用文件：按章节结构提取实体 + 关系
# ==============================

def _chapter_structure_extract(input_pdf_path, data, prefix_url, device):
    """
    调用 chapter_segmenter.segment_chapters 切分章节，再将每个章节转为实体+关系。

    普通章节：一级标题 → 实体，二级标题内容 → 属性字段
    特殊章节：每个二级节各自独立成为实体
    返回: (entities, relations)
    """
    chapters: List[Chapter] = segment_chapters(data)

    all_entities = []
    all_relations = []
    filename = os.path.basename(input_pdf_path)

    for chapter in chapters:
        if chapter.is_skipped:
            continue

        if chapter.is_special:
            # 特殊章节：每个二级节独立成为实体
            for l2 in chapter.l2_sections:
                content_text = "\n".join(
                    p for p in (_record_to_text(r) for r in l2.records) if p
                ).strip()
                if not content_text:
                    continue
                mapped_l2_type = L2_TYPE_MAPPING.get(l2.title, l2.title)
                entity_name = f"{device}的{mapped_l2_type}"
                hl = _collect_highlights(l2.records)
                url = _build_highlight_url(input_pdf_path, hl, prefix_url)
                entity = {
                    "type": mapped_l2_type,
                    "properties": {
                        "名称": entity_name,
                        "内容": content_text,
                        "knowledge_source": json.dumps(
                            [{"filename": filename, "info": content_text, "url": url}],
                            ensure_ascii=False
                        )
                    }
                }
                _register_entity_type(mapped_l2_type, properties=["名称", "内容", "knowledge_source"])
                all_entities.append(entity)
                all_relations.append({
                    "type": _get_relation_type(mapped_l2_type),
                    "from_entity": device,
                    "to_entity": entity_name
                })
                print(f"  → [特殊] 实体 type=【{mapped_l2_type}】名称=【{entity_name}】")

        else:
            # 普通章节：一级标题 → 实体
            mapped_title = _map_l1_title(chapter.title)
            entity_name = f"{device}的{mapped_title}"

            properties = {"名称": entity_name}

            # 一级标题到第一个二级标题之间的内容 → 概述
            loose_text = "\n".join(
                p for p in (_record_to_text(r) for r in chapter.loose_records) if p
            ).strip()
            if loose_text:
                properties["概述"] = loose_text

            # 各二级节：命中 L2_TYPE_MAPPING 的独立成实体，否则作为属性
            for l2 in chapter.l2_sections:
                content_text = "\n".join(
                    p for p in (_record_to_text(r) for r in l2.records) if p
                ).strip()
                if not content_text:
                    continue
                if l2.title in L2_TYPE_MAPPING:
                    mapped_l2_type = L2_TYPE_MAPPING[l2.title]
                    l2_entity_name = f"{device}的{mapped_l2_type}"
                    hl = _collect_highlights(l2.records)
                    l2_url = _build_highlight_url(input_pdf_path, hl, prefix_url)
                    l2_entity = {
                        "type": mapped_l2_type,
                        "properties": {
                            "名称": l2_entity_name,
                            "内容": content_text,
                            "knowledge_source": json.dumps(
                                [{"filename": filename, "info": content_text, "url": l2_url}],
                                ensure_ascii=False
                            )
                        }
                    }
                    _register_entity_type(mapped_l2_type, properties=["名称", "内容", "knowledge_source"])
                    all_entities.append(l2_entity)
                    all_relations.append({
                        "type": _get_relation_type(mapped_l2_type),
                        "from_entity": device,
                        "to_entity": l2_entity_name
                    })
                    print(f"  → [L2独立] 实体 type=【{mapped_l2_type}】名称=【{l2_entity_name}】")
                properties[l2.title] = content_text

            # 整章原文
            raw_parts = [_record_to_text(r) for r in chapter.records]
            full_content = "\n".join(p for p in raw_parts if p).strip()
            print(f"  [DEBUG内容] 章节={mapped_title} records={len(chapter.records)} 内容长度={len(full_content)}")
            if full_content:
                properties["内容"] = full_content

            hl = _collect_highlights(chapter.records)
            url = _build_highlight_url(input_pdf_path, hl, prefix_url)
            properties["knowledge_source"] = json.dumps(
                [{"filename": filename, "info": full_content, "url": url}],
                ensure_ascii=False
            )

            prop_names = [k for k in properties if k not in ("名称", "knowledge_source")]
            _register_entity_type(mapped_title, properties=["名称"] + prop_names + ["knowledge_source"])
            all_entities.append({"type": mapped_title, "properties": properties})
            all_relations.append({
                "type": _get_relation_type(mapped_title),
                "from_entity": device,
                "to_entity": entity_name
            })
            print(f"  → 实体 type=【{mapped_title}】名称=【{entity_name}】属性: {prop_names}")

    print(f"\n[章节提取] 共生成实体: {len(all_entities)}，关系: {len(all_relations)}")
    return all_entities, all_relations


# ==============================
# 对外唯一接口
# ==============================

async def get_filtertext_node_relation(
    input_pdf_path: str,
    data: list,
    prefix_url: str = "http://192.168.0.111:9085",
    device: str = DEVICE_NAME
):
    """
    对外唯一接口。

    自动判断文件名：
      - 含"备品备件"或"图解目录" → 直接跳过，返回空
      - 含 接线图/原理图/总图/连接图/芯线表/随机文件目录 → 直接生成对应实体+关系
      - 其他文件 → 按一级标题=实体、二级标题=属性 的规则提取，同时生成关系

    :param input_pdf_path: PDF文件路径
    :param data: content_list JSON数据
    :param prefix_url: highlight文件服务地址
    :param device: 设备名称，默认"发动机"
    :return: (all_entity_results, all_relation_results)
    """
    # ---- 1. 跳过不需要解析的文件 ----
    if _should_skip(input_pdf_path):
        print(f"[跳过] 文件名命中跳过关键词，不做任何解析: {os.path.basename(input_pdf_path)}")
        return [], []

    # ---- 公共：从封面提取设计单位 ----
    design_unit_entities = []
    design_unit_relations = []
    design_unit_name = _extract_design_unit(data)
    if not design_unit_name:
        p0_texts = [ins.get("text", "")[:40] for ins in data if ins.get("page_idx", -1) == 0]
        print(f"  [设计单位] 未找到，page_idx=0 共{len(p0_texts)}条: {p0_texts}")
    if design_unit_name:
        design_unit_entities.append({
            "type": "设计单位",
            "properties": {
                "名称": design_unit_name,
                "knowledge_source": ""
            }
        })
        design_unit_relations.append({
            "type": "设计",
            "from_entity": device,
            "to_entity": design_unit_name
        })
        _register_entity_type("设计单位", properties=["名称", "knowledge_source"])
        print(f"[设计单位] {design_unit_name}")

    # ---- 2. 图/表/目录类 ----
    category, keyword = _detect_nomodel_category(input_pdf_path)
    if category is not None:
        print(f"[图/表/目录] 文件类别={category}, 关键词={keyword}, device={device}")
        entities, relations = _special_file_generate_entities(
            input_pdf_path, prefix_url, device, category, keyword
        )
        entities = design_unit_entities + entities
        relations = design_unit_relations + relations
        print(f"完成，共生成实体: {len(entities)}，关系: {len(relations)}")
        return entities, relations

    # ---- 3. 其他文件：按章节结构提取 ----
    print(f"[章节提取] 按一级标题→实体、二级标题→属性 提取, device={device}")
    entities, relations = _chapter_structure_extract(
        input_pdf_path, data, prefix_url, device
    )
    entities = design_unit_entities + entities
    relations = design_unit_relations + relations
    print(f"全部完成，共生成实体: {len(entities)}，关系: {len(relations)}")
    return entities, relations


# ==============================
# 本地测试入口
# ==============================

if __name__ == "__main__":
    import asyncio

    test_cases = [
        {
            "pdf": "/app/kg_build/test/122-06A0014-B01001_发动机-维修手册.pdf",
            "json": "/app/kg_build/test/122-06A0014-B01001_发动机-维修手册_content_list.json",
        },
    ]

    all_results = {"entities": [], "relations": []}

    for t in test_cases:
        pdf_path = t["pdf"]
        json_path = t["json"]

        if not os.path.exists(json_path):
            print(f"跳过（文件不存在）: {json_path}")
            continue

        print(f"\n{'='*60}")
        print(f"处理: {pdf_path}")
        print(f"{'='*60}")

        try:
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
        except Exception as e:
            print(f"读取JSON出错: {e}")
            continue

        entities, relations = asyncio.run(
            get_filtertext_node_relation(
                pdf_path, data=data, device=DEVICE_NAME
            )
        )
        all_results["entities"].extend(entities)
        all_results["relations"].extend(relations)

    output_path = "/app/kg_build/test/过滤文本抽取结果.json"
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(all_results, f, ensure_ascii=False, indent=4)
        print(f"\n结果已保存至: {output_path}")
    except Exception as e:
        print(f"保存结果出错: {e}")

