"""一次性脚本：将 SQLite 数据迁移到 MySQL"""
import sqlite3
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from config import DATABASE_URL

SQLITE_PATH = "airlabs.db"

# 按外键依赖顺序排列
TABLES = [
    "roles",
    "users",
    "projects",
    "project_milestones",
    "submissions",
    "submission_history",
    "ai_tool_costs",
    "ai_tool_cost_allocations",
    "outsource_costs",
    "cost_overrides",
    "overhead_costs",
]


def migrate():
    if not DATABASE_URL.startswith("mysql"):
        print("ERROR: DATABASE_URL 不是 MySQL，请检查 .env 配置")
        return

    # 连接 SQLite
    sqlite_conn = sqlite3.connect(SQLITE_PATH)
    sqlite_conn.row_factory = sqlite3.Row

    # 连接 MySQL — 先建表
    from database import Base, engine
    import models  # noqa: F401 — 确保所有模型已注册
    Base.metadata.create_all(bind=engine)
    print("[OK] MySQL tables created")

    mysql_engine = engine
    Session = sessionmaker(bind=mysql_engine)
    session = Session()

    try:
        for table in TABLES:
            rows = sqlite_conn.execute(f"SELECT * FROM {table}").fetchall()
            if not rows:
                print(f"  {table}: 0 rows (skip)")
                continue

            cols = rows[0].keys()
            col_list = ", ".join(cols)
            param_list = ", ".join(f":{c}" for c in cols)
            insert_sql = text(f"INSERT INTO {table} ({col_list}) VALUES ({param_list})")

            # 清空目标表（避免重复运行冲突）
            session.execute(text(f"DELETE FROM {table}"))

            for row in rows:
                data = dict(row)
                session.execute(insert_sql, data)

            session.commit()
            print(f"  {table}: {len(rows)} rows migrated")

        print("\n[DONE] Migration complete!")
    except Exception as e:
        session.rollback()
        print(f"\n[ERROR] Migration failed: {e}")
        raise
    finally:
        session.close()
        sqlite_conn.close()


if __name__ == "__main__":
    migrate()