#!/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 --phase load_snapshot --snapshot_path /path/to/my_custom_snapshot.json
   
  # 全流程一次跑完
  python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase all --workers 4
"""

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,
)
import entity_registry as _entity_registry
import importlib
from modelsAPI.model_api import OpenaiAPI
from kg_build.kg_save import MilitaryKnowledgeGraph
from kg_build.filter_node_relation import filter_node_relationship
from filename_proceess_and_kgquery import get_entity
from kg_build.extract_filtertext import get_filtertext_node_relation
from extract_text_node_relation import process_file_content, check_filename_ocr
import os
import json
import logging
import argparse
import asyncio
import traceback
import time
from datetime import datetime
from multiprocessing import Process, Queue
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict, Tuple, Any, Optional

EXTRACT_TIMEOUT = 600  # 单文件抽取超时秒数
EMBED_WORKERS = 8      # Embedding 并发线程数
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}")


# ── 项目内模块 ──────────────────────────────────────────────────────────────


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://192.168.2.12:57687")
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "zdht123@")
MINERU_INSTANCES = [
    # ("http://192.168.0.64:18000/analyze-pdf"),
    #("http://192.168.0.111:9977/analyze-pdf/"),
    #   ("http://192.168.0.111:9979/analyze-pdf"),
    ("http://192.168.2.12:59988/analyze-pdf/"),
]
PREFIX_URL = os.getenv("PREFIX_URL", "http://192.168.2.12:59085")

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, api_url: str = MINERU_INSTANCES[0]) -> Dict:
    """
    对单个文件做实体/关系抽取，结果保存到 output_dir/<filename>.json。
    已存在 JSON 则跳过（断点续跑）。
    返回 {"file": ..., "status": "ok"/"skip"/"error", "entity_count": ..., "rel_count": ...}
    """
    logger.info(f'[api_url] {api_url}')
    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=api_url,
                check_cancelled_callback=lambda: False,
                device=lower_entity_device,
                device_type=lower_entity_type,
            )

        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 = [], []
            logger.info("无需进行ocr解析，直接构建实体")
        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, api_url: str = MINERU_INSTANCES[0]):
    """子进程入口：抽取单文件，结果放入队列"""
    try:
        result = extract_single_file(file_path, output_dir, api_url)
        result_queue.put(result)
    except Exception as e:
        result_queue.put({"file": os.path.basename(file_path), "status": "error", "error": str(e)})


def _run_file_with_timeout(args: Tuple) -> Dict:
    """在线程中启动子进程并监控超时，实现滚动窗口并发。"""
    fp, output_dir, inst = args
    file_name = os.path.basename(fp)
    q = Queue()
    p = Process(target=_extract_worker, args=(fp, output_dir, q, inst))
    p.start()
    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)
        return {"file": file_name, "status": "timeout"}
    try:
        return q.get_nowait()
    except Exception:
        return {"file": file_name, "status": "error"}


def phase_extract(folder: str, output_dir: str, workers: int):
    """阶段1：滚动窗口多进程抽取，每完成一个立即补充下一个。"""
    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}")
    os.makedirs(output_dir, exist_ok=True)

    tasks = [
        (fp, output_dir, MINERU_INSTANCES[i % len(MINERU_INSTANCES)])
        for i, fp in enumerate(files)
    ]

    with ThreadPoolExecutor(max_workers=workers) as executor:
        results = list(executor.map(_run_file_with_timeout, tasks))

    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(f"生成 Embedding（并发 {EMBED_WORKERS} 线程）...")
    t0 = time.time()

    def _embed_one(entity: Dict) -> int:
        name = str(entity.get("properties", {}).get("名称", "") or "")
        if not name or "未知" in name:
            return 0
        try:
            entity["properties"]["embedding"] = OpenaiAPI.generate_embedding(name)
            return 1
        except Exception as e:
            logger.warning(f"Embedding 失败 [{name}]: {e}")
            return 0

    with ThreadPoolExecutor(max_workers=EMBED_WORKERS) as executor:
        embedded = sum(executor.map(_embed_one, filter_data["entities"]))
    logger.info(f"Embedding 完成  {embedded} 个实体  耗时:{time.time() - t0:.1f}s")

    # ── 3.5 本地保存带 Embedding 的实体关系 ──────────────────────────────────
    snapshot_path = os.path.join(os.path.dirname(output_dir) if output_dir != "/" else "/app", "kg_snapshot.json")
    try:
        with open(snapshot_path, "w", encoding="utf-8") as f:
            json.dump(filter_data, f, ensure_ascii=False, indent=2)
        logger.info(f"本地快照已保存: {snapshot_path}  实体:{len(filter_data['entities'])}  关系:{len(filter_data['relationships'])}")
    except Exception as e:
        logger.warning(f"本地快照保存失败: {e}")

    # ── 4. 写入 Neo4j ─────────────────────────────────────────────────────────
    logger.info("写入 Neo4j...")
    # ontology, ontology_desc, node_types = _build_ontology_config()

    # 从实体数据中提取动态类型，追加到静态 ontology 末尾（静态配置优先，数据类型补充）
    data_types = list(dict.fromkeys(
        e["type"] for e in filter_data["entities"] if e.get("type")
    ))
    merged_ontology = list(dict.fromkeys( data_types))
    # new_types = [t for t in data_types if t not in set(ontology)]
    # if new_types:
    #     logger.info(f"从实体数据补充新类型 {len(new_types)} 个: {new_types}")

    t0 = time.time()
    kg = MilitaryKnowledgeGraph(
        uri=NEO4J_URI,
        auth=(NEO4J_USER, NEO4J_PASSWORD),
        Ontology=merged_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 phase_load_snapshot(snapshot_path: str):
    """
    Directly read kg_snapshot.json and store entities/relations into Neo4j.
    Skips extraction and embedding generation steps.
    """
    if not os.path.exists(snapshot_path):
        logger.error(f"Snapshot file not found: {snapshot_path}")
        return

    logger.info(f"Loading snapshot from: {snapshot_path}")
    t0 = time.time()
    
    try:
        with open(snapshot_path, "r", encoding="utf-8") as f:
            filter_data = json.load(f)
        
        entities = filter_data.get("entities", [])
        relations = filter_data.get("relationships", [])
        logger.info(f"Loaded snapshot: {len(entities)} entities, {len(relations)} relations")
        
        if not entities and not relations:
            logger.warning("Snapshot contains no data to load.")
            return

        # Extract dynamic types from entities to ensure Ontology coverage
        data_types = list(dict.fromkeys(
            e["type"] for e in entities if e.get("type")
        ))
        merged_ontology = list(dict.fromkeys(data_types))
        
        logger.info(f"Initializing MilitaryKnowledgeGraph with ontology types: {merged_ontology[:5]}...") # Log first few for brevity
        
        kg = MilitaryKnowledgeGraph(
            uri=NEO4J_URI,
            auth=(NEO4J_USER, NEO4J_PASSWORD),
            Ontology=merged_ontology,
        )
        
        try:
            entity_count, rel_count = kg.process_data(filter_data)
            elapsed = time.time() - t0
            logger.info(f"Neo4j write completed from snapshot. Entities: {entity_count}, Relations: {rel_count}. Time: {elapsed:.1f}s")
        finally:
            kg.close()
            
    except json.JSONDecodeError:
        logger.error(f"Failed to decode JSON from {snapshot_path}")
    except Exception as e:
        logger.error(f"Error loading snapshot: {e}", exc_info=True)

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


def main():
    parser = argparse.ArgumentParser(description="批量知识图谱构建")
    parser.add_argument("--folder", default="", help="待处理文件夹（阶段1必填）")
    parser.add_argument("--output_dir", default="/app/kg_output", help="JSON 输出/读取目录")
    # Update choices to include 'load_snapshot'
    parser.add_argument("--phase", default="all", choices=["extract", "merge", "load_snapshot", "all"], 
                        help="extract=只抽取  merge=只入库(含抽取结果汇总)  load_snapshot=直接加载快照  all=全流程")
    parser.add_argument("--workers", type=int, default=1, help="阶段1并发线程数")
    parser.add_argument("--snapshot_path", default="", help="快照文件路径 (用于 load_snapshot 阶段)")
    args = parser.parse_args()

    # Determine default snapshot path if not provided and phase is load_snapshot
    if args.phase == "load_snapshot" and not args.snapshot_path:
        # Default to the location where phase_merge saves it
        args.snapshot_path = os.path.join(os.path.dirname(args.output_dir) if args.output_dir != "/" else "/app", "kg_snapshot.json")

    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)
        
    if args.phase == "load_snapshot":
        phase_load_snapshot(args.snapshot_path)


if __name__ == "__main__":
    main()
# ontology, ontology_desc, node_types = _build_ontology_config()
# print(ontology)