#!/usr/bin/env python3
"""Notion Full Export → SQLite brain.db importer"""
import sqlite3, os, re

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

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

def main():
    conn = sqlite3.connect(DB)
    c = conn.cursor()
    c.execute("DELETE FROM documents")
    conn.commit()

    # 第一遍: 建立目录→UUID映射
    # Notion export: page .md filename = "Title UUID.md", child directory = "Title/"
    # Directory name may NOT include UUID, only the title
    dir_to_uuid = {}
    # Build a mapping from title-only to uuid
    title_to_uuid = {}
    for walk_root, dirs, files in os.walk(CONTENT):
        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
            # Map this page's uuid to any subdirectory whose basename matches the title
            # Also check the full filename (without .md) as fallback
            title_to_uuid[page['title'].lower().strip()] = page['uuid']

    # Second pass on dirs: match dir names to page titles
    for walk_root, dirs, files in os.walk(CONTENT):
        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:
                # Try prefix match: directory name might be truncated
                for title, uuid in title_to_uuid.items():
                    if title.startswith(dname) or dname.startswith(title):
                        if len(dname) > 3 and len(title) > 3:
                            dir_to_uuid[dpath] = uuid
                            break

    # 第二遍: 导入所有页面
    pages = []
    for walk_root, dirs, files in sorted(os.walk(CONTENT)):
        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)
            depth = max(0, rel.count('/'))
            parent_id = dir_to_uuid.get(walk_root)
            pages.append({**page, "depth": depth, "parent_id": parent_id, "path": rel})

    # 批量插入
    c.execute("BEGIN TRANSACTION")
    for i, p in enumerate(pages):
        c.execute("""
            INSERT OR REPLACE INTO documents (document_id, title, content, parent_id, path, content_type)
            VALUES (?, ?, ?, ?, ?, 'notion')
        """, (p['uuid'], p['title'], p['content'], p['parent_id'], p['path']))
        if (i+1) % 1000 == 0:
            conn.commit()
            c.execute("BEGIN TRANSACTION")
            print(f"  {i+1}/{len(pages)}...")
    conn.commit()

    # 统计
    c.execute("SELECT count(*) FROM documents")
    total = c.fetchone()[0]
    c.execute("SELECT count(*) FROM documents WHERE parent_id IS NULL")
    roots = c.fetchone()[0]
    c.execute("SELECT count(DISTINCT parent_id) FROM documents WHERE parent_id IS NOT NULL")
    parents = c.fetchone()[0]

    print(f"\n导入完成: {total}页 ({roots}根页面, {parents}父节点)")

    c.execute("SELECT document_id, title FROM documents WHERE parent_id IS NULL ORDER BY title")
    for row in c.fetchall()[:30]:
        print(f"  [{row[0][:8]}..] {row[1][:70]}")

    c.execute("SELECT count(*) FROM sqlite_master WHERE type='table'")
    tables = c.fetchone()[0]
    conn.close()
    print(f"Done. {tables} tables total.")

if __name__ == '__main__':
    main()