backend/app/database/seed.py

"""
Database Seed Script — Populates initial roles, permissions, and the super admin user.

Run after migrations:
    python -m app.database.seed

Or from the project root:
    python backend/app/database/seed.py
"""

import asyncio
import os
import sys
import uuid
import logging

# Ensure backend is on path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.database.session import get_session_factory, dispose_engine
from app.models.rbac import Role, Permission, RolePermission
from app.models.user import User, UserPreferences
from app.models.rbac import UserRole
from app.core.security import hash_password

logger = logging.getLogger(__name__)

# ── Role Definitions ────────────────────────────────────────────────

ROLES = [
    {
        "name": "super_admin",
        "display_name": "Super Admin",
        "description": "Full system access. Can manage all users, roles, and settings.",
        "hierarchy_level": 100,
        "is_system": True,
    },
    {
        "name": "admin",
        "display_name": "Admin",
        "description": "Administrative access. Can manage users and platform settings.",
        "hierarchy_level": 80,
        "is_system": True,
    },
    {
        "name": "manager",
        "display_name": "Manager",
        "description": "Team management. Can manage projects, models, and datasets.",
        "hierarchy_level": 60,
        "is_system": True,
    },
    {
        "name": "ml_engineer",
        "display_name": "ML Engineer",
        "description": "Build, train, and deploy ML models.",
        "hierarchy_level": 40,
        "is_system": True,
    },
    {
        "name": "data_scientist",
        "display_name": "Data Scientist",
        "description": "Explore data, create experiments, and build models.",
        "hierarchy_level": 30,
        "is_system": True,
    },
    {
        "name": "viewer",
        "display_name": "Viewer",
        "description": "Read-only access to projects, models, and datasets.",
        "hierarchy_level": 10,
        "is_system": True,
    },
]

# ── Permission Definitions ──────────────────────────────────────────

PERMISSIONS = [
    # Users
    ("users:read", "View Users", "View user profiles", "users", "read"),
    ("users:write", "Edit Users", "Edit user profiles", "users", "write"),
    ("users:delete", "Delete Users", "Soft-delete user accounts", "users", "delete"),
    # Projects
    ("projects:read", "View Projects", "View project details", "projects", "read"),
    ("projects:write", "Edit Projects", "Create and edit projects", "projects", "write"),
    ("projects:delete", "Delete Projects", "Delete projects", "projects", "delete"),
    # Models
    ("models:read", "View Models", "View ML models", "models", "read"),
    ("models:write", "Edit Models", "Create and edit ML models", "models", "write"),
    ("models:deploy", "Deploy Models", "Deploy ML models to production", "models", "deploy"),
    # Datasets
    ("datasets:read", "View Datasets", "View datasets", "datasets", "read"),
    ("datasets:write", "Edit Datasets", "Upload and edit datasets", "datasets", "write"),
    ("datasets:delete", "Delete Datasets", "Delete datasets", "datasets", "delete"),
    # Experiments
    ("experiments:read", "View Experiments", "View experiments", "experiments", "read"),
    ("experiments:write", "Edit Experiments", "Create and run experiments", "experiments", "write"),
    # Admin
    ("admin:dashboard", "Admin Dashboard", "Access admin dashboard", "admin", "dashboard"),
    ("admin:audit_logs", "Audit Logs", "View audit logs", "admin", "audit_logs"),
    # System
    ("system:settings", "System Settings", "Manage system settings", "system", "settings"),
]

# ── Role → Permission Mapping ──────────────────────────────────────

ROLE_PERMISSIONS = {
    "super_admin": ["*"],  # Gets all permissions
    "admin": [
        "users:read", "users:write", "users:delete",
        "projects:read", "projects:write", "projects:delete",
        "models:read", "models:write", "models:deploy",
        "datasets:read", "datasets:write", "datasets:delete",
        "experiments:read", "experiments:write",
        "admin:dashboard", "admin:audit_logs",
        "system:settings",
    ],
    "manager": [
        "users:read",
        "projects:read", "projects:write",
        "models:read", "models:write", "models:deploy",
        "datasets:read", "datasets:write",
        "experiments:read", "experiments:write",
    ],
    "ml_engineer": [
        "projects:read",
        "models:read", "models:write", "models:deploy",
        "datasets:read", "datasets:write",
        "experiments:read", "experiments:write",
    ],
    "data_scientist": [
        "projects:read",
        "models:read", "models:write",
        "datasets:read", "datasets:write",
        "experiments:read", "experiments:write",
    ],
    "viewer": [
        "projects:read",
        "models:read",
        "datasets:read",
        "experiments:read",
    ],
}


async def seed_roles(db: AsyncSession) -> dict[str, uuid.UUID]:
    """Seed system roles. Returns {role_name: role_id}."""
    role_map = {}

    for role_data in ROLES:
        result = await db.execute(
            select(Role).where(Role.name == role_data["name"])
        )
        existing = result.scalar_one_or_none()

        if existing:
            role_map[existing.name] = existing.id
            print(f"  Role '{existing.name}' already exists")
        else:
            role = Role(**role_data)
            db.add(role)
            await db.flush()
            role_map[role.name] = role.id
            print(f"  Created role '{role.name}' (level={role.hierarchy_level})")

    return role_map


async def seed_permissions(db: AsyncSession) -> dict[str, uuid.UUID]:
    """Seed permissions. Returns {codename: permission_id}."""
    perm_map = {}

    for codename, display_name, description, resource, action in PERMISSIONS:
        result = await db.execute(
            select(Permission).where(Permission.codename == codename)
        )
        existing = result.scalar_one_or_none()

        if existing:
            perm_map[existing.codename] = existing.id
        else:
            perm = Permission(
                codename=codename,
                display_name=display_name,
                description=description,
                resource=resource,
                action=action,
            )
            db.add(perm)
            await db.flush()
            perm_map[perm.codename] = perm.id
            print(f"  Created permission '{codename}'")

    return perm_map


async def seed_role_permissions(
    db: AsyncSession,
    role_map: dict[str, uuid.UUID],
    perm_map: dict[str, uuid.UUID],
):
    """Assign permissions to roles."""
    all_perm_codenames = list(perm_map.keys())

    for role_name, permission_codenames in ROLE_PERMISSIONS.items():
        role_id = role_map.get(role_name)
        if not role_id:
            continue

        # super_admin gets everything
        codenames = all_perm_codenames if "*" in permission_codenames else permission_codenames

        for codename in codenames:
            perm_id = perm_map.get(codename)
            if not perm_id:
                continue

            # Check if already assigned
            result = await db.execute(
                select(RolePermission).where(
                    RolePermission.role_id == role_id,
                    RolePermission.permission_id == perm_id,
                )
            )
            if result.scalar_one_or_none():
                continue

            db.add(RolePermission(role_id=role_id, permission_id=perm_id))

    await db.flush()
    print("  Role-permission assignments complete")


async def seed_super_admin(db: AsyncSession, role_map: dict[str, uuid.UUID]):
    """Create the initial super admin user if none exists."""

    # Check if any super admin exists
    result = await db.execute(
        select(User)
        .join(UserRole, UserRole.user_id == User.id)
        .join(Role, Role.id == UserRole.role_id)
        .where(Role.name == "super_admin")
    )
    if result.scalar_one_or_none():
        print("  Super admin already exists, skipping")
        return

    admin_email = os.environ.get("ADMIN_EMAIL", "admin@datavision.ai")
    admin_password = os.environ.get("ADMIN_PASSWORD", "Admin@DataVision2024!")

    hashed, algorithm = hash_password(admin_password)

    user = User(
        email=admin_email,
        hashed_password=hashed,
        password_hash_algorithm=algorithm,
        full_name="System Administrator",
        is_active=True,
        is_verified=True,
    )
    db.add(user)
    await db.flush()

    # Assign super_admin role
    super_admin_role_id = role_map.get("super_admin")
    if super_admin_role_id:
        db.add(UserRole(user_id=user.id, role_id=super_admin_role_id, is_active=True))

    # Create default preferences
    db.add(UserPreferences(user_id=user.id))

    await db.flush()
    print(f"  Created super admin: {admin_email}")
    print(f"  DEFAULT PASSWORD: {admin_password}")
    print(f"  >>> CHANGE THIS IN PRODUCTION! <<<")


async def run_seed():
    """Run all seed operations."""
    print("\n=== DataVision Database Seed ===\n")

    factory = get_session_factory()
    async with factory() as db:
        try:
            print("[1/4] Seeding roles...")
            role_map = await seed_roles(db)

            print("[2/4] Seeding permissions...")
            perm_map = await seed_permissions(db)

            print("[3/4] Assigning permissions to roles...")
            await seed_role_permissions(db, role_map, perm_map)

            print("[4/4] Creating super admin user...")
            await seed_super_admin(db, role_map)

            await db.commit()
            print("\n=== Seed complete ===\n")

        except Exception as e:
            await db.rollback()
            print(f"\nSeed failed: {e}")
            raise
        finally:
            await dispose_engine()


if __name__ == "__main__":
    # Import Permission model (needed for forward reference)
    from app.models.rbac import Permission  # noqa: F811
    asyncio.run(run_seed())