#!/usr/bin/env python3
"""
将图数据 CSV 转换为 Neo4j 导入格式

输入: output/06_图数据_v2/{dt}.csv
输出: output/neo4j/{dt}/
    ├── nodes_user.csv      # User 节点
    ├── nodes_video.csv     # Video 节点
    ├── nodes_share.csv     # Share 节点
    ├── rels_created.csv    # CREATED 关系 (User → Share)
    ├── rels_of.csv         # OF 关系 (Share → Video)
    └── rels_clicked.csv    # CLICKED 关系 (User → Share)

Neo4j 导入命令:
    neo4j-admin database import full \
        --nodes=User=nodes_user.csv \
        --nodes=Video=nodes_video.csv \
        --nodes=Share=nodes_share.csv \
        --relationships=CREATED=rels_created.csv \
        --relationships=OF=rels_of.csv \
        --relationships=CLICKED=rels_clicked.csv \
        neo4j
"""

import pandas as pd
from pathlib import Path
import argparse


def convert_to_neo4j(input_csv: Path, output_dir: Path):
    """将原始 CSV 转换为 Neo4j 导入格式"""

    print(f"读取: {input_csv}")
    df = pd.read_csv(input_csv)
    print(f"总记录数: {len(df)}")

    output_dir.mkdir(parents=True, exist_ok=True)

    # 分离 share 和 click 事件
    df_share = df[df['action_type'] == 'share'].copy()
    df_click = df[df['action_type'] == 'click'].copy()

    print(f"  share 事件: {len(df_share)}")
    print(f"  click 事件: {len(df_click)}")

    # ========== 节点文件 ==========

    # 1. User 节点 (收集所有 mid: actor_mid + sharer_mid)
    all_mids = set()
    all_mids.update(df['actor_mid'].dropna().unique())
    all_mids.update(df['sharer_mid'].dropna().unique())
    all_users = pd.Series(list(all_mids))

    users_df = pd.DataFrame({
        'mid:ID(User)': all_users
    })
    users_df.to_csv(output_dir / 'nodes_user.csv', index=False)
    print(f"  User 节点: {len(users_df)}")

    # 2. Video 节点
    all_vids = df['vid'].dropna().unique()
    videos_df = pd.DataFrame({
        'vid:ID(Video)': all_vids
    })
    videos_df.to_csv(output_dir / 'nodes_video.csv', index=False)
    print(f"  Video 节点: {len(videos_df)}")

    # 3. Share 节点 (只从 share 事件创建，包含元数据)
    shares_df = df_share[['shareid', 'rootshareid', 'sharedepth']].drop_duplicates('shareid')
    shares_df = shares_df.rename(columns={
        'shareid': 'shareid:ID(Share)',
        'rootshareid': 'rootshareid',
        'sharedepth': 'depth:int'
    })
    shares_df.to_csv(output_dir / 'nodes_share.csv', index=False)
    print(f"  Share 节点: {len(shares_df)}")

    # ========== 关系文件 ==========

    # 1. CREATED 关系 (User → Share)
    # share 事件: sharer_mid 创建了 shareid
    created_df = df_share[['sharer_mid', 'shareid', 'action_ts']].dropna(subset=['sharer_mid', 'shareid'])
    created_df = created_df.drop_duplicates(['sharer_mid', 'shareid'])
    created_df = created_df.rename(columns={
        'sharer_mid': ':START_ID(User)',
        'shareid': ':END_ID(Share)',
        'action_ts': 'action_ts:long'
    })
    created_df.to_csv(output_dir / 'rels_created.csv', index=False)
    print(f"  CREATED 关系: {len(created_df)}")

    # 2. OF 关系 (Share → Video)
    # 每个 shareid 对应一个 vid
    of_df = df_share[['shareid', 'vid']].dropna()
    of_df = of_df.drop_duplicates('shareid')
    of_df = of_df.rename(columns={
        'shareid': ':START_ID(Share)',
        'vid': ':END_ID(Video)'
    })
    of_df.to_csv(output_dir / 'rels_of.csv', index=False)
    print(f"  OF 关系: {len(of_df)}")

    # 3. CLICKED 关系 (User → Share)
    # click 事件: actor_mid 点击了 shareid
    clicked_df = df_click[['actor_mid', 'shareid', 'action_ts']].dropna(subset=['actor_mid', 'shareid'])
    # 同一个用户可能多次点击同一个 share，保留所有记录还是去重？
    # 这里先去重，如果需要保留多次点击可以注释下一行
    clicked_df = clicked_df.drop_duplicates(['actor_mid', 'shareid'])
    clicked_df = clicked_df.rename(columns={
        'actor_mid': ':START_ID(User)',
        'shareid': ':END_ID(Share)',
        'action_ts': 'action_ts:long'
    })
    clicked_df.to_csv(output_dir / 'rels_clicked.csv', index=False)
    print(f"  CLICKED 关系: {len(clicked_df)}")

    print(f"\n输出目录: {output_dir}")


def main():
    parser = argparse.ArgumentParser(description='将图数据 CSV 转换为 Neo4j 导入格式')
    parser.add_argument('input_csv', type=Path, help='输入 CSV 文件路径')
    parser.add_argument('-o', '--output', type=Path, help='输出目录 (默认: output/neo4j/{dt}/)')

    args = parser.parse_args()

    if not args.input_csv.exists():
        print(f"错误: 文件不存在 {args.input_csv}")
        return

    # 默认输出目录
    if args.output:
        output_dir = args.output
    else:
        # 从文件名提取日期
        dt = args.input_csv.stem  # e.g., "20260120"
        output_dir = Path('output/neo4j') / dt

    convert_to_neo4j(args.input_csv, output_dir)


if __name__ == '__main__':
    main()
