scripts/migrate_to_main.py

#!/usr/bin/env python3
"""
B-11 migration: 清测试数据 + 导入 0603.md 到 thread='main'

跑法（dry-run、不真动 DB）：
    DATABASE_URL=postgresql://... python3 scripts/migrate_to_main.py

跑法（真执行）：
    DATABASE_URL=postgresql://... python3 scripts/migrate_to_main.py --commit

数据清理范围（[[zhiyu-migration-project]] B-11 决议）：
- conversations 表 thread_id='test' 全删
- memories 表 created_at >= '2026-05-29' 全删（保留 chat 端真历史）
- mido_messages / dreams / messages_board / diary_entries / images 全删（全是测试期产生）

0603.md 导入：
- thread_id='main'
- created_at 从 2026-06-03 06:00 起、每 turn +1 分钟（合并算作"一天"）
- session_id='b11-migration-0603' 统一
- 思维链入 metadata.thinking JSON 字段
- 用 `|` 分隔思维链和正文
"""

import asyncio
import asyncpg
import os
import sys
import re
import json
from datetime import datetime, timezone, timedelta
from pathlib import Path

DATABASE_URL = os.environ.get('DATABASE_URL', '')
ZERO3_PATH = Path(__file__).parent.parent.parent / 'memory' / 'chat_history' / '0603.md'

DRY_RUN = '--commit' not in sys.argv

BJ_TZ = timezone(timedelta(hours=8))
BASE_TIME = datetime(2026, 6, 3, 6, 0, 0, tzinfo=BJ_TZ)
TURN_INTERVAL_MIN = 1
SESSION_ID = 'b11-migration-0603'


def parse_0603(text: str) -> list:
    """解析 0603.md、返回 [{role, content, thinking}, ...]

    格式：
        2026.06.03
        **User:**
        ...
        ---
        **Claude:**
        思维链多段...
        最后一行思维链 |

        正文多段...
        ---
    """
    text = text.lstrip()
    if text.startswith('2026.06.03'):
        text = text[len('2026.06.03'):].lstrip()

    raw_turns = re.split(r'\n---\n', text)
    turns = []
    for raw in raw_turns:
        raw = raw.strip()
        if not raw:
            continue
        if raw.startswith('**User:**'):
            content = raw[len('**User:**'):].strip()
            turns.append({'role': 'user', 'content': content, 'thinking': None})
        elif raw.startswith('**Claude:**'):
            content = raw[len('**Claude:**'):].strip()
            if '|' in content:
                idx = content.rfind('|')
                thinking = content[:idx].rstrip()
                body = content[idx + 1:].lstrip()
            else:
                thinking = None
                body = content
            turns.append({'role': 'assistant', 'content': body, 'thinking': thinking})
    return turns


async def clean_test_data(conn, dry_run=True):
    operations = [
        ("conversations.thread_id='test'",
         "SELECT COUNT(*) FROM conversations WHERE thread_id = 'test'",
         "DELETE FROM conversations WHERE thread_id = 'test'"),
        ("memories created_at >= '2026-05-29'",
         "SELECT COUNT(*) FROM memories WHERE created_at >= '2026-05-29'",
         "DELETE FROM memories WHERE created_at >= '2026-05-29'"),
        ("mido_messages 全部",
         "SELECT COUNT(*) FROM mido_messages",
         "DELETE FROM mido_messages"),
        ("dreams 全部",
         "SELECT COUNT(*) FROM dreams",
         "DELETE FROM dreams"),
        ("messages_board 全部",
         "SELECT COUNT(*) FROM messages_board",
         "DELETE FROM messages_board"),
        ("diary_entries 全部",
         "SELECT COUNT(*) FROM diary_entries",
         "DELETE FROM diary_entries"),
        ("images 全部",
         "SELECT COUNT(*) FROM images",
         "DELETE FROM images"),
    ]

    for label, count_sql, delete_sql in operations:
        try:
            count = await conn.fetchval(count_sql)
        except Exception as e:
            print(f"  [SKIP] {label}: {e!r}")
            continue
        if dry_run:
            print(f"  [DRY] {label}: 会删 {count} 条")
        else:
            res = await conn.execute(delete_sql)
            print(f"  [DEL] {label}: {res}（之前 {count} 条）")


async def import_0603(conn, turns, dry_run=True):
    n_imported = 0
    for i, turn in enumerate(turns):
        created_at = BASE_TIME + timedelta(minutes=i * TURN_INTERVAL_MIN)
        metadata = json.dumps({'thinking': turn['thinking']}, ensure_ascii=False) if turn['thinking'] else None

        if dry_run:
            preview = turn['content'][:80].replace('\n', ' ')
            if len(turn['content']) > 80:
                preview += '...'
            t_mark = '✦' if turn['thinking'] else ' '
            print(f"  [DRY] {created_at.strftime('%H:%M')} {t_mark} {turn['role']:>9}: {preview}")
        else:
            await conn.execute("""
                INSERT INTO conversations (session_id, role, content, model, metadata, thread_id, created_at)
                VALUES ($1, $2, $3, $4, $5, $6, $7)
            """, SESSION_ID, turn['role'], turn['content'], 'claude.ai-export', metadata, 'main', created_at)
            n_imported += 1
    return n_imported


async def main():
    if not DATABASE_URL:
        print("⚠️  DATABASE_URL 没设、跑不了。export DATABASE_URL=postgresql://... 再来", file=sys.stderr)
        sys.exit(1)
    if not ZERO3_PATH.exists():
        print(f"⚠️  找不到 {ZERO3_PATH}", file=sys.stderr)
        sys.exit(1)

    text = ZERO3_PATH.read_text(encoding='utf-8')
    turns = parse_0603(text)
    n_user = sum(1 for t in turns if t['role'] == 'user')
    n_asst = sum(1 for t in turns if t['role'] == 'assistant')
    n_with_thinking = sum(1 for t in turns if t.get('thinking'))

    print(f"=== 解析 0603.md ===")
    print(f"  总 turn 数：{len(turns)}（user {n_user} / assistant {n_asst}）")
    print(f"  带思维链：{n_with_thinking} 条")
    if turns:
        print(f"  第一条：{turns[0]['role']} - {turns[0]['content'][:60]!r}")
        print(f"  最后一条：{turns[-1]['role']} - {turns[-1]['content'][:60]!r}")
    print(f"  时间戳：{BASE_TIME.strftime('%Y-%m-%d %H:%M')} 起、每 turn +{TURN_INTERVAL_MIN} 分钟")
    print(f"  最后一条会落在 {(BASE_TIME + timedelta(minutes=(len(turns)-1) * TURN_INTERVAL_MIN)).strftime('%H:%M')}")
    print()

    if DRY_RUN:
        print("=== DRY RUN（不会真删 / 真灌）===\n")
    else:
        print("=== COMMIT MODE（实际执行 + 不可逆）===\n")
        confirm = input("确认要真执行吗？输入 'YES' 继续: ")
        if confirm != 'YES':
            print("取消")
            return

    pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=2, statement_cache_size=0)
    try:
        async with pool.acquire() as conn:
            print("[1/2] 清测试数据")
            await clean_test_data(conn, dry_run=DRY_RUN)

            print(f"\n[2/2] 导入 0603.md 到 thread='main'")
            n = await import_0603(conn, turns, dry_run=DRY_RUN)
            if not DRY_RUN:
                print(f"  导入 {n} 条")
    finally:
        await pool.close()

    print(f"\n{'(dry-run、没动 DB)' if DRY_RUN else 'Done、DB 已变更'}")


if __name__ == '__main__':
    asyncio.run(main())