calminer/services/bootstrap.py

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Callable

from config.settings import AdminBootstrapSettings
from models import User
from services.pricing import PricingMetadata
from services.repositories import (
    PricingSettingsSeedResult,
    ensure_default_roles,
)
from services.unit_of_work import UnitOfWork


logger = logging.getLogger(__name__)


@dataclass(slots=True)
class RoleBootstrapResult:
    created: int
    ensured: int


@dataclass(slots=True)
class AdminBootstrapResult:
    created_user: bool
    updated_user: bool
    password_rotated: bool
    roles_granted: int


@dataclass(slots=True)
class PricingBootstrapResult:
    seed: PricingSettingsSeedResult
    projects_assigned: int


def bootstrap_admin(
    *,
    settings: AdminBootstrapSettings,
    unit_of_work_factory: Callable[[], UnitOfWork] = UnitOfWork,
) -> tuple[RoleBootstrapResult, AdminBootstrapResult]:
    """Ensure default roles and administrator account exist."""

    with unit_of_work_factory() as uow:
        assert uow.roles is not None and uow.users is not None

        role_result = _bootstrap_roles(uow)
        admin_result = _bootstrap_admin_user(uow, settings)

    logger.info(
        "Admin bootstrap result: created_user=%s updated_user=%s password_rotated=%s roles_granted=%s",
        admin_result.created_user,
        admin_result.updated_user,
        admin_result.password_rotated,
        admin_result.roles_granted,
    )
    return role_result, admin_result


def _bootstrap_roles(uow: UnitOfWork) -> RoleBootstrapResult:
    assert uow.roles is not None
    before = {role.name for role in uow.roles.list()}
    ensure_default_roles(uow.roles)
    after = {role.name for role in uow.roles.list()}
    created = len(after - before)
    return RoleBootstrapResult(created=created, ensured=len(after))


def _bootstrap_admin_user(
    uow: UnitOfWork,
    settings: AdminBootstrapSettings,
) -> AdminBootstrapResult:
    assert uow.users is not None and uow.roles is not None

    created_user = False
    updated_user = False
    password_rotated = False
    roles_granted = 0

    user = uow.users.get_by_email(settings.email, with_roles=True)
    if user is None:
        user = User(
            email=settings.email,
            username=settings.username,
            password_hash=User.hash_password(settings.password),
            is_active=True,
            is_superuser=True,
        )
        uow.users.create(user)
        created_user = True
    else:
        if user.username != settings.username:
            user.username = settings.username
            updated_user = True
        if not user.is_active:
            user.is_active = True
            updated_user = True
        if not user.is_superuser:
            user.is_superuser = True
            updated_user = True
        if settings.force_reset:
            user.password_hash = User.hash_password(settings.password)
            password_rotated = True
            updated_user = True
        uow.users.session.flush()

    user = uow.users.get(user.id, with_roles=True)
    assert user is not None

    existing_roles = {role.name for role in user.roles}
    for role_name in settings.roles:
        role = uow.roles.get_by_name(role_name)
        if role is None:
            logger.warning(
                "Bootstrap admin role '%s' is not defined; skipping assignment",
                role_name,
            )
            continue
        if role.name in existing_roles:
            continue
        uow.users.assign_role(
            user_id=user.id,
            role_id=role.id,
            granted_by=user.id,
        )
        roles_granted += 1
        existing_roles.add(role.name)

    uow.users.session.flush()

    return AdminBootstrapResult(
        created_user=created_user,
        updated_user=updated_user,
        password_rotated=password_rotated,
        roles_granted=roles_granted,
    )


def bootstrap_pricing_settings(
    *,
    metadata: PricingMetadata,
    unit_of_work_factory: Callable[[], UnitOfWork] = UnitOfWork,
    default_slug: str = "default",
) -> PricingBootstrapResult:
    """Ensure baseline pricing settings exist and projects reference them."""

    with unit_of_work_factory() as uow:
        seed_result = uow.ensure_default_pricing_settings(
            metadata=metadata,
            slug=default_slug,
        )

        assigned = 0
        if uow.projects:
            default_settings = seed_result.settings
            projects = uow.projects.list(with_pricing=True)
            for project in projects:
                if project.pricing_settings is None:
                    uow.set_project_pricing_settings(project, default_settings)
                    assigned += 1

    logger.info(
        "Pricing bootstrap result: slug=%s created=%s updated_fields=%s impurity_upserts=%s projects_assigned=%s",
        seed_result.settings.slug,
        seed_result.created,
        seed_result.updated_fields,
        seed_result.impurity_upserts,
        assigned,
    )
    return PricingBootstrapResult(seed=seed_result, projects_assigned=assigned)