#!/usr/bin/env python3
"""
批量知识图谱构建脚本 —— 三阶段 ETL

阶段1 (extract): 对文件夹内所有文件做实体抽取，每个文件结果存为 JSON
阶段2 (merge):   汇总所有 JSON，去重，生成 embedding，写入 Neo4j

用法:
  # 只跑抽取（4个容器各跑各的文件夹）
  python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase extract --workers 4

  # 只跑入库（所有容器跑完后，跑一次）
  python batch_kg_build.py --output_dir /app/kg_output --phase merge

  # 全流程一次跑完
  python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase all --workers 4
"""

import os
import json
import logging
import argparse
import asyncio
import traceback
import time
from datetime import datetime
from multiprocessing import Process, Queue
from typing import List, Dict, Tuple, Any, Optional

EXTRACT_TIMEOUT = 600  # 单文件抽取超时秒数
FAULT_DIR       = "/app/fault_wenjian"


def _record_fault(file_name: str):
    """将超时文件名追加到 /app/fault_wenjian/timeout_YYYYMMDD.txt"""
    os.makedirs(FAULT_DIR, exist_ok=True)
    date_str   = datetime.now().strftime("%Y%m%d")
    fault_path = os.path.join(FAULT_DIR, f"timeout_{date_str}.txt")
    with open(fault_path, "a", encoding="utf-8") as f:
        f.write(file_name + "\n")
    logger.warning(f"[FAULT] 已记录超时文件: {fault_path}")

# ── 项目内模块 ──────────────────────────────────────────────────────────────
from extract_text_node_relation import process_file_content, check_filename_ocr
from kg_build.extract_filtertext import get_filtertext_node_relation
from filename_proceess_and_kgquery import get_entity
from kg_build.filter_node_relation import filter_node_relationship
from kg_build.kg_save import MilitaryKnowledgeGraph
from modelsAPI.model_api import OpenaiAPI
import importlib
import entity_registry as _entity_registry
from default_ontology_config import (
    ONTOLOGY          as _BASE_ONTOLOGY,
    ONTOLOGY_DESCRIPTION as _BASE_ONTOLOGY_DESC,
    NODE_TYPES        as _BASE_NODE_TYPES,
    RELATIONSHIPS     as RELATIONSHIPS,
    RELATIONSHIP_TYPES as RELATIONSHIP_TYPES,
)


def _build_ontology_config():
    """
    与 app.py build_kg_config 保持一致：
    default_ontology_config 的静态配置 + entity_registry.json 动态注册的实体类型合并。
    """
    importlib.reload(_entity_registry)
    reg_list = _entity_registry.get_ontology_list()
    reg_desc = _entity_registry.get_ontology_description()
    reg_nodes = _entity_registry.get_node_types()

    # 合并：静态配置优先（手写配置不被覆盖）
    ontology = list(dict.fromkeys(_BASE_ONTOLOGY + reg_list))       # 去重保序
    ontology_desc = {**reg_desc, **_BASE_ONTOLOGY_DESC}
    node_types = {**reg_nodes, **_BASE_NODE_TYPES}
    return ontology, ontology_desc, node_types


# 模块加载时构建一次（抽取阶段开始前已有 registry 内容）
ONTOLOGY, ONTOLOGY_DESC, NODE_TYPES = _build_ontology_config()

# ── 配置（与 app.py 保持一致，按需修改） ────────────────────────────────────
NEO4J_URI        = os.getenv("NEO4J_URI",      "bolt://localhost:7687")
NEO4J_USER       = os.getenv("NEO4J_USER",     "neo4j")
NEO4J_PASSWORD   = os.getenv("NEO4J_PASSWORD", "zdht123@")
MINERU_ANALYZE   = os.getenv("MINERU_ANALYZE", "http://192.168.1.64:18000/analyze-pdf")
MINERU_CACHE     = os.getenv("MINERU_CACHE",   "http://192.168.1.64:18000/get_content_list")
PREFIX_URL       = os.getenv("PREFIX_URL",     "http://192.168.1.64:9085")

SUPPORTED_EXT    = {".pdf", ".md", ".docx", ".xlsx"}

