import os
import json
from typing import List, Tuple, Dict, Any

# --- 修复了拼写错误的导入路径 ---
from kg_build_v2.extract_wxanli_node_rela import get_wxanli_entity_rela
from kg_build_v2.extract_operation_node_rela import get_operation_node_rela
from kg_build_v2.extract_repair_node_rela import get_repair_node_rela
# 修正拼写：extrat_normal -> extract_normal (请确保文件名实际为 extract_normal_node_rela.py)
from kg_build_v2.extrat_normal_node_rela import get_nromal_node_rela 
from langchain_text_splitters import MarkdownHeaderTextSplitter, CharacterTextSplitter
from kg_build_v2.get_text_node_reala import extract_text_node_relationship
from kg_build_v2.utils import remove_duplicates, deduplicate_relationships
from modelsAPI.model_api import OpenaiAPI
import logging

# --- 日志配置 ---
logger = logging.getLogger(__name__)
if not logger.handlers:
    handler = logging.StreamHandler()
    fmt = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(fmt)
    logger.addHandler(handler)
logger.setLevel(logging.INFO)


def _split_markdown(text: str, chunk_size: int = 2048, chunk_overlap: int = 50) -> Tuple[List[str], List[str]]:
    """
    将 Markdown 文本按一级标题分割，并对过长文本进行二次分块。
    """
    headers_to_split_on = [("#", "Header_1")]
    splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on, strip_headers=True)
    chunks = splitter.split_text(text)

    char_splitter = CharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separator="\n\n"
    )

    final_chunks = []
    start_chars = {"1", "2", "3", "4", "5", "6", "7", "8", "9", "0", "第", "十", "一", "二", "三", "四", "五", "六", "七", "八", "九", "故障"}
    
    # 提取有效标题
    valid_headers = {
        chunk.metadata.get('Header_1', '') 
        for chunk in chunks 
        if chunk.metadata.get('Header_1', '') and chunk.metadata['Header_1'][0] in start_chars
    }

    # 构建带标题前缀的块
    for chunk in chunks:
        header_1 = chunk.metadata.get('Header_1', '')
        content = chunk.page_content
        
        if len(content) <= 20000:
            chunk_with_header = f"{header_1}\n{content}" if header_1 else content
            final_chunks.append(chunk_with_header)
        else:
            split_texts = char_splitter.split_documents([chunk])
            for split_chunk in split_texts:
                chunk_with_header = f"{header_1}\n{split_chunk.page_content}" if header_1 else split_chunk.page_content
                final_chunks.append(chunk_with_header)

    # 重新分组合并
    rechunked = []
    rechunked_headers = []
    current_content = ""
    current_header = ""

    for chunk in final_chunks:
        # 检查是否为新章节
        new_header = next((h for h in valid_headers if chunk.startswith(h + "\n")), None)
        
        if new_header:
            if current_content:
                rechunked.append(current_content)
                rechunked_headers.append(current_header)
            current_content = chunk
            current_header = new_header
        else:
            current_content += "\n" + chunk if current_content else chunk

    # 添加最后一段
    if current_content:
        rechunked.append(current_content)
        rechunked_headers.append(current_header)

    return rechunked, rechunked_headers


