import os
import json
from typing import List, Tuple, Dict, Any, Optional
import re
from kg_build.extract_compose_table import get_compose_node_relation
from kg_build.extract_compose_tujie_table import get_tujiecompose_node_relation
from kg_build.extract_guzhang_table import get_guzhang_node_relation
from kg_build.extract_operation_table import get_operation_node_relation
from kg_build.extract_repair_table import get_repair_node_relation
from kg_build.extract_wxanli_table import get_wxanli_node_relation
from kg_build.get_text_node_reala import extract_text_node_relationship
from kg_build.extract_beipinbeijian_table import get_beipinbeijian_node_relation
from kg_build.extract_filtertext import get_filtertext_node_relation
import requests
from doc2pdf import Doc2PDF
import traceback
from extract_excel_node_relation import get_excel_node_relation

import logging
import asyncio  # ← 新增导入
import time
# --- 日志配置 ---
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)

converter = Doc2PDF()

def build_repair_relationships(entities: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
    """
    根据维修项目编号 OR 名称在故障维修方案中的出现情况，建立“修理”关系。
    匹配规则：只要编号或名称任意一个出现在维修方案文本中，即视为匹配。
    """
    # 1. 数据分类：分离故障模式和维修项目
    faults = [e for e in entities if e.get("type") == "故障模式"]
    maintenances = [e for e in entities if e.get("type") == "维修项目"]
    devices = [e for e in entities if e.get("type") == "设备"]

    
    relationships = []
    entitys = []
    # 2. 遍历每一个维修项目
    for mt in maintenances:
        props = mt.get("properties", {})
        mt_code = props.get("编号")
        mt_name = props.get("名称")
        
        # 如果连名称和编号都没有，无法进行任何匹配，直接跳过
        if not mt_code and not mt_name:
            continue
        
        # 3. 遍历每一个故障模式
        for fault in faults:
            f_props = fault.get("properties", {})
            fault_name = f_props.get("名称")
            repair_text = f_props.get("维修方案", "")
            
            if not fault_name or not repair_text:
                continue
                
            # 4. 核心匹配逻辑修改 (OR 关系)
            match_by_code = mt_code in repair_text if mt_code else False
            match_by_name = mt_name in repair_text if mt_name else False
            
            # 只要满足任一条件，即建立关系
            if match_by_code or match_by_name:
                relationships.append({
                    "type": "修理",
                    "from_entity": mt_name,      # 关系起点：维修项目名称
                    "to_entity": fault_name      # 关系终点：故障名称
                })
    for mt in devices:
        props = mt.get("properties", {})
        mt_code = props.get("维修项目名称")
        mt_name = props.get("维修项目编号")
        devicename = props.get("名称")
        # 如果连名称和编号都没有，无法进行任何匹配，直接跳过
        if not mt_code and not mt_name:
            continue
        
        # 3. 遍历每一个故障模式
        for fault in faults:
            f_props = fault.get("properties", {})
            fault_name = f_props.get("名称")
            repair_text = f_props.get("维修方案", "")
            
            if not fault_name or not repair_text:
                continue
                
            # 4. 核心匹配逻辑修改 (OR 关系)
            match_by_code = mt_code in repair_text if mt_code else False
            match_by_name = mt_name in repair_text if mt_name else False
            
            # 只要满足任一条件，即建立关系
            if match_by_code or match_by_name:
                entitys.append({
                    "type": "维修工作",
                    "properties": {"名称": f"{devicename}的维修工作"}
                })
                relationships.append({
                    "type":        "发生故障使用维修工作",
                    "from_entity": devicename,
                    "to_entity":   f"{devicename}的维修工作"
                })
                relationships.append({
                    "type":        "包含故障",
                    "from_entity": f"{devicename}的维修工作",
                    "to_entity":   fault_name
                })
                relationships.append({
                    "type": "修理",
                    "from_entity": mt_name,      # 关系起点：维修项目名称
                    "to_entity": fault_name 
                })
                
    
    
    return entitys,relationships

async def get_md_node_relation(
    content_list: str,
    input_pdf_path: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]],
    PREFIX_URL : str = "http://192.168.0.111:9085",
    device : str = "",
    device_type : str = ""
) -> Tuple[List[Dict], List[Dict]]:
    """
    核心函数：协调各模块完成知识抽取、清洗、向量化。
    """
    start_time = time.time()
    final_all_entity = []
    final_all_relation = []

    if not content_list:
        logger.error("输入的 Markdown 内容为空")
        return [], []
    if "维修手册" in input_pdf_path or "维修说明书" in input_pdf_path or "培训手册" in input_pdf_path:
        t = time.time()
        all_entity_results_guzhang ,all_relation_results_guzhang  =await get_guzhang_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
        length_entity_guzhang = len(all_entity_results_guzhang)
        length_relation_guzhang = len(all_relation_results_guzhang)
        logger.info(f"故障抽取 - 实体: {length_entity_guzhang}, 关系: {length_relation_guzhang}, 耗时: {time.time()-t:.2f}s")
        if length_entity_guzhang > 0 :
            final_all_entity.extend(all_entity_results_guzhang)
        if length_relation_guzhang > 0:
            final_all_relation.extend(all_relation_results_guzhang)

        t = time.time()
        all_entity_results_repair ,all_relation_results_repair  =await get_repair_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL)
        length_entity_repair = len(all_entity_results_repair)
        length_relation_repair = len(all_relation_results_repair)
        logger.info(f"维修项目表抽取 - 实体: {length_entity_repair}, 关系: {length_relation_repair}, 耗时: {time.time()-t:.2f}s")
        if length_entity_repair > 0 :
            final_all_entity.extend(all_entity_results_repair)
        if length_relation_repair > 0:
            final_all_relation.extend(all_relation_results_repair)

        t = time.time()
        all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
        length_entity_text= len(all_entity_results_text)
        length_relation_text  = len(all_relation_results_text)
        logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
        if length_entity_text > 0 :
            final_all_entity.extend(all_entity_results_text)
        if length_relation_text > 0:
            final_all_relation.extend(all_relation_results_text)
    elif "图解目录" in input_pdf_path:
        # 2. 处理图解组成表
        t = time.time()
        all_entity_results_tujiecompose ,all_relation_results_tujiecompose =await get_tujiecompose_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device,lower_entity_type=device_type)
        length_entity_tujiecompose = len(all_entity_results_tujiecompose)
        length_relation_tujiecompose  = len(all_relation_results_tujiecompose)
        logger.info(f"图解组成表抽取 - 实体: {length_entity_tujiecompose}, 关系: {length_relation_tujiecompose}, 耗时: {time.time()-t:.2f}s")
        if length_entity_tujiecompose > 0 :
            final_all_entity.extend(all_entity_results_tujiecompose)
        if length_relation_tujiecompose > 0:
            final_all_relation.extend(all_relation_results_tujiecompose)
    
    elif "操作使用手册" in input_pdf_path:
        # 6. 处理操作项目表
        t = time.time()
        all_entity_results_operation ,all_relation_results_goperation =await get_operation_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
        length_entity_operation= len(all_entity_results_operation)
        length_relation_goperation  = len(all_relation_results_goperation)
        logger.info(f"操作项目表抽取 - 实体: {length_entity_operation}, 关系: {length_relation_goperation}, 耗时: {time.time()-t:.2f}s")
        if length_entity_operation > 0 :
            final_all_entity.extend(all_entity_results_operation)
        if length_relation_goperation > 0:
            final_all_relation.extend(all_relation_results_goperation)

        t = time.time()
        all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
        length_entity_text= len(all_entity_results_text)
        length_relation_text  = len(all_relation_results_text)
        logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
        if length_entity_text > 0 :
            final_all_entity.extend(all_entity_results_text)
        if length_relation_text > 0:
            final_all_relation.extend(all_relation_results_text)

    elif "备品备件" in input_pdf_path:
        # 5. 处理备品备件表
        t = time.time()
        all_entity_results_beipinbeijian ,all_relation_results_beipinbeijian =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
        length_entity_beipinbeijian= len(all_entity_results_beipinbeijian)
        length_relation_beipinbeijian  = len(all_relation_results_beipinbeijian)
        logger.info(f"备品备件表抽取 - 实体: {length_entity_beipinbeijian}, 关系: {length_relation_beipinbeijian}, 耗时: {time.time()-t:.2f}s")
        if length_entity_beipinbeijian > 0 :
            final_all_entity.extend(all_entity_results_beipinbeijian)
        if length_relation_beipinbeijian > 0:
            final_all_relation.extend(all_relation_results_beipinbeijian)
    elif "故障案例" in input_pdf_path:

        # 7. 处理维修案例
        t = time.time()
        all_entity_results_wxanli ,all_relation_results_wxanli =await get_wxanli_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL)
        length_entity_wxanli= len(all_entity_results_wxanli)
        length_relation_wxanli  = len(all_relation_results_wxanli)
        logger.info(f"维修案例抽取 - 实体: {length_entity_wxanli}, 关系: {length_relation_wxanli}, 耗时: {time.time()-t:.2f}s")
        if length_entity_wxanli > 0 :
            final_all_entity.extend(all_entity_results_wxanli)
        if length_relation_wxanli > 0:
            final_all_relation.extend(all_relation_results_wxanli)
    else:
        logger.warning("未识别的文档类型，默认进行文本抽取")
        t = time.time()
        all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
        length_entity_text= len(all_entity_results_text)
        length_relation_text  = len(all_relation_results_text)
        logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
        if length_entity_text > 0 :
            final_all_entity.extend(all_entity_results_text)
        if length_relation_text > 0:
            final_all_relation.extend(all_relation_results_text)


    # 8. 建立修复关系
    reparientitys,repair_relations = build_repair_relationships(final_all_entity)
    final_all_entity.extend(reparientitys)
    final_all_relation.extend(repair_relations)
    logger.info(f"构建故障修复实体数量: {len(reparientitys)}")
    logger.info(f"构建故障修复关系数量: {len(repair_relations)}")
    


    logger.info(f"总实体: {len(final_all_entity)}, 总关系: {len(final_all_relation)}")
    result = {
        "entities": final_all_entity,
        "relations": final_all_relation
    }
    output_json_path= "/app/wxproject.json"
    # 保存为 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}")
    end_time = time.time()
    print(f"总耗时: {end_time - start_time:.2f} 秒")
    return final_all_entity, final_all_relation

def get_content_list_json(filename, folder_path):
    """
    根据传入的文件名，在指定文件夹下查找对应的 _content_list.json 文件并读取内容。
    
    参数:
    filename (str): 原始文件名，例如 "163-06A0015-B01001_发动机-操作使用手册.pdf"
    folder_path (str): 目标文件夹路径
    
    返回:
    dict or None: 如果找到文件则返回 JSON 内容字典，否则返回 None
    """
    # 提取文件名主体（去掉扩展名）
    base_name = os.path.splitext(filename)[0]
    
    # 构造目标 JSON 文件名
    target_json_filename = f"{base_name}_content_list.json"
    
    # 构造完整路径
    target_path = os.path.join(folder_path, target_json_filename)
    
    # 检查文件是否存在
    if not os.path.exists(target_path):
        print(f"未找到文件: {target_path}")
        return None
    
    # 读取并解析 JSON
    try:
        with open(target_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return data
    except json.JSONDecodeError as e:
        print(f"JSON 解析错误: {e}")
        return None
    except Exception as e:
        print(f"读取文件时发生错误: {e}")
        return None


def process_file_content(
    file_name: str,
    file_path: str,
    task_id: str,
    data_dir: str,
    ontology: Any,
    ontology_desc: str,
    relationships: List,
    node_types: List,
    relationship_types: List,
    prefix_url: str,
    api_url: str,
    check_cancelled_callback: callable,
    device : str,
    device_type : str,
) -> Tuple[List[Dict], List[Dict]]:
    """
    根据文件类型解析内容并抽取实体和关系。
    
    Returns:
        tuple: (entities, relationships)
    """
    file_extension = os.path.splitext(file_name)[1].lower()
    merged_output_data = {"entities": [], "relationships": []}
    from fileparse_util import process_document
    
    # --- 处理 PDF 和 MD ---
    if file_extension in ['.pdf', '.md']:
        async def _run_process_doc():
            return await process_document(file_name)
        doc_result = asyncio.run(_run_process_doc())
        if doc_result is not None:
            content_list, images_dict, md_content= doc_result

        else:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件未找到: {file_path}")

            with open(file_path, "rb") as f:
                files = {"file": (os.path.basename(file_path), f, "application/pdf")}
                logger.info(file_path)
                response = requests.post(api_url, files=files, timeout=12000)

            if response.status_code != 200:
                logger.error(f"❌ API 请求失败，状态码: {response.status_code}, 响应: {response.text}")
                raise Exception(f"PDF/MD 解析 API 失败: {response.status_code}")

            result = response.json()
            data = result.get("data", {})
            content_list = data.get("content_list")
        
        logger.info(f"[{task_id}] 开始文件 {file_name} 的实体和关系抽取")
        if check_cancelled_callback():
            raise InterruptedError("任务已取消")

        async def run_extraction():
            return await get_md_node_relation(
                content_list=content_list,
                input_pdf_path=file_path,
                Ontology=ontology,
                Ontology_DESCRIPTION=ontology_desc,
                Relationships=relationships,
                RELATIONSHIP_TYPES=relationship_types,
                NODE_TYPES=node_types,
                PREFIX_URL=prefix_url,
                device = device,
                device_type = device_type
            )

        final_all_entity, final_all_relation = asyncio.run(run_extraction())
        merged_output_data = {"entities": final_all_entity, "relationships": final_all_relation}

    # --- 处理 DOCX ---
    elif file_extension in ['.docx']:
        base_name = os.path.splitext(file_name)[0]
        pdf_filename = f"{base_name}.pdf"
        pdf_path = os.path.join(data_dir, pdf_filename)

        logger.info(f"[{task_id}] 转换 DOCX -> PDF: {file_name}")

        try:
            converter.convert(file_path, output_dir=data_dir)

            if not os.path.exists(pdf_path):
                raise FileNotFoundError(f"PDF 转换失败: {pdf_path}")

            # 与 PDF 分支保持一致：先尝试 process_document，失败则上传 API
            async def _run_process_doc_docx():
                return await process_document(pdf_filename)
            doc_result = asyncio.run(_run_process_doc_docx())
            if doc_result is not None:
                content_list, _ = doc_result
            else:
                with open(pdf_path, "rb") as f:
                    files = {"file": (pdf_filename, f, "application/pdf")}
                    logger.info(pdf_path)
                    response = requests.post(api_url, files=files, timeout=12000)

                if response.status_code != 200:
                    logger.error(f"❌ API 请求失败，状态码: {response.status_code}, 响应: {response.text}")
                    raise Exception(f"PDF 解析 API 失败: {response.status_code}")

                result = response.json()
                data = result.get("data", {})
                content_list = data.get("content_list")

            logger.info(f"[{task_id}] 开始文件 {file_name} 的实体和关系抽取")
            if check_cancelled_callback():
                raise InterruptedError("任务已取消")

            async def run_extraction_docx():
                return await get_md_node_relation(
                    content_list=content_list,
                    input_pdf_path=pdf_path,
                    Ontology=ontology,
                    Ontology_DESCRIPTION=ontology_desc,
                    Relationships=relationships,
                    RELATIONSHIP_TYPES=relationship_types,
                    NODE_TYPES=node_types,
                    PREFIX_URL=prefix_url,
                    device=device,
                    device_type=device_type,
                )

            final_all_entity, final_all_relation = asyncio.run(run_extraction_docx())
            merged_output_data = {"entities": final_all_entity, "relationships": final_all_relation}

        except Exception as e:
            logger.error(f"[{task_id}] DOCX 处理失败: {str(e)}")
            logger.error(traceback.format_exc())
            raise
        finally:
            pass

    # --- 处理 XLSX ---
    elif file_extension in ['.xlsx']:
        base_name = os.path.splitext(file_name)[0]
        pdf_filename = f"{base_name}.pdf"
        logger.info(f"[{task_id}] 转换 xlsx -> PDF (仅生成索引引用): {file_name}")
        # 注意：原代码中这里也调用了 converter，但后续逻辑主要依赖 Excel 解析
        try:
            converter.convert(file_path, output_dir=os.path.join(data_dir, "uploads_v1"))
        except Exception as e:
            logger.warning(f"Excel 转 PDF 失败，但不影响数据提取: {e}")
        
        all_entities, all_relationships = get_excel_node_relation(file_path, sheet_name="Sheet1")
        
        # 补充知识源信息
        for entity in all_entities:
            knowledge_source = []
            qiepian = entity["properties"].get("切片")
            url = f"{prefix_url}/files/uploads/{pdf_filename}"
            knowledge_source.append({
                "filename": file_name,
                "info": qiepian,
                "url": url
            })
            entity["properties"]["knowledge_source"] = knowledge_source
            
        merged_output_data = {"entities": all_entities, "relationships": all_relationships}

    else:
        raise ValueError(f"不支持的文件类型: {file_extension}")

    # 统一取消检查
    if check_cancelled_callback():
        raise InterruptedError("任务已取消")

    # 保存缓存 (可选，保持原有逻辑)
    try:
        safe_fname = file_name.replace(os.sep, "_").replace(" ", "_")
        cache_dir = os.path.join(data_dir, "kg_cache")
        os.makedirs(cache_dir, exist_ok=True)
        cache_path = os.path.join(cache_dir, f"merged_output_{task_id}_{safe_fname}.json")
        with open(cache_path, "w", encoding="utf-8") as cf:
            json.dump(merged_output_data, cf, ensure_ascii=False, indent=2)
        logger.info(f"[{task_id}] Saved merged_output_data cache: {cache_path}")
    except Exception as e:
        logger.warning(f"[{task_id}] Failed to save merged_output_data cache: {e}")

    return merged_output_data["entities"], merged_output_data["relationships"]

def check_filename_ocr(filename: str) -> Optional[str]:
    """
    根据文件名关键词判断是否需要 OCR。

    Returns:
        "ocr_true"  — 文件名包含维修/操作类关键词
        "ocr_false" — 文件名包含原理/接线/总图类关键词
        None        — 未命中任何关键词
    """
    #ocr_true_keywords  = ["维修手册", "维修说明书", "操作使用手册", "图解目录","培训手册",""]
    # ocr_false_keywords = ["原理图", "接线图", "总图","随机文件目录","安装说明书","芯线表","链接图","备品备件及工具清单"]
    ocr_false_keywords = ["原理图", "接线图", "总图","连接图","线图","布置图","安装图","曲线图","标志图","综合图","结构图","基座图","型线图","展开图","编号图","剖面图","吊装图","总装图","曲线图","代号图","视界图","射角图","管路图","周界图","外形图","关系图","装配图","覆盖图","尺寸图","完工图","芯线表","明细表","接线表","编制表","值表册","线路表","矩阵表","汇总表","信息表","连接表","随机文件目录", "随机文件","备品备件","图解目录","电流图","接线图和芯线表","外形图","外形图、外部接线图、安装图","设备总图", "线路图"]

   # for kw in ocr_true_keywords:
    #    if kw in filename:
    #        return "ocr_true"

    for kw in ocr_false_keywords:
        if kw in filename:
            return "ocr_false"
#        else:
 #           return "ocr_true"

    #return None
    return "ocr_true"


