#!/usr/bin/env python
# coding=utf-8
"""
将分享数据转换为 Neo4j 导入格式

图模型: (User) -[:SHARED {vid, ts}]-> (User)

用法:
    python export_neo4j.py output/05_图数据/20260111.csv
    python export_neo4j.py output/05_图数据/*.csv  # 多文件
"""

import sys
import csv
from pathlib import Path

def export_neo4j(input_files, output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    users = set()
    relationships = []

    # 读取所有输入文件
    for input_file in input_files:
        print(f"读取: {input_file}")
        with open(input_file, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            for row in reader:
                from_mid = row['from_mid']
                vid = row['vid']
                target_mid = row['target_mid']
                ts = row['ts']

                users.add(from_mid)
                users.add(target_mid)
                relationships.append((from_mid, target_mid, vid, ts))

    print(f"用户数: {len(users)}, 分享关系数: {len(relationships)}")

    # 1. 导出用户节点 (neo4j-admin 格式)
    users_file = output_dir / 'users.csv'
    with open(users_file, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(['mid:ID', ':LABEL'])
        for mid in users:
            writer.writerow([mid, 'User'])
    print(f"写入: {users_file}")

    # 2. 导出分享关系 (neo4j-admin 格式)
    rels_file = output_dir / 'shared.csv'
    with open(rels_file, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow([':START_ID', ':END_ID', 'vid', 'ts:long', ':TYPE'])
        for from_mid, target_mid, vid, ts in relationships:
            writer.writerow([from_mid, target_mid, vid, ts, 'SHARED'])
    print(f"写入: {rels_file}")

    # 3. 生成导入命令
    cmd = f"""# Neo4j 导入命令 (停止数据库后执行):
neo4j-admin database import full \\
    --nodes={users_file.absolute()} \\
    --relationships={rels_file.absolute()} \\
    --overwrite-destination \\
    neo4j
"""
    print(cmd)

    cmd_file = output_dir / 'import_cmd.sh'
    with open(cmd_file, 'w') as f:
        f.write(cmd.strip())
    print(f"写入: {cmd_file}")

if __name__ == '__main__':
    import glob

    if len(sys.argv) < 2:
        print("用法: python export_neo4j.py <csv文件或通配符>")
        sys.exit(1)

    input_files = []
    for pattern in sys.argv[1:]:
        input_files.extend(glob.glob(pattern))

    if not input_files:
        print(f"找不到文件: {sys.argv[1:]}")
        sys.exit(1)

    output_dir = Path(input_files[0]).parent / 'neo4j'
    export_neo4j(input_files, output_dir)