def read_markdown_file(file_path: str) -> str:
    """
    读取 Markdown 文件内容。
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    except Exception as e:
        logger.error(f"读取文件失败 [{file_path}]: {e}")
        return ""


def get_html_node_rela(md_content: str) -> Tuple[List[Dict], List[Dict], str]:
    """
    从 HTML/Markdown 表格中提取实体和关系。
    """
    if not md_content:
        return [], [], ""

    all_entities = []
    all_relations = []
    cleaned_content = ""

    # 1. 提取维修项目表
    try:
        repair_entities = get_repair_node_rela(md_content)
        if isinstance(repair_entities, list):
            all_entities.extend(repair_entities)
            # --- 修复日志调用：使用 f-string ---
            logger.info(f"维修项目表 - 实体数量: {len(repair_entities)}")
    except Exception as e:
        logger.error(f"提取维修项目表时发生错误: {e}")

    # 2. 提取操作项目表
    try:
        result_operation = get_operation_node_rela(md_content)
        operation_entities = []
        operation_relations = []
        
        if isinstance(result_operation, dict):
            operation_entities = result_operation.get("entities", [])
            operation_relations = result_operation.get("relationships", [])
        elif isinstance(result_operation, list):
            logger.warning("get_operation_node_rela 返回了列表而非字典")

        all_entities.extend(operation_entities)
        all_relations.extend(operation_relations)
        # --- 修复日志调用 ---
        logger.info(f"操作项目表 - 实体: {len(operation_entities)}, 关系: {len(operation_relations)}")
        
    except Exception as e:
        logger.error(f"提取操作项目表时发生错误: {e}")

    # 3. 提取常规表格 (普通表、故障表等)
    # --- 修复拼写错误：get_nromal_node_rela -> get_nromal_node_rela ---
    try:
        result_normal, content_normal = get_nromal_node_rela(md_content) 
        normal_entities = []
        normal_relations = []
        
        if isinstance(result_normal, list) and result_normal:
            item = result_normal[0] 
            if isinstance(item, dict):
                normal_entities = item.get("entities", [])
                normal_relations = item.get("relationships", [])
        elif isinstance(result_normal, dict):
            normal_entities = result_normal.get("entities", [])
            normal_relations = result_normal.get("relationships", [])

        all_entities.extend(normal_entities)
        all_relations.extend(normal_relations)
        cleaned_content = content_normal or ""
        
        # --- 修复日志调用 ---
        logger.info(f"常规表格 - 实体: {len(normal_entities)}, 关系: {len(normal_relations)}")
        
    except Exception as e:
        logger.error(f"提取常规表格时发生错误: {e}")

    return all_entities, all_relations, cleaned_content


def get_wxanli_node_rela(md_content: str) -> Tuple[List[Dict], List[Dict]]:
    """
    从维修案例中提取实体和关系。
    """
    try:
        result_wxanli = get_wxanli_entity_rela(md_content)
        if isinstance(result_wxanli, dict):
            wxanli_entities = result_wxanli.get("entities", [])
            wxanli_relations = result_wxanli.get("relationships", [])
        else:
            wxanli_entities = []
            wxanli_relations = []
            logger.warning("维修案例返回数据格式异常，使用空列表代替")
    except Exception as e:
        logger.error(f"提取维修案例时发生错误: {e}")
        wxanli_entities = []
        wxanli_relations = []

    return wxanli_entities, wxanli_relations


def build_repair_relationships(entities: List[Dict]) -> List[Dict]:
    """
    根据维修项目编号和故障描述建立“修理”关系。
    """
    faults = [e for e in entities if e.get("type") == "故障"]
    maintenances = [e for e in entities if e.get("type") == "维修项目"]
    
    relationships = []
    
    for mt in maintenances:
        mt_code = mt.get("properties", {}).get("维修项目编号")
        mt_name = mt.get("properties", {}).get("名称")
        
        if not mt_code or not mt_name:
            continue
        
        for fault in faults:
            fault_name = fault.get("properties", {}).get("名称")
            repair_text = fault.get("properties", {}).get("维修项目", "")
            
            if not fault_name:
                continue
                
            if mt_code in repair_text:
                relationships.append({
                    "type": "修理",
                    "from_entity": mt_name,
                    "to_entity": fault_name
                })
    
    return relationships


def get_md_node_relation(
    md_content: str,
    Ontology: List[str],
    Ontology_DESCRIPTION: Dict[str, List[str]],
    Relationships: List[str],
    RELATIONSHIP_TYPES: Dict[str, List[str]],
    NODE_TYPES: Dict[str, List[str]]
) -> Tuple[List[Dict], List[Dict]]:
    """
    核心函数：协调各模块完成知识抽取、清洗、向量化。
    """
    final_all_entity = []
    final_all_relation = []

    if not md_content:
        logger.error("输入的 Markdown 内容为空")
        return [], []

    # 1. 处理表格类数据
    table_entities, table_relations, cleaned_content = get_html_node_rela(md_content)
    logger.info(f"表格类抽取 - 实体: {len(table_entities)}, 关系: {len(table_relations)}")

    # 2. 处理维修案例
    case_entities, case_relations = get_wxanli_node_rela(md_content)
    logger.info(f"维修案例抽取 - 实体: {len(case_entities)}, 关系: {len(case_relations)}")

    # 合并表格与案例
    all_entities = table_entities + case_entities
    all_relations = table_relations + case_relations
    final_all_entity.extend(all_entities)
    final_all_relation.extend(all_relations)

    # 3. 建立修复关系
    repair_relations = build_repair_relationships(all_entities)
    final_all_relation.extend(repair_relations)
    logger.info(f"构建故障修复关系数量: {len(repair_relations)}")

    # 4. 处理纯文本
    if cleaned_content.strip():
        chunks, _ = _split_markdown(cleaned_content)
        text_entities, text_relationships = extract_text_node_relationship(
            chunks, 
            ontology=Ontology,
            Ontology_DESCRIPTION=Ontology_DESCRIPTION,
            relationships=Relationships,
            relationship_types=RELATIONSHIP_TYPES,
            node_types=NODE_TYPES
        )
        final_all_entity.extend(text_entities)
        final_all_relation.extend(text_relationships)
        logger.info(f"文本抽取 - 实体: {len(text_entities)}, 关系: {len(text_relationships)}")
    else:
        logger.info("未检测到有效清洗文本，跳过文本抽取步骤")

    # 5. 去重
    final_all_entity = remove_duplicates(final_all_entity)
    final_all_relation = deduplicate_relationships(final_all_relation)
    logger.info(f"去重后 - 总实体: {len(final_all_entity)}, 总关系: {len(final_all_relation)}")

    # 6. 实体向量化
#    for entity in final_all_entity:
 #       name = entity["properties"].get("名称")
  #      if name:
            # --- 优化变量名 ---
   #         embedding_vector = OpenaiAPI.generate_embedding(name)
    #        entity["properties"]["embedding"] = embedding_vector
   # logger.info("实体向量化完成")

    return final_all_entity, final_all_relation

if __name__ == "__main__":
    input_directory = "/app/测试FDJwx手册_shenghna_output11_full_content.md"
    final_all_entity = []
    final_all_relation = []
    # 调用目录处理函数
    mdcontent = read_markdown_file(input_directory)
    
    if mdcontent is None:
        print("无法读取文件，程序退出")
        exit(1)
    Ontology = [
        "舰艇", "系统", "设备", "零部件", "故障", "修理方案", "维修项目", "操作项目", "保养", "备品备件"
    ]
    Ontology_DESC = {
    "舰艇": ["代表一艘完整海上作战或保障平台的实体，是所有系统、设备和零部件的最高层级载体。区别于“系统”或“设备”，舰艇具有独立的型号、弦号和部队归属，不可作为其他舰艇的组成部分。"],
    "系统": ["舰艇中实现某一类功能（如动力、导航、通信）的集成化功能单元，由多个设备协同组成。与“设备”不同，系统本身不直接对应物理装置，而是逻辑/功能层面的集合；与“舰艇”不同，系统不具备独立航行能力或部队归属。"],
    "设备": ["安装在舰艇上、具有独立型号和运行状态的物理装置，可被启停、监控运行时间，并可能产生故障。区别于“零部件”，设备是可独立更换或维修的功能单元；区别于“系统”，设备是具体硬件而非功能集合。"],
    "零部件": ["构成设备的不可再拆分的基础物理元件（如轴承、电路板、密封圈），通常无独立运行状态或启用时间。与“设备”关键区别在于：零部件不能单独运行，也不记录运行时间或技术状态变更。"],
    "故障": ["设备或系统在运行中出现的异常状态或失效事件，必须包含具体现象（如“无法启动”）、原因（如“电源短路”）及影响。区别于“修理”或“保养”，故障是问题本身，而非应对措施；不包含操作步骤或人员信息。"],
    "修理方案": ["针对已发生故障所执行的一次性修复行为，包含具体时间、人员、工具、方案和所关联的维修项目。区别于“保养”，修理是响应式、非周期性的；区别于“维修项目”，修理是实际发生的实例，而非标准化规程。"],
    "维修项目": ["标准化、可重复执行的维修任务规程，规定了某类设备或系统在特定条件下应如何维修，包含编号、步骤、材料、安全要求等。区别于“修理”，维修项目是模板或规范，而非具体事件；不包含实际发生时间或人员。"],
    "操作项目": ["舰艇正常运行中需执行的标准操作流程（如启动主机、校准雷达），强调操作步骤与安全保障。区别于“维修项目”，操作使用项目用于日常使用而非维护"],
    "保养" : ["按计划周期性执行的预防性维护活动（如润滑、清洁、检测），目的是延缓性能退化。"]
}
    relationships = [
        "包含", "组成", "解决", "用于", "使用", "具有","存储备用"
    ]
    relationship_types = {
        "包含": ["舰艇包含系统", "系统包含设备", "设备包含零部件", "修理包含维修项目"],
        "解决": ["修理解决故障", "维修项目解决故障"],
        "使用": ["设备使用操作项目", "设备使用保养", "设备使用维修项目"],
        "具有": ["设备具有故障", "零部件具有故障"],
        "存储备用" : ["设备存储备用备品备件"]
    }
    node_types = {
        "舰艇": ["名称", "型号", "弦号", "设计单位", "总装工厂", "所属部队"],
        "系统": ["名称", "技术责任单位", "生产单位", "级别"],
        "设备": ["名称", "型号", "技术责任单位", "生产单位", "技术状态", "生产日期", "启用时间", "运行时间", "运行条件"],
        "零部件": ["名称", "型号", "技术责任单位", "生产单位"],
        "故障": ["名称", "故障原因", "故障影响", "故障类型", "发生时间", "处置情况", "修复情况"],
        "修理方案": ["名称", "修理人员", "修理时间", "修理工具", "备品备件", "修理类型", "维修项目"],
        "维修项目": ["名称", "维修项目编号", "组成编码", "维修设备名称", "维修间隔期", "维修级别", "页码", "安全防护要求", "安全注意事项", "工具", "备品备件", "消耗材料", "人力人员", "修理步骤"],
        "操作项目": ["名称", "操作项目编号", "安全防护", "注意事项", "保障设施、保障设备（工具）", "人力人员", "操作步骤"],
        "保养": ["保养人员", "保养时间", "保养工具", "保养频率"],
        "备品备件" : ["名称", "规格", "型号", "图号或标准号或代号","计量单位", "单机重量", "尺寸（长X宽X高mm）", "单机装配数量","实物图片电子文档名称","MTBF(b)","生成厂家",
        "出厂日期",
        "供货周期（天）",
        "有效期（月）",
        "使用寿命（h)",
        "备件箱号",
        "储存要求_环境",
        "储存要求_期限",
        "停产后供应方式"],
    }
    final_all_entity,final_all_relation = get_md_node_relation(mdcontent,Ontology=Ontology,Ontology_DESCRIPTION=Ontology_DESC,Relationships=relationships,RELATIONSHIP_TYPES=relationship_types,NODE_TYPES=node_types)
