"""
数据迁移脚本：创建默认工作区，将所有现有数据归入默认工作区。

运行方式: cd backend && ./venv/Scripts/python scripts/migrate_add_workspace.py
"""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import uuid
from datetime import datetime
from app.core.database import SessionLocal, engine
from app.core.config import settings

def migrate():
    db = SessionLocal()
    try:
        # 1. 创建默认工作区
        from app.models.workspace import Workspace, WorkspaceMembership
        from app.models.user import User

        # 检查是否已有默认工作区
        existing = db.query(Workspace).filter(Workspace.is_default == 1).first()
        if existing:
            print(f"默认工作区已存在: {existing.id} ({existing.name})")
            default_ws_id = existing.id
            default_ws = existing
        else:
            # 找到第一个 admin 用户作为 owner
            admin_user = db.query(User).filter(User.role == "admin").first()
            if not admin_user:
                admin_user = db.query(User).first()
            if not admin_user:
                print("错误: 数据库中没有用户，请先创建管理员账户")
                return

            default_ws = Workspace(
                id=str(uuid.uuid4()),
                name="默认工作区",
                description="系统自动创建的默认工作区，包含所有历史数据",
                is_default=1,
                owner_id=admin_user.id,
                max_members=200,
                status="active",
                created_at=datetime.utcnow(),
                updated_at=datetime.utcnow(),
            )
            db.add(default_ws)
            db.flush()
            default_ws_id = default_ws.id
            print(f"已创建默认工作区: {default_ws_id} ({default_ws.name})")

        # 2. 将所有现有用户加入默认工作区
        all_users = db.query(User).all()
        added_count = 0
        for user in all_users:
            existing_membership = db.query(WorkspaceMembership).filter(
                WorkspaceMembership.workspace_id == default_ws_id,
                WorkspaceMembership.user_id == user.id,
            ).first()
            if not existing_membership:
                membership = WorkspaceMembership(
                    id=str(uuid.uuid4()),
                    workspace_id=default_ws_id,
                    user_id=user.id,
                    role="admin",  # 现有用户默认给admin角色
                    joined_at=datetime.utcnow(),
                )
                db.add(membership)
                added_count += 1

        db.flush()
        print(f"已将 {added_count} 个用户加入默认工作区")

        # 3. 回填所有 Tier 1 表的 workspace_id
        tables_to_backfill = [
            ("agents", "user_id"),
            ("workflows", "user_id"),
            ("goals", "creator_id"),
            ("knowledge_bases", "user_id"),
            ("data_sources", "user_id"),
            ("model_configs", "user_id"),
            ("tools", "user_id"),
            ("agent_schedules", "user_id"),
            ("workflow_templates", "user_id"),
            ("orchestration_templates", "user_id"),
            ("executions", None),  # No user_id, backfill all
            ("tasks", None),
        ]

        from sqlalchemy import text
        for table, _user_col in tables_to_backfill:
            result = db.execute(
                text(f"UPDATE {table} SET workspace_id = :ws_id WHERE workspace_id IS NULL"),
                {"ws_id": default_ws_id}
            )
            print(f"  {table}: 已回填 {result.rowcount} 行")

        db.commit()
        print("\n迁移完成!")
        print(f"  默认工作区: {default_ws_id}")
        print(f"  成员数: {len(all_users)}")
        print(f"  所有现有数据已归入默认工作区")

    except Exception as e:
        db.rollback()
        print(f"迁移失败: {e}")
        raise
    finally:
        db.close()


if __name__ == "__main__":
    migrate()