#!/usr/bin/env python3
"""FAISS → LanceDB 迁移脚本

将 Python 织忆的 FAISS 索引转换为 Go 织忆的 LanceDB 格式。

用法:
    python3 scripts/migrate.py \
        --data-dir ~/projects/zhiyi/data \
        --output-dir /var/lib/zhiyi/data

要求:
    pip install lancedb numpy
"""

import argparse
import json
import os
import sys
import time

def parse_args():
    p = argparse.ArgumentParser(description='FAISS → LanceDB 迁移')
    p.add_argument('--data-dir', required=True, help='旧Python织忆data目录')
    p.add_argument('--output-dir', default='/var/lib/zhiyi/data', help='LanceDB输出目录')
    p.add_argument('--dry-run', action='store_true', help='只统计，不执行')
    return p.parse_args()

def main():
    args = parse_args()
    
    # 1. 读取旧数据
    distilled_file = os.path.join(args.data_dir, 'hermes-main', 'distilled', '2026-05.jsonl')
    episodes_file = os.path.join(args.data_dir, 'hermes-main', 'episodes', '2026-05.jsonl')
    vectors_file = os.path.join(args.data_dir, 'sbert_vectors.jsonl')
    meta_file = os.path.join(args.data_dir, 'sbert_meta.json')
    
    distilled = []
    episodes = []
    
    for fname in [distilled_file, episodes_file]:
        if os.path.exists(fname):
            with open(fname) as f:
                for line in f:
                    distilled.append(json.loads(line.strip()))
    
    # 2. 读取向量
    vectors = []
    meta = {}
    if os.path.exists(meta_file):
        with open(meta_file) as f:
            meta = json.load(f)
    if os.path.exists(vectors_file):
        with open(vectors_file) as f:
            for line in f:
                vectors.append(json.loads(line.strip()))
    
    total = len(distilled)
    vector_count = len(vectors)
    print(f"📊 已统计: {total} 条 distilled, {vector_count} 条向量")
    
    if args.dry_run:
        print("   (dry-run 模式，不执行)")
        return
    
    # 3. 写入 LanceDB
    try:
        import lancedb
        import numpy as np
    except ImportError:
        print("❌ 需要 lancedb: pip install lancedb numpy")
        sys.exit(1)
    
    os.makedirs(args.output_dir, exist_ok=True)
    db = lancedb.connect(args.output_dir)
    
    # 创建 memories 表
    try:
        table = db.create_table("memories", [{
            "id": "init",
            "agent_id": "system",
            "namespace": "default",
            "content": "init",
            "category": "system_fact",
            "vector": [0.0]*1024,
            "tier": "normal",
            "quality_score": 0.0,
            "recall_count": 0,
            "freshness": "fresh",
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
            "updated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
            "is_deleted": False,
        }])
    except Exception:
        table = db.open_table("memories")
    
    # 迁移数据
    migrated = 0
    batch = []
    
    for i, d in enumerate(distilled):
        content = d.get('content', '') or d.get('summary', '') or ''
        if not content.strip():
            continue
        
        vec = vectors[i] if i < len(vectors) else [0.0] * 1024
        if isinstance(vec, dict):
            vec = vec.get('vector', [0.0]*1024)
        
        # 确保1024维
        if len(vec) != 1024:
            if len(vec) < 1024:
                vec = vec + [0.0] * (1024 - len(vec))
            else:
                vec = vec[:1024]
        
        record = {
            "id": d.get('id', f'migrated_{i}'),
            "agent_id": d.get('agent_id', 'hermes-main'),
            "namespace": d.get('namespace', 'hermes-main'),
            "content": content[:2000],
            "category": d.get('category', 'general'),
            "vector": vec,
            "tier": "core" if d.get('importance', 0) >= 3 else "normal",
            "quality_score": float(d.get('importance', 0.5)) / 5.0,
            "recall_count": d.get('recall_count', 0),
            "freshness": "fresh",
            "created_at": d.get('created_at', time.strftime("%Y-%m-%dT%H:%M:%SZ")),
            "updated_at": d.get('updated_at', d.get('created_at', time.strftime("%Y-%m-%dT%H:%M:%SZ"))),
            "is_deleted": False,
        }
        batch.append(record)
        
        if len(batch) >= 100:
            table.add(batch)
            migrated += len(batch)
            batch = []
            print(f"  ✓ 已迁移 {migrated}/{total}")
    
    if batch:
        table.add(batch)
        migrated += len(batch)
    
    print(f"\n✅ 迁移完成: {migrated} 条记录 → {args.output_dir}")

if __name__ == '__main__':
    main()