# ── 日志 ────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("batch_kg_build.log", encoding="utf-8"),
    ],
)
logger = logging.getLogger(__name__)


# ════════════════════════════════════════════════════════════════════════════
# 阶段 1：单文件抽取
# ════════════════════════════════════════════════════════════════════════════

def extract_single_file(file_path: str, output_dir: str) -> Dict:
    """
    对单个文件做实体/关系抽取，结果保存到 output_dir/<filename>.json。
    已存在 JSON 则跳过（断点续跑）。
    返回 {"file": ..., "status": "ok"/"skip"/"error", "entity_count": ..., "rel_count": ...}
    """
    file_name = os.path.basename(file_path)
    ext = os.path.splitext(file_name)[1].lower()
    out_path = os.path.join(output_dir, file_name + ".json")

    # ── 断点续跑：已有结果直接跳过 ──────────────────────────────────────────
    if os.path.exists(out_path):
        logger.info(f"[SKIP] 已有结果，跳过: {file_name}")
        return {"file": file_name, "status": "skip"}

    if ext not in SUPPORTED_EXT:
        logger.warning(f"[SKIP] 不支持的文件类型: {file_name}")
        return {"file": file_name, "status": "skip"}

    t0 = time.time()
    entities: List[Dict] = []
    relations: List[Dict] = []
    lower_entity_device: Optional[str] = None
    lower_entity_type:   Optional[str] = None

    try:
        # ── 1. 文件名实体提取 ────────────────────────────────────────────────
        try:
            fn_result = get_entity(file_name)
            fn_entities  = fn_result.get("entities", [])
            fn_relations = fn_result.get("relationships", [])
            lower_entity = fn_result.get("low_level")
            for e in fn_entities:
                if e.get("properties", {}).get("名称") == lower_entity:
                    lower_entity_device = e["properties"].get("名称")
                    lower_entity_type   = e.get("type")
                    break
        except Exception as e:
            logger.warning(f"[{file_name}] 文件名实体提取失败: {e}")
            fn_entities, fn_relations = [], []

        # ── 2. 正文实体抽取 ──────────────────────────────────────────────────
        check = check_filename_ocr(file_name)

        if check == "ocr_true":
            entities, relations = process_file_content(
                file_name             = file_name,
                file_path             = file_path,
                task_id               = "batch",
                data_dir              = os.path.dirname(file_path),
                ontology              = ONTOLOGY,
                ontology_desc         = ONTOLOGY_DESC,
                relationships         = RELATIONSHIPS,
                node_types            = NODE_TYPES,
                relationship_types    = RELATIONSHIP_TYPES,
                prefix_url            = PREFIX_URL,
                api_url               = MINERU_ANALYZE,
                check_cancelled_callback = lambda: False,
                device                = lower_entity_device,
                device_type           = lower_entity_type,
                get_contentlist_url   = MINERU_CACHE,
            )

        elif check == "ocr_false":
            async def _run_filtertext():
                return await get_filtertext_node_relation(
                    input_pdf_path = file_path,
                    data           = [],
                    prefix_url     = PREFIX_URL,
                    device         = lower_entity_device,
                )
            result = asyncio.run(_run_filtertext())
            if isinstance(result, tuple) and len(result) >= 2:
                entities, relations = result[0], result[1]
            else:
                entities, relations = [], []

        else:
            logger.warning(f"[{file_name}] check_filename_ocr 返回 None，跳过正文抽取")

        # ── 3. 合并文件名实体 ────────────────────────────────────────────────
        all_entities  = entities  + fn_entities
        all_relations = relations + fn_relations
        # 过滤掉没有名称的实体
        all_entities = [e for e in all_entities if e.get("properties", {}).get("名称")]

        # ── 4. 保存 JSON ─────────────────────────────────────────────────────
        os.makedirs(output_dir, exist_ok=True)
        payload = {"entities": all_entities, "relationships": all_relations}
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)

        cost = time.time() - t0
        logger.info(
            f"[OK] {file_name}  实体:{len(all_entities)}  关系:{len(all_relations)}  耗时:{cost:.1f}s"
        )
        return {
            "file": file_name,
            "status": "ok",
            "entity_count": len(all_entities),
            "rel_count": len(all_relations),
        }

    except Exception as e:
        logger.error(f"[ERROR] {file_name}: {e}\n{traceback.format_exc()}")
        return {"file": file_name, "status": "error", "error": str(e)}


