#!/usr/bin/env python3
"""
步骤 3：批量更新 Neo4j 中所有节点 knowledge_source 字段里的旧服务 URL
在【迁移目标机】上执行本脚本
配置项请修改 migration_config.py
"""

import json
import sys
import os

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

from neo4j import GraphDatabase


def update_urls(driver, old_url: str, new_url: str):
    updated = 0
    skipped = 0
    errors = 0

    with driver.session() as session:
        result = session.run(
            "MATCH (n) WHERE n.knowledge_source IS NOT NULL "
            "RETURN id(n) AS node_id, n.名称 AS name, n.knowledge_source AS ks"
        )
        records = list(result)

    print(f"共找到 {len(records)} 个含 knowledge_source 的节点")

    for record in records:
        node_id = record["node_id"]
        name    = record["name"]
        ks_raw  = record["ks"]

        try:
            if isinstance(ks_raw, str):
                ks_list = json.loads(ks_raw)
            elif isinstance(ks_raw, list):
                ks_list = ks_raw
            else:
                skipped += 1
                continue
        except json.JSONDecodeError as e:
            print(f"  [WARN] 节点 {name}(id={node_id}) JSON 解析失败: {e}")
            errors += 1
            continue

        has_old_url = any(
            isinstance(item, dict) and old_url in item.get("url", "")
            for item in ks_list
        )
        if not has_old_url:
            skipped += 1
            continue

        for item in ks_list:
            if isinstance(item, dict) and "url" in item:
                item["url"] = item["url"].replace(old_url, new_url)

        new_ks_str = json.dumps(ks_list, ensure_ascii=False)

        with driver.session() as session:
            session.run(
                "MATCH (n) WHERE id(n) = $node_id SET n.knowledge_source = $ks",
                node_id=node_id,
                ks=new_ks_str,
            )

        print(f"  [OK] 更新节点: {name}")
        updated += 1

    return updated, skipped, errors


def main():
    print("==========================================")
    print("  kgrag Neo4j URL 批量更新脚本 - 步骤 3")
    print(f"  旧 URL 前缀: {OLD_BASE_URL}")
    print(f"  新 URL 前缀: {NEW_BASE_URL}")
    print(f"  Neo4j: {NEO4J_URI}")
    print("==========================================")

    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
    try:
        updated, skipped, errors = update_urls(driver, OLD_BASE_URL, NEW_BASE_URL)
    finally:
        driver.close()

    print("")
    print("==========================================")
    print(f"  完成！更新: {updated}  跳过: {skipped}  失败: {errors}")
    print("  下一步：执行 bash 04_update_config.sh")
    print("==========================================")


if __name__ == "__main__":
    main()