#!/usr/bin/env python3
"""
按舰船号拆分 kg_snapshot.json、kg_output/、files/ 等目录
输出结果放在 OUTPUT_BASE_DIR 下，每艘舰一个子目录

文件名规则与舰船的对应关系只需修改下方「变量配置区」，逻辑不用动。
"""

import json
import os
import re
import shutil
import sys
from collections import defaultdict
from typing import Dict, List, Optional, Set

# 从统一配置文件读取路径
sys.path.insert(0, os.path.dirname(__file__))
from migration_config import DST_CODE_DIR, DST_FILES_DIR

# ===================== 变量配置区（路径已从 migration_config.py 读取） ========================

# 输入路径
SNAPSHOT_PATH  = os.path.join(DST_CODE_DIR, "kg_snapshot.json")   # 全量快照
KG_OUTPUT_DIR  = os.path.join(DST_CODE_DIR, "kg_output")          # 抽取结果 JSON 目录
FILES_DIR      = DST_FILES_DIR                                      # 原始文档目录
UPLOADS_V1_DIR = os.path.join(DST_FILES_DIR, "uploads_v1")        # 高亮 PDF 目录

# 输出根目录（每艘舰一个子目录）
OUTPUT_BASE_DIR = os.path.join(DST_CODE_DIR, "migration_split")

# --------------------------------------------------------------------------
# 舰船规则表：定义「哪些文件名属于哪艘舰」
#
# 每条规则是一个 dict，支持以下匹配方式（可以组合，任意一条命中即匹配）：
#   prefix   : 文件名以此开头（不含路径）
#   suffix   : 文件名以此结尾（含扩展名）
#   contains : 文件名包含此字符串
#   regex    : 正则表达式，匹配整个文件名（不含路径）
#
# 示例：
#   {"prefix": "122-"}         → 以 "122-" 开头
#   {"prefix": "163-"}         → 以 "163-" 开头
#   {"contains": "发动机"}     → 包含"发动机"
#   {"regex": r"^JZ\d+_.+"}   → 正则匹配
# --------------------------------------------------------------------------
SHIP_RULES: Dict[str, List[Dict]] = {
    "101舰": [
        {"prefix": "101-"},
    ],
    "122舰": [
        {"prefix": "122-"},
    ],
    "163舰": [
        {"prefix": "163-"},
    ],
    # 新增舰船示例（取消注释并修改）：
    # "200舰": [
    #     {"prefix": "200-"},
    #     {"contains": "某特殊前缀"},
    # ],
}

# 共享节点（同时属于多艘舰）的处理策略：
#   "all"      : 每艘涉及的舰都保留这个节点（knowledge_source 保持完整）
#   "majority" : 归属来源文件更多的那艘舰，其他舰不保留
SHARED_NODE_POLICY = "all"

# ==========================================================================


def match_ship(filename: str) -> Optional[str]:
    """
    根据 SHIP_RULES 判断文件名属于哪艘舰。
    返回舰船名，或 None（没有匹配）。
    第一条命中的规则生效（规则之间互斥时注意顺序）。
    """
    bare = os.path.basename(filename)
    for ship, rules in SHIP_RULES.items():
        for rule in rules:
            if "prefix" in rule and bare.startswith(rule["prefix"]):
                return ship
            if "suffix" in rule and bare.endswith(rule["suffix"]):
                return ship
            if "contains" in rule and rule["contains"] in bare:
                return ship
            if "regex" in rule and re.match(rule["regex"], bare):
                return ship
    return None


# --------------------------------------------------------------------------
# 快照拆分
# --------------------------------------------------------------------------

def _get_node_ships(entities: List[Dict]) -> Dict[str, Set[str]]:
    """返回 {节点名: {归属舰集合}}"""
    node_ships: Dict[str, Set[str]] = defaultdict(set)
    for e in entities:
        name = e.get("properties", {}).get("名称", "")
        if not name:
            continue
        ks_raw = e.get("properties", {}).get("knowledge_source", "")
        try:
            ks_list = json.loads(ks_raw) if isinstance(ks_raw, str) else (ks_raw or [])
        except Exception:
            ks_list = []
        for ks in ks_list:
            ship = match_ship(ks.get("filename", ""))
            if ship:
                node_ships[name].add(ship)
    return node_ships


def split_snapshot(snapshot_path: str, output_base: str):
    print(f"\n[快照拆分] 读取 {snapshot_path} ...")
    with open(snapshot_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    entities      = data.get("entities", [])
    relationships = data.get("relationships", [])
    print(f"  总节点: {len(entities)}  总关系: {len(relationships)}")

    node_ships = _get_node_ships(entities)

    # 统计共享情况
    all_ships = set(SHIP_RULES.keys())
    stats = defaultdict(int)
    for name, ships in node_ships.items():
        if len(ships) == 0:
            stats["无来源"] += 1
        elif len(ships) == 1:
            stats[next(iter(ships))] += 1
        else:
            stats["共享"] += 1
    print(f"  节点归属统计: { dict(stats) }")

    # 为每艘舰生成节点/关系子集
    for ship in SHIP_RULES:
        ship_entities = []
        for e in entities:
            name  = e.get("properties", {}).get("名称", "")
            ships = node_ships.get(name, set())

            if ship in ships:
                ship_entities.append(e)
            elif not ships and SHARED_NODE_POLICY == "all":
                # 无来源文件的节点（如顶层舰艇/系统节点）每艘舰都保留
                ship_entities.append(e)

        # 只保留两端节点都在本舰集合里的关系
        ship_node_names = {e.get("properties", {}).get("名称", "") for e in ship_entities}
        ship_rels = [
            r for r in relationships
            if r.get("from_entity") in ship_node_names
            and r.get("to_entity") in ship_node_names
        ]

        out_dir = os.path.join(output_base, ship)
        os.makedirs(out_dir, exist_ok=True)
        out_path = os.path.join(out_dir, "kg_snapshot.json")
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump({"entities": ship_entities, "relationships": ship_rels},
                      f, ensure_ascii=False, indent=2)
        print(f"  [{ship}] 节点:{len(ship_entities)}  关系:{len(ship_rels)}  → {out_path}")


# --------------------------------------------------------------------------
# kg_output 拆分（每个 .json 文件对应一个原始文档）
# --------------------------------------------------------------------------

def split_kg_output(kg_output_dir: str, output_base: str):
    if not os.path.isdir(kg_output_dir):
        print(f"\n[kg_output 拆分] 目录不存在，跳过: {kg_output_dir}")
        return

    print(f"\n[kg_output 拆分] 扫描 {kg_output_dir} ...")
    files = [f for f in os.listdir(kg_output_dir) if f.endswith(".json")]
    print(f"  共 {len(files)} 个文件")

    counts = defaultdict(int)
    unmatched = []

    for fname in files:
        # kg_output 文件名格式：原始文件名.json（含扩展名，如 122-xxx.pdf.json）
        # 去掉最后的 .json 拿到原始文件名
        original = fname[:-5] if fname.endswith(".json") else fname
        ship = match_ship(original)

        if ship:
            dst_dir = os.path.join(output_base, ship, "kg_output")
            os.makedirs(dst_dir, exist_ok=True)
            shutil.copy2(os.path.join(kg_output_dir, fname),
                         os.path.join(dst_dir, fname))
            counts[ship] += 1
        else:
            unmatched.append(fname)

    for ship, cnt in counts.items():
        print(f"  [{ship}] {cnt} 个 JSON")
    if unmatched:
        print(f"  [未匹配] {len(unmatched)} 个文件: {unmatched[:5]}")


# --------------------------------------------------------------------------
# files/ 原始文档拆分
# --------------------------------------------------------------------------

def split_files(files_dir: str, output_base: str):
    if not os.path.isdir(files_dir):
        print(f"\n[files 拆分] 目录不存在，跳过: {files_dir}")
        return

    print(f"\n[files 拆分] 扫描 {files_dir} ...")
    all_files = [f for f in os.listdir(files_dir) if os.path.isfile(os.path.join(files_dir, f))]
    print(f"  共 {len(all_files)} 个文件")

    counts = defaultdict(int)
    unmatched = []

    for fname in all_files:
        # files/ 里有两种格式：
        #   UUID_原文件名.pdf  → 去掉 UUID 前缀再匹配
        #   原文件名.pdf       → 直接匹配
        bare = fname
        # 尝试去掉 UUID 前缀（格式：8-4-4-4-12 字符的十六进制）
        uuid_pattern = r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}_(.+)$'
        m = re.match(uuid_pattern, fname, re.IGNORECASE)
        if m:
            bare = m.group(1)

        ship = match_ship(bare)
        if ship:
            dst_dir = os.path.join(output_base, ship, "files")
            os.makedirs(dst_dir, exist_ok=True)
            shutil.copy2(os.path.join(files_dir, fname),
                         os.path.join(dst_dir, fname))
            counts[ship] += 1
        else:
            unmatched.append(fname)

    for ship, cnt in counts.items():
        print(f"  [{ship}] {cnt} 个文件")
    if unmatched:
        print(f"  [未匹配] {len(unmatched)} 个: {unmatched[:5]}")


# --------------------------------------------------------------------------
# uploads_v1/ 高亮 PDF 拆分（通过快照里的 knowledge_source 反查归属）
# --------------------------------------------------------------------------

def split_uploads(snapshot_path: str, uploads_dir: str, output_base: str):
    if not os.path.isdir(uploads_dir):
        print(f"\n[uploads 拆分] 目录不存在，跳过: {uploads_dir}")
        return

    print(f"\n[uploads 拆分] 通过快照反查高亮 PDF 归属 ...")

    with open(snapshot_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    # 从快照里收集：高亮PDF文件名 → 归属舰船
    highlight_ship: Dict[str, str] = {}
    for e in data.get("entities", []):
        ks_raw = e.get("properties", {}).get("knowledge_source", "")
        try:
            ks_list = json.loads(ks_raw) if isinstance(ks_raw, str) else (ks_raw or [])
        except Exception:
            continue
        for ks in ks_list:
            url = ks.get("url", "")
            src_file = ks.get("filename", "")
            ship = match_ship(src_file)
            if not ship or not url:
                continue
            # URL 格式：http://host/upload/<highlight_filename>#page=N
            if "/upload/" in url:
                highlight_fname = url.split("/upload/")[-1].split("#")[0]
                # 共享节点下同一个高亮文件只记录第一个匹配舰船
                if highlight_fname not in highlight_ship:
                    highlight_ship[highlight_fname] = ship

    counts = defaultdict(int)
    missing = []

    for fname, ship in highlight_ship.items():
        src = os.path.join(uploads_dir, fname)
        if os.path.exists(src):
            dst_dir = os.path.join(output_base, ship, "files", "uploads_v1")
            os.makedirs(dst_dir, exist_ok=True)
            shutil.copy2(src, os.path.join(dst_dir, fname))
            counts[ship] += 1
        else:
            missing.append(fname)

    for ship, cnt in counts.items():
        print(f"  [{ship}] {cnt} 个高亮 PDF")
    if missing:
        print(f"  [本地缺失] {len(missing)} 个（需从服务器 uploads_v1 补全）: {missing[:3]}")


# --------------------------------------------------------------------------
# 入口
# --------------------------------------------------------------------------

def main():
    print("=" * 50)
    print("  kgrag 按舰船号拆分工具")
    print(f"  舰船规则: { {k: v for k, v in SHIP_RULES.items()} }")
    print(f"  共享节点策略: {SHARED_NODE_POLICY}")
    print("=" * 50)

    os.makedirs(OUTPUT_BASE_DIR, exist_ok=True)

    split_snapshot(SNAPSHOT_PATH, OUTPUT_BASE_DIR)
    split_kg_output(KG_OUTPUT_DIR, OUTPUT_BASE_DIR)
    split_files(FILES_DIR, OUTPUT_BASE_DIR)
    split_uploads(SNAPSHOT_PATH, UPLOADS_V1_DIR, OUTPUT_BASE_DIR)

    print("\n" + "=" * 50)
    print(f"  完成！输出目录: {OUTPUT_BASE_DIR}")
    print("  每艘舰的子目录结构：")
    print("    <舰名>/kg_snapshot.json  ← 导入新 Neo4j 用")
    print("    <舰名>/kg_output/        ← 抽取结果 JSON")
    print("    <舰名>/files/            ← 原始文档")
    print("    <舰名>/files/uploads_v1/ ← 高亮 PDF")
    print("=" * 50)


if __name__ == "__main__":
    main()