def _extract_worker(file_path: str, output_dir: str, result_queue: Queue):
    """子进程入口：抽取单文件，结果放入队列"""
    try:
        result = extract_single_file(file_path, output_dir)
        result_queue.put(result)
    except Exception as e:
        result_queue.put({"file": os.path.basename(file_path), "status": "error", "error": str(e)})


def phase_extract(folder: str, output_dir: str, workers: int):
    """阶段1：多进程抽取，每个文件独立进程，超时后真正 kill"""
    files = [
        os.path.join(folder, f)
        for f in os.listdir(folder)
        if os.path.isfile(os.path.join(folder, f))
        and os.path.splitext(f)[1].lower() in SUPPORTED_EXT
    ]
    logger.info(f"发现 {len(files)} 个文件，并发数: {workers}")

    results = []
    # 按 workers 大小分批，每批同时跑 workers 个进程
    for i in range(0, len(files), workers):
        batch = files[i:i + workers]
        procs = []
        for fp in batch:
            q = Queue()
            p = Process(target=_extract_worker, args=(fp, output_dir, q))
            p.start()
            procs.append((p, fp, q))

        for p, fp, q in procs:
            file_name = os.path.basename(fp)
            p.join(timeout=EXTRACT_TIMEOUT)
            if p.is_alive():
                p.kill()
                p.join()
                logger.warning(f"[TIMEOUT] {file_name} 超过 {EXTRACT_TIMEOUT}s，已跳过")
                _record_fault(file_name)
                results.append({"file": file_name, "status": "timeout"})
            else:
                try:
                    results.append(q.get_nowait())
                except Exception:
                    results.append({"file": file_name, "status": "error"})

    ok      = sum(1 for r in results if r["status"] == "ok")
    skip    = sum(1 for r in results if r["status"] == "skip")
    error   = sum(1 for r in results if r["status"] == "error")
    timeout = sum(1 for r in results if r["status"] == "timeout")
    logger.info(f"抽取完成  成功:{ok}  跳过:{skip}  失败:{error}  超时:{timeout}")
    return results


# ════════════════════════════════════════════════════════════════════════════
# 阶段 2：汇总 → 去重 → Embedding → 写 Neo4j
# ════════════════════════════════════════════════════════════════════════════

def phase_merge(output_dir: str):
    """阶段2：读所有 JSON → 去重 → 向量化 → 写 Neo4j"""

    # ── 1. 读取所有抽取结果 ──────────────────────────────────────────────────
    all_entities:  List[Dict] = []
    all_relations: List[Dict] = []

    json_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]
    logger.info(f"读取 {len(json_files)} 个抽取结果文件")

    for fname in json_files:
        fpath = os.path.join(output_dir, fname)
        try:
            with open(fpath, "r", encoding="utf-8") as f:
                data = json.load(f)
            all_entities.extend(data.get("entities", []))
            all_relations.extend(data.get("relationships", []))
        except Exception as e:
            logger.warning(f"读取失败 {fname}: {e}")

    logger.info(f"汇总: 实体 {len(all_entities)}  关系 {len(all_relations)}")

    # ── 2. 去重（filter_node_relationship 只跑这一次） ───────────────────────
    merged_input = {"entities": all_entities, "relationships": all_relations}
    logger.info("开始全局去重（filter_node_relationship）...")
    t0 = time.time()
    try:
        filter_data = filter_node_relationship(
            URI      = NEO4J_URI,
            USERNAME = NEO4J_USER,
            PASSWORD = NEO4J_PASSWORD,
            merged_output_data = merged_input,
        )
    except Exception as e:
        logger.error(f"去重失败，使用原始数据: {e}")
        filter_data = merged_input
    logger.info(f"去重完成  实体:{len(filter_data['entities'])}  关系:{len(filter_data['relationships'])}  耗时:{time.time()-t0:.1f}s")

    # ── 3. 生成 Embedding ────────────────────────────────────────────────────
    logger.info("生成 Embedding...")
    t0 = time.time()
    embedded = 0
    for entity in filter_data["entities"]:
        name = entity.get("properties", {}).get("名称", "")
        if not name or "未知" in name:
            continue
        try:
            entity["properties"]["embedding"] = OpenaiAPI.generate_embedding(name)
            embedded += 1
        except Exception as e:
            logger.warning(f"Embedding 失败 [{name}]: {e}")
    logger.info(f"Embedding 完成  {embedded} 个实体  耗时:{time.time()-t0:.1f}s")

    # ── 4. 写入 Neo4j ─────────────────────────────────────────────────────────
    logger.info("写入 Neo4j...")
    t0 = time.time()
    kg = MilitaryKnowledgeGraph(
        uri   = NEO4J_URI,
        auth  = (NEO4J_USER, NEO4J_PASSWORD),
        Ontology = ONTOLOGY,
    )
    try:
        entity_count, rel_count = kg.process_data(filter_data)
        logger.info(
            f"Neo4j 写入完成  实体:{entity_count}  关系:{rel_count}  耗时:{time.time()-t0:.1f}s"
        )
    finally:
        kg.close()


