#!/usr/bin/env python3
"""Batch import multiple Notion exports into brain.db (append mode)"""
import sqlite3, os, re

DB = "/Users/bingshuolingdianyuanhe/WorkBuddy/2026-06-04-13-44-41/persona-brain-db/brain.db"

EXPORTS = [
    ("人格记忆总索引", "/tmp/batch-import/人格记忆总索引/export"),
    ("光湖世界入口导航", "/tmp/batch-import/光湖世界入口导航/export"),
    ("曜冥纪元", "/tmp/batch-import/曜冥纪元/export"),
]

def parse_md(filepath):
    fn = os.path.basename(filepath).replace('.md', '')
    uuid = None
    m = re.search(r'([a-f0-9]{32})$', fn)
    if m:
        uuid = m.group(1)
    # Strip UUID suffix from title
    if uuid and fn.endswith(uuid):
        title = fn[:fn.rindex(uuid)].strip().rstrip(' ·').strip()
    else:
        title = fn
    with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
        raw = f.read()[:12000]
    return {"title": title, "uuid": uuid, "content": raw}

def process_export(name, content_dir):
    conn = sqlite3.connect(DB)
    c = conn.cursor()

    # Build title → uuid + dir → uuid mappings
    title_to_uuid = {}
    dir_to_uuid = {}

    for walk_root, dirs, files in os.walk(content_dir):
        for f in files:
            if not f.endswith('.md'): continue
            fpath = os.path.join(walk_root, f)
            page = parse_md(fpath)
            if not page['uuid'] or not page['title']: continue
            title_to_uuid[page['title'].lower().strip()] = page['uuid']

    for walk_root, dirs, files in os.walk(content_dir):
        for d in dirs:
            dpath = os.path.join(walk_root, d)
            dname = d.lower().strip()
            if dname in title_to_uuid:
                dir_to_uuid[dpath] = title_to_uuid[dname]
            else:
                for title, uuid in title_to_uuid.items():
                    if len(dname) > 4 and len(title) > 4:
                        if title.startswith(dname) or dname.startswith(title):
                            dir_to_uuid[dpath] = uuid
                            break

    # Import pages
    new_count = 0
    skip_count = 0
    pages = []

    for walk_root, dirs, files in sorted(os.walk(content_dir)):
        for f in sorted(files):
            if not f.endswith('.md'): continue
            fpath = os.path.join(walk_root, f)
            page = parse_md(fpath)
            if not page['uuid']: continue
            rel = os.path.relpath(fpath, content_dir)
            parent_id = dir_to_uuid.get(walk_root)
            pages.append({**page, "parent_id": parent_id, "path": rel})

    c.execute("BEGIN TRANSACTION")
    for i, p in enumerate(pages):
        c.execute("SELECT 1 FROM documents WHERE document_id=?", (p['uuid'],))
        if c.fetchone():
            skip_count += 1
            continue
        c.execute("""
            INSERT INTO documents (document_id, title, content, parent_id, path, content_type)
            VALUES (?, ?, ?, ?, ?, 'notion')
        """, (p['uuid'], p['title'], p['content'], p['parent_id'], p['path']))
        new_count += 1
        if new_count % 500 == 0:
            conn.commit()
            c.execute("BEGIN TRANSACTION")
            print(f"  {new_count} new...")
    conn.commit()

    c.execute("SELECT count(*) FROM documents")
    total = c.fetchone()[0]
    conn.close()
    print(f"  {name}: {new_count} new + {skip_count} skipped (总{total}页)")
    return new_count, skip_count

def main():
    total_new = 0
    total_skip = 0
    for name, path in EXPORTS:
        if not os.path.isdir(path):
            print(f"  SKIP {name}: {path} not found")
            continue
        new, skip = process_export(name, path)
        total_new += new
        total_skip += skip

    conn = sqlite3.connect(DB)
    c = conn.cursor()
    c.execute("SELECT count(*) FROM documents")
    final = c.fetchone()[0]
    c.execute("SELECT count(*) FROM documents WHERE parent_id IS NULL")
    roots = c.fetchone()[0]
    conn.close()
    print(f"\n总计: {total_new}新增 + {total_skip}跳过 => 数据库共{final}页({roots}入口)")
    import os as _os
    print(f"数据库大小: {_os.path.getsize(DB)/1024/1024:.1f}MB")

if __name__ == '__main__':
    main()