#!/usr/bin/env python3
"""
步骤 5：验证迁移结果
在【迁移目标机】上执行本脚本
配置项请修改 migration_config.py
"""

import os
import json
import sys

sys.path.insert(0, os.path.dirname(__file__))
from migration_config import (
    OLD_IP, DST_CODE_DIR, DST_FILES_DIR, NEW_BASE_URL,
    NEO4J_URI, NEO4J_USER, NEO4J_PASSWORD
)

from neo4j import GraphDatabase

URL_CHECK_SAMPLE = 5

PASS = "\033[92m[PASS]\033[0m"
FAIL = "\033[91m[FAIL]\033[0m"
WARN = "\033[93m[WARN]\033[0m"


def check_files():
    print("\n--- 1. 关键文件检查 ---")
    checks = [
        (DST_FILES_DIR,                                  "原始文档目录"),
        (f"{DST_FILES_DIR}/uploads_v1",                  "高亮 PDF 目录"),
        (f"{DST_CODE_DIR}/kg_output",                    "抽取结果 JSON 目录"),
        (f"{DST_CODE_DIR}/kg_snapshot.json",             "全量快照 kg_snapshot.json"),
        (f"{DST_CODE_DIR}/entity_registry.json",         "实体注册表"),
        (f"{DST_CODE_DIR}/app.py",                       "主程序 app.py"),
        (f"{DST_CODE_DIR}/config.py",                    "配置文件 config.py"),
    ]
    all_ok = True
    for path, desc in checks:
        if os.path.exists(path):
            if os.path.isdir(path):
                count = len(os.listdir(path))
                print(f"  {PASS} {desc}: {path}  ({count} 个文件)")
            else:
                size = os.path.getsize(path) / 1024 / 1024
                print(f"  {PASS} {desc}: {path}  ({size:.1f} MB)")
        else:
            print(f"  {FAIL} 缺失: {desc}: {path}")
            all_ok = False
    return all_ok


def check_neo4j():
    print("\n--- 2. Neo4j 图谱检查 ---")
    try:
        driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
        with driver.session() as session:
            node_count = session.run("MATCH (n) RETURN count(n) AS c").single()["c"]
            rel_count  = session.run("MATCH ()-[r]->() RETURN count(r) AS c").single()["c"]
            ks_count   = session.run(
                "MATCH (n) WHERE n.knowledge_source IS NOT NULL RETURN count(n) AS c"
            ).single()["c"]
        driver.close()
        print(f"  {PASS} Neo4j 连接正常")
        print(f"  {PASS} 节点数: {node_count}  关系数: {rel_count}  含溯源节点: {ks_count}")
        return True, ks_count
    except Exception as e:
        print(f"  {FAIL} Neo4j 连接失败: {e}")
        return False, 0


def check_old_ip(neo4j_ok, ks_count):
    print(f"\n--- 3. 旧 IP 残留检查（{OLD_IP}）---")

    config_files = ["config.py", "batch_kg_build.py", "app.py"]
    residual_found = False
    for fname in config_files:
        fpath = os.path.join(DST_CODE_DIR, fname)
        if not os.path.exists(fpath):
            continue
        with open(fpath, "r", encoding="utf-8", errors="ignore") as f:
            content = f.read()
        lines = [i+1 for i, l in enumerate(content.splitlines()) if OLD_IP in l]
        if lines:
            print(f"  {WARN} {fname} 第 {lines} 行仍含旧 IP")
            residual_found = True

    if neo4j_ok and ks_count > 0:
        try:
            driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
            with driver.session() as session:
                result = session.run(
                    f"MATCH (n) WHERE n.knowledge_source CONTAINS '{OLD_IP}' "
                    "RETURN count(n) AS c"
                )
                old_url_count = result.single()["c"]
            driver.close()
            if old_url_count > 0:
                print(f"  {WARN} Neo4j 中仍有 {old_url_count} 个节点含旧 URL，请重新执行步骤 3")
                residual_found = True
            else:
                print(f"  {PASS} Neo4j 中无旧 IP 残留")
        except Exception as e:
            print(f"  {WARN} 无法检查 Neo4j 中的旧 URL: {e}")

    if not residual_found:
        print(f"  {PASS} 配置文件中无旧 IP 残留")


def check_urls(ks_count):
    print(f"\n--- 4. 溯源 URL 可访问性抽查（抽查 {URL_CHECK_SAMPLE} 个）---")

    if ks_count == 0:
        print(f"  {WARN} Neo4j 不可用，跳过 URL 检查")
        return

    try:
        driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
        with driver.session() as session:
            result = session.run(
                "MATCH (n) WHERE n.knowledge_source IS NOT NULL "
                f"RETURN n.名称 AS name, n.knowledge_source AS ks LIMIT {URL_CHECK_SAMPLE}"
            )
            records = list(result)
        driver.close()
    except Exception as e:
        print(f"  {WARN} 无法读取 Neo4j 节点: {e}")
        return

    for record in records:
        name   = record["name"]
        ks_raw = record["ks"]
        try:
            ks_list = json.loads(ks_raw) if isinstance(ks_raw, str) else ks_raw
            url = ks_list[0].get("url", "") if ks_list else ""
            if not url:
                print(f"  {WARN} 节点 {name}: knowledge_source 无 url 字段")
                continue
            filename = url.split("/upload/")[-1].split("#")[0]
            local_path = os.path.join(DST_FILES_DIR, "uploads_v1", filename)
            if os.path.exists(local_path):
                print(f"  {PASS} {name}: 高亮 PDF 存在 → {filename}")
            else:
                print(f"  {FAIL} {name}: 高亮 PDF 缺失 → {local_path}")
        except Exception as e:
            print(f"  {WARN} 节点 {name} 检查失败: {e}")


def main():
    print("==========================================")
    print("  kgrag 迁移验证脚本 - 步骤 5")
    print(f"  代码目录: {DST_CODE_DIR}")
    print(f"  文件目录: {DST_FILES_DIR}")
    print(f"  新服务地址: {NEW_BASE_URL}")
    print("==========================================")

    files_ok         = check_files()
    neo4j_ok, ks_cnt = check_neo4j()
    check_old_ip(neo4j_ok, ks_cnt)
    check_urls(ks_cnt)

    print("\n==========================================")
    if files_ok and neo4j_ok:
        print("  总体结果：迁移基本完成，请检查上方 WARN/FAIL 项")
    else:
        print("  总体结果：存在关键问题，请修复 FAIL 项后重新验证")
    print("==========================================")


if __name__ == "__main__":
    main()