# ════════════════════════════════════════════════════════════════════════════
# 入口
# ════════════════════════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="批量知识图谱构建")
    parser.add_argument("--folder",     default="",           help="待处理文件夹（阶段1必填）")
    parser.add_argument("--output_dir", default="/app/kg_output", help="JSON 输出/读取目录")
    parser.add_argument("--phase",      default="all",        choices=["extract", "merge", "all"],
                        help="extract=只抽取  merge=只入库  all=全流程")
    parser.add_argument("--workers",    type=int, default=1,  help="阶段1并发线程数")
    args = parser.parse_args()

    if args.phase in ("extract", "all"):
        if not args.folder:
            parser.error("--phase extract/all 需要指定 --folder")
        phase_extract(args.folder, args.output_dir, args.workers)

    if args.phase in ("merge", "all"):
        phase_merge(args.output_dir)


def _fake_worker(file_path: str, output_dir: str, result_queue: Queue):
    """测试用子进程：file_B 睡 20s 模拟超时，其余 2s"""
    file_name = os.path.basename(file_path)
    sleep_time = 20 if file_name == "file_B" else 2
    time.sleep(sleep_time)
    out_path = os.path.join(output_dir, file_name + ".json")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump({"file": file_name}, f)
    result_queue.put({"file": file_name, "status": "ok"})


def _test_timeout():
    """
    模拟超时跳过逻辑，验证完删掉此函数及 _fake_worker、下面的调用。
    file_A 正常(2s)，file_B 超时(20s)，file_C 正常(2s)
    期望总耗时约 TIMEOUT+2s，file_B 不写 JSON
    """
    import tempfile

    global EXTRACT_TIMEOUT
    EXTRACT_TIMEOUT = 5  # 测试用 5 秒

    output_dir = tempfile.mkdtemp()
    logger.info(f"测试输出目录: {output_dir}")

    files = ["file_A", "file_B", "file_C"]
    results = []
    t_start = time.time()

    for i in range(0, len(files), 2):
        batch = files[i:i + 2]
        procs = []
        for f in batch:
            q = Queue()
            p = Process(target=_fake_worker, args=(f, output_dir, q))
            p.start()
            procs.append((p, f, q))

        for p, f, q in procs:
            p.join(timeout=EXTRACT_TIMEOUT)
            if p.is_alive():
                p.kill()
                p.join()
                logger.warning(f"[TIMEOUT] {f} 超时，进程已 kill，不写 JSON")
                _record_fault(f)
                results.append({"file": f, "status": "timeout"})
            else:
                results.append(q.get_nowait())

    logger.info(f"总耗时: {time.time()-t_start:.1f}s")
    logger.info(f"结果: {results}")
    logger.info(f"output_dir 中的文件: {os.listdir(output_dir)}")


if __name__ == "__main__":
    _test_timeout()
    # main()