from __future__ import annotations

import uuid
from collections.abc import Callable, Container, Iterable, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from typing import (
    TYPE_CHECKING,
    Any,
    Self,
    cast,
    overload,
    override,
)

import attr
import sqlalchemy as sa
import trafaret as t
from sqlalchemy.dialects import postgresql as pgsql
from sqlalchemy.engine.row import Row
from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
from sqlalchemy.ext.asyncio import AsyncSession as SASession
from sqlalchemy.orm import (
    Mapped,
    foreign,
    joinedload,
    load_only,
    mapped_column,
    relationship,
    selectinload,
)
from sqlalchemy.sql.expression import true

from ai.backend.common import validators as tx
from ai.backend.common.config import agent_selector_config_iv
from ai.backend.common.types import (
    AgentSelectionStrategy,
    JSONSerializableMixin,
    SessionTypes,
)
from ai.backend.manager.data.scaling_group.types import ScalingGroupData
from ai.backend.manager.models.base import (
    GUID,
    Base,
    PydanticColumn,
    StructuredJSONObjectColumn,
)
from ai.backend.manager.models.group import resolve_group_name_or_id, resolve_groups
from ai.backend.manager.models.rbac import (
    AbstractPermissionContext,
    AbstractPermissionContextBuilder,
    DomainScope,
    ProjectScope,
    RBACModel,
    ScopeType,
    UserScope,
    get_predefined_roles_in_scope,
)
from ai.backend.manager.models.rbac.context import ClientContext
from ai.backend.manager.models.rbac.permission_defs import ScalingGroupPermission
from ai.backend.manager.models.scaling_group.types import FairShareScalingGroupSpec
from ai.backend.manager.models.types import QueryCondition
from ai.backend.manager.models.user import UserRole
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine

if TYPE_CHECKING:
    from ai.backend.manager.models.agent import AgentRow
    from ai.backend.manager.models.domain import DomainRow
    from ai.backend.manager.models.group import GroupRow
    from ai.backend.manager.models.keypair import KeyPairRow
    from ai.backend.manager.models.resource_preset import ResourcePresetRow
    from ai.backend.manager.models.session import SessionRow

__all__: Sequence[str] = (
    # table defs
    "scaling_groups",
    "ScalingGroupOpts",
    "ScalingGroupRow",
    "sgroups_for_domains",
    "sgroups_for_groups",
    "sgroups_for_keypairs",
    # functions
    "query_allowed_sgroups",
)


@attr.define(slots=True)
class ScalingGroupOpts(JSONSerializableMixin):
    allowed_session_types: list[SessionTypes] = attr.Factory(
        lambda: [
            SessionTypes.INTERACTIVE,
            SessionTypes.BATCH,
            SessionTypes.INFERENCE,
        ],
    )
    pending_timeout: timedelta = timedelta(seconds=0)
    config: Mapping[str, Any] = attr.field(factory=dict)

    # Scheduler has a dedicated database column to store its name,
    # but agent selector configuration is stored as a part of the scheduler_opts column.
    agent_selection_strategy: AgentSelectionStrategy = AgentSelectionStrategy.DISPERSED
    agent_selector_config: Mapping[str, Any] = attr.field(factory=dict)

    # Only used in the ConcentratedAgentSelector
    enforce_spreading_endpoint_replica: bool = False

    allow_fractional_resource_fragmentation: bool = True
    """If set to false, agent will refuse to start kernel when they are forced to fragment fractional resource request"""

    route_cleanup_target_statuses: list[str] = attr.field(factory=lambda: ["unhealthy"])
    """List of route statuses that should be automatically cleaned up. Valid values: healthy, unhealthy, degraded"""

    def to_json(self) -> dict[str, Any]:
        return {
            "allowed_session_types": [item.value for item in self.allowed_session_types],
            "pending_timeout": self.pending_timeout.total_seconds(),
            "config": self.config,
            "agent_selection_strategy": self.agent_selection_strategy,
            "agent_selector_config": self.agent_selector_config,
            "enforce_spreading_endpoint_replica": self.enforce_spreading_endpoint_replica,
            "allow_fractional_resource_fragmentation": self.allow_fractional_resource_fragmentation,
            "route_cleanup_target_statuses": self.route_cleanup_target_statuses,
        }

    @classmethod
    def from_json(cls, obj: Mapping[str, Any]) -> ScalingGroupOpts:
        return cls(**cls.as_trafaret().check(obj))

    @classmethod
    def as_trafaret(cls) -> t.Trafaret:
        return t.Dict({
            t.Key("allowed_session_types", default=["interactive", "batch"]): t.List(
                tx.Enum(SessionTypes), min_length=1
            ),
            t.Key("pending_timeout", default=0): tx.TimeDuration(allow_negative=False),
            # Each scheduler impl refers an additional "config" key.
            t.Key("config", default={}): t.Mapping(t.String, t.Any),
            t.Key("agent_selection_strategy", default=AgentSelectionStrategy.DISPERSED): tx.Enum(
                AgentSelectionStrategy
            ),
            t.Key("agent_selector_config", default={}): agent_selector_config_iv,
            t.Key("enforce_spreading_endpoint_replica", default=False): t.ToBool,
            t.Key("allow_fractional_resource_fragmentation", default=True): t.ToBool,
            t.Key("route_cleanup_target_statuses", default=["unhealthy"]): t.List(
                t.Enum("healthy", "unhealthy", "degraded")
            ),
        }).allow_extra("*")


# When scheduling, we take the union of allowed scaling groups for
# each domain, group, and keypair.


class ScalingGroupForDomainRow(Base):  # type: ignore[misc]
    __tablename__ = "sgroups_for_domains"
    id: Mapped[uuid.UUID] = mapped_column(
        "id", GUID, primary_key=True, server_default=sa.text("uuid_generate_v4()")
    )
    scaling_group: Mapped[str] = mapped_column(
        "scaling_group",
        sa.ForeignKey("scaling_groups.name", onupdate="CASCADE", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    domain: Mapped[str] = mapped_column(
        "domain",
        sa.ForeignKey("domains.name", onupdate="CASCADE", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    __table_args__ = (
        # constraint
        sa.UniqueConstraint("scaling_group", "domain", name="uq_sgroup_domain"),
    )
    sgroup_row: Mapped[ScalingGroupRow] = relationship(
        "ScalingGroupRow",
        back_populates="sgroup_for_domains_rows",
    )
    domain_row: Mapped[DomainRow] = relationship(
        "DomainRow",
        back_populates="sgroup_for_domains_rows",
    )


# For compatibility
sgroups_for_domains = ScalingGroupForDomainRow.__table__


class ScalingGroupForProjectRow(Base):  # type: ignore[misc]
    __tablename__ = "sgroups_for_groups"
    id: Mapped[uuid.UUID] = mapped_column(
        "id", GUID, primary_key=True, server_default=sa.text("uuid_generate_v4()")
    )
    scaling_group: Mapped[str] = mapped_column(
        "scaling_group",
        sa.ForeignKey("scaling_groups.name", onupdate="CASCADE", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    group: Mapped[uuid.UUID] = mapped_column(
        "group",
        sa.ForeignKey("groups.id", onupdate="CASCADE", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )

    __table_args__ = (
        # constraint
        sa.UniqueConstraint("scaling_group", "group", name="uq_sgroup_ugroup"),
    )
    sgroup_row: Mapped[ScalingGroupRow] = relationship(
        "ScalingGroupRow",
        back_populates="sgroup_for_groups_rows",
    )
    project_row: Mapped[GroupRow] = relationship(
        "GroupRow",
        back_populates="sgroup_for_groups_rows",
    )


# For compatibility
sgroups_for_groups = ScalingGroupForProjectRow.__table__


class ScalingGroupForKeypairsRow(Base):  # type: ignore[misc]
    __tablename__ = "sgroups_for_keypairs"
    id: Mapped[uuid.UUID] = mapped_column(
        "id", GUID, primary_key=True, server_default=sa.text("uuid_generate_v4()")
    )
    scaling_group: Mapped[str] = mapped_column(
        "scaling_group",
        sa.ForeignKey("scaling_groups.name", onupdate="CASCADE", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    access_key: Mapped[str] = mapped_column(
        "access_key",
        sa.ForeignKey("keypairs.access_key", onupdate="CASCADE", ondelete="CASCADE"),
        index=True,
        nullable=False,
    )
    __table_args__ = (
        # constraint
        sa.UniqueConstraint("scaling_group", "access_key", name="uq_sgroup_akey"),
    )
    sgroup_row: Mapped[ScalingGroupRow] = relationship(
        "ScalingGroupRow",
        back_populates="sgroup_for_keypairs_rows",
    )
    keypair_row: Mapped[KeyPairRow] = relationship(
        "KeyPairRow",
        back_populates="sgroup_for_keypairs_rows",
    )


# For compatibility
sgroups_for_keypairs = ScalingGroupForKeypairsRow.__table__


def _get_resource_preset_join_condition() -> Any:
    from ai.backend.manager.models.resource_preset import ResourcePresetRow

    return ScalingGroupRow.name == foreign(ResourcePresetRow.scaling_group_name)


class ScalingGroupRow(Base):  # type: ignore[misc]
    __tablename__ = "scaling_groups"
    name: Mapped[str] = mapped_column("name", sa.String(length=64), primary_key=True)
    description: Mapped[str | None] = mapped_column("description", sa.String(length=512))
    is_active: Mapped[bool | None] = mapped_column(
        "is_active", sa.Boolean, index=True, default=True
    )
    is_public: Mapped[bool] = mapped_column(
        "is_public", sa.Boolean, index=True, default=True, server_default=true(), nullable=False
    )
    created_at: Mapped[datetime | None] = mapped_column(
        "created_at", sa.DateTime(timezone=True), server_default=sa.func.now()
    )
    wsproxy_addr: Mapped[str | None] = mapped_column(
        "wsproxy_addr", sa.String(length=1024), nullable=True
    )
    wsproxy_api_token: Mapped[str | None] = mapped_column(
        "wsproxy_api_token", sa.String(length=128), nullable=True
    )
    driver: Mapped[str] = mapped_column("driver", sa.String(length=64), nullable=False)
    driver_opts: Mapped[dict[str, Any]] = mapped_column(
        "driver_opts", pgsql.JSONB(), nullable=False, default={}
    )
    scheduler: Mapped[str] = mapped_column("scheduler", sa.String(length=64), nullable=False)
    use_host_network: Mapped[bool] = mapped_column(
        "use_host_network", sa.Boolean, nullable=False, default=False
    )
    scheduler_opts: Mapped[ScalingGroupOpts] = mapped_column(
        "scheduler_opts",
        StructuredJSONObjectColumn(ScalingGroupOpts),
        nullable=False,
        default={},
    )
    fair_share_spec: Mapped[FairShareScalingGroupSpec | None] = mapped_column(
        "fair_share_spec",
        PydanticColumn(FairShareScalingGroupSpec),
        nullable=True,
        default=None,
    )

    sessions: Mapped[list[SessionRow]] = relationship("SessionRow", back_populates="scaling_group")
    agents: Mapped[list[AgentRow]] = relationship("AgentRow", back_populates="scaling_group_row")

    sgroup_for_domains_rows: Mapped[list[ScalingGroupForDomainRow]] = relationship(
        "ScalingGroupForDomainRow",
        back_populates="sgroup_row",
    )
    sgroup_for_groups_rows: Mapped[list[ScalingGroupForProjectRow]] = relationship(
        "ScalingGroupForProjectRow",
        back_populates="sgroup_row",
    )
    sgroup_for_keypairs_rows: Mapped[list[ScalingGroupForKeypairsRow]] = relationship(
        "ScalingGroupForKeypairsRow",
        back_populates="sgroup_row",
    )
    resource_preset_rows: Mapped[list[ResourcePresetRow]] = relationship(
        "ResourcePresetRow",
        back_populates="scaling_group_row",
        primaryjoin=_get_resource_preset_join_condition,
    )

    def to_dataclass(self) -> ScalingGroupData:
        """Convert Row to domain model data."""
        from ai.backend.manager.data.scaling_group.types import (
            ScalingGroupDriverConfig,
            ScalingGroupMetadata,
            ScalingGroupNetworkConfig,
            ScalingGroupSchedulerConfig,
            ScalingGroupSchedulerOptions,
            ScalingGroupStatus,
            SchedulerType,
        )

        return ScalingGroupData(
            name=self.name,
            status=ScalingGroupStatus(
                is_active=self.is_active if self.is_active is not None else True,
                is_public=self.is_public,
            ),
            metadata=ScalingGroupMetadata(
                description=self.description or "",
                created_at=self.created_at or datetime.now(UTC),
            ),
            network=ScalingGroupNetworkConfig(
                wsproxy_addr=self.wsproxy_addr or "",
                wsproxy_api_token=self.wsproxy_api_token or "",
                use_host_network=self.use_host_network,
            ),
            driver=ScalingGroupDriverConfig(
                name=self.driver,
                options=self.driver_opts,
            ),
            scheduler=ScalingGroupSchedulerConfig(
                name=SchedulerType(self.scheduler),
                options=ScalingGroupSchedulerOptions(
                    allowed_session_types=self.scheduler_opts.allowed_session_types,
                    pending_timeout=self.scheduler_opts.pending_timeout,
                    config=self.scheduler_opts.config,
                    agent_selection_strategy=self.scheduler_opts.agent_selection_strategy,
                    agent_selector_config=self.scheduler_opts.agent_selector_config,
                    enforce_spreading_endpoint_replica=self.scheduler_opts.enforce_spreading_endpoint_replica,
                    allow_fractional_resource_fragmentation=self.scheduler_opts.allow_fractional_resource_fragmentation,
                    route_cleanup_target_statuses=self.scheduler_opts.route_cleanup_target_statuses,
                ),
            ),
            fair_share_spec=self.fair_share_spec or FairShareScalingGroupSpec(),
        )

    @classmethod
    async def list_by_condition(
        cls,
        conditions: Iterable[QueryCondition],
        *,
        db: ExtendedAsyncSAEngine,
    ) -> list[Self]:
        stmt = sa.select(ScalingGroupRow)
        for cond in conditions:
            stmt = cond(stmt)
        async with db.begin_readonly_session() as db_session:
            return list((await db_session.scalars(stmt)).all())


def and_names(names: Iterable[str]) -> Callable[..., sa.sql.Select[Any]]:
    return lambda query_stmt: query_stmt.where(ScalingGroupRow.name.in_(names))


# For compatibility
scaling_groups = ScalingGroupRow.__table__


@dataclass
class ScalingGroupModel(RBACModel[ScalingGroupPermission]):
    name: str
    description: str | None
    is_active: bool
    is_public: bool
    created_at: datetime

    wsproxy_addr: str | None
    wsproxy_api_token: str | None
    driver: str
    driver_opts: dict[str, Any]
    scheduler: str
    use_host_network: bool
    scheduler_opts: ScalingGroupOpts

    orm_obj: ScalingGroupRow
    _permissions: frozenset[ScalingGroupPermission] = field(default_factory=frozenset)

    @property
    def permissions(self) -> Container[ScalingGroupPermission]:
        return self._permissions

    @classmethod
    def from_row(cls, row: ScalingGroupRow, permissions: Iterable[ScalingGroupPermission]) -> Self:
        return cls(
            name=row.name,
            description=row.description,
            is_active=row.is_active if row.is_active is not None else True,
            is_public=row.is_public,
            created_at=row.created_at or datetime.now(UTC),
            wsproxy_addr=row.wsproxy_addr,
            wsproxy_api_token=row.wsproxy_api_token,
            driver=row.driver,
            driver_opts=row.driver_opts,
            scheduler=row.scheduler,
            use_host_network=row.use_host_network,
            scheduler_opts=row.scheduler_opts,
            _permissions=frozenset(permissions),
            orm_obj=row,
        )


@overload
async def query_allowed_sgroups(
    db_conn: SAConnection,
    domain_name: str,
    group: uuid.UUID,
    access_key: str,
) -> Sequence[Row[Any]]: ...


@overload
async def query_allowed_sgroups(
    db_conn: SAConnection,
    domain_name: str,
    group: Iterable[uuid.UUID],
    access_key: str,
) -> Sequence[Row[Any]]: ...


@overload
async def query_allowed_sgroups(
    db_conn: SAConnection,
    domain_name: str,
    group: str,
    access_key: str,
) -> Sequence[Row[Any]]: ...


@overload
async def query_allowed_sgroups(
    db_conn: SAConnection,
    domain_name: str,
    group: Iterable[str],
    access_key: str,
) -> Sequence[Row[Any]]: ...


async def query_allowed_sgroups(
    db_conn: SAConnection,
    domain_name: str,
    group: uuid.UUID | Iterable[uuid.UUID] | str | Iterable[str],
    access_key: str,
) -> Sequence[Row[Any]]:
    query = sa.select(sgroups_for_domains).where(sgroups_for_domains.c.domain == domain_name)
    result = await db_conn.execute(query)
    from_domain = {row.scaling_group for row in result}

    group_ids: Iterable[uuid.UUID] = []
    match group:
        case uuid.UUID() | str():
            if group_id := await resolve_group_name_or_id(db_conn, domain_name, group):
                group_ids = [group_id]
            else:
                group_ids = []
        case list() | tuple() | set():
            group_ids = await resolve_groups(db_conn, domain_name, cast(Iterable[Any], group))
    from_group: set[str]
    if not group_ids:
        from_group = set()  # empty
    else:
        group_cond = sgroups_for_groups.c.group.in_(group_ids)
        query = sa.select(sgroups_for_groups).where(group_cond)
        result = await db_conn.execute(query)
        from_group = {row.scaling_group for row in result}

    query = sa.select(sgroups_for_keypairs).where(sgroups_for_keypairs.c.access_key == access_key)
    result = await db_conn.execute(query)
    from_keypair = {row.scaling_group for row in result}

    sgroups = from_domain | from_group | from_keypair
    query = (
        sa.select(scaling_groups)
        .where(
            (scaling_groups.c.name.in_(sgroups)) & (scaling_groups.c.is_active),
        )
        .order_by(scaling_groups.c.name)
    )
    result = await db_conn.execute(query)
    return [row for row in result]


ALL_SCALING_GROUP_PERMISSIONS: frozenset[ScalingGroupPermission] = frozenset([
    perm for perm in ScalingGroupPermission
])

OWNER_PERMISSIONS: frozenset[ScalingGroupPermission] = ALL_SCALING_GROUP_PERMISSIONS
ADMIN_PERMISSIONS: frozenset[ScalingGroupPermission] = frozenset({
    # Admin permissions
    ScalingGroupPermission.READ_ATTRIBUTE,
    # sub-scope permissions
    ScalingGroupPermission.AGENT_PERMISSIONS,
    ScalingGroupPermission.COMPUTE_SESSION_PERMISSIONS,
    ScalingGroupPermission.INFERENCE_SERVICE_PERMISSIONS,
    ScalingGroupPermission.STORAGE_HOST_PERMISSIONS,
})
MONITOR_PERMISSIONS: frozenset[ScalingGroupPermission] = frozenset({
    # Admin permissions
    ScalingGroupPermission.READ_ATTRIBUTE,
    # sub-scope permissions
    ScalingGroupPermission.AGENT_PERMISSIONS,
    ScalingGroupPermission.COMPUTE_SESSION_PERMISSIONS,
    ScalingGroupPermission.INFERENCE_SERVICE_PERMISSIONS,
    ScalingGroupPermission.STORAGE_HOST_PERMISSIONS,
})
PRIVILEGED_MEMBER_PERMISSIONS: frozenset[ScalingGroupPermission] = frozenset({
    ScalingGroupPermission.AGENT_PERMISSIONS,
    ScalingGroupPermission.COMPUTE_SESSION_PERMISSIONS,
    ScalingGroupPermission.INFERENCE_SERVICE_PERMISSIONS,
    ScalingGroupPermission.STORAGE_HOST_PERMISSIONS,
})
MEMBER_PERMISSIONS: frozenset[ScalingGroupPermission] = frozenset({
    ScalingGroupPermission.AGENT_PERMISSIONS,
    ScalingGroupPermission.COMPUTE_SESSION_PERMISSIONS,
    ScalingGroupPermission.INFERENCE_SERVICE_PERMISSIONS,
    ScalingGroupPermission.STORAGE_HOST_PERMISSIONS,
})

ScalingGroupToPermissionMap = Mapping[str, frozenset[ScalingGroupPermission]]

type WhereClauseType = sa.sql.expression.BinaryExpression[Any] | sa.sql.expression.BooleanClauseList


@dataclass
class ScalingGroupPermissionContext(AbstractPermissionContext[ScalingGroupPermission, str, str]):
    @property
    def sgroup_to_permissions_map(self) -> ScalingGroupToPermissionMap:
        return self.object_id_to_additional_permission_map

    @property
    def query_condition(self) -> WhereClauseType | None:
        cond: WhereClauseType | None = None

        def _OR_coalesce(
            base_cond: WhereClauseType | None,
            _cond: sa.sql.expression.BinaryExpression[Any],
        ) -> WhereClauseType:
            return base_cond | _cond if base_cond is not None else _cond

        if self.object_id_to_additional_permission_map:
            cond = _OR_coalesce(
                cond, ScalingGroupRow.name.in_(self.object_id_to_additional_permission_map.keys())
            )
        if self.object_id_to_overriding_permission_map:
            cond = _OR_coalesce(
                cond, ScalingGroupRow.name.in_(self.object_id_to_overriding_permission_map.keys())
            )
        return cond

    async def build_query(self) -> sa.sql.Select[Any] | None:
        cond = self.query_condition
        if cond is None:
            return None
        return sa.select(ScalingGroupRow).where(cond)

    async def calculate_final_permission(self, rbac_obj: str) -> frozenset[ScalingGroupPermission]:
        host_name = rbac_obj
        return self.object_id_to_additional_permission_map.get(host_name, frozenset())


class ScalingGroupPermissionContextBuilder(
    AbstractPermissionContextBuilder[ScalingGroupPermission, ScalingGroupPermissionContext]
):
    db_session: SASession

    def __init__(self, db_session: SASession) -> None:
        self.db_session = db_session

    @override
    async def calculate_permission(
        self,
        ctx: ClientContext,
        target_scope: ScopeType,
    ) -> frozenset[ScalingGroupPermission]:
        roles = await get_predefined_roles_in_scope(ctx, target_scope, self.db_session)
        permissions = await self._calculate_permission_by_predefined_roles(roles)
        permissions |= await self.apply_customized_role(ctx, target_scope)
        return permissions

    async def apply_customized_role(
        self,
        ctx: ClientContext,
        _target_scope: ScopeType,
    ) -> frozenset[ScalingGroupPermission]:
        if ctx.user_role == UserRole.SUPERADMIN:
            return ALL_SCALING_GROUP_PERMISSIONS
        return frozenset()

    @override
    async def build_ctx_in_system_scope(
        self,
        ctx: ClientContext,
    ) -> ScalingGroupPermissionContext:
        from ai.backend.manager.models.domain import DomainRow

        perm_ctx = ScalingGroupPermissionContext()
        _domain_query_stmt = sa.select(DomainRow).options(load_only(DomainRow.name))
        for row in await self.db_session.scalars(_domain_query_stmt):
            to_be_merged = await self.build_ctx_in_domain_scope(ctx, DomainScope(row.name))
            perm_ctx.merge(to_be_merged)
        return perm_ctx

    @override
    async def build_ctx_in_domain_scope(
        self,
        ctx: ClientContext,
        scope: DomainScope,
    ) -> ScalingGroupPermissionContext:
        from ai.backend.manager.models.domain import DomainRow

        permissions = await self.calculate_permission(ctx, scope)
        if not permissions:
            # User is not part of the domain.
            return ScalingGroupPermissionContext()

        stmt = (
            sa.select(DomainRow)
            .where(DomainRow.name == scope.domain_name)
            .options(selectinload(DomainRow.sgroup_for_domains_rows))
        )
        domain_row = await self.db_session.scalar(stmt)
        if domain_row is None:
            return ScalingGroupPermissionContext()
        scaling_groups = domain_row.sgroup_for_domains_rows
        return ScalingGroupPermissionContext(
            object_id_to_additional_permission_map={
                row.scaling_group: permissions for row in scaling_groups
            }
        )

    @override
    async def build_ctx_in_project_scope(
        self,
        ctx: ClientContext,
        scope: ProjectScope,
    ) -> ScalingGroupPermissionContext:
        from ai.backend.manager.models.group import GroupRow

        project_permissions = await self.calculate_permission(ctx, scope)
        if not project_permissions:
            # User is not part of the domain.
            return ScalingGroupPermissionContext()

        stmt = (
            sa.select(GroupRow)
            .where(GroupRow.id == scope.project_id)
            .options(selectinload(GroupRow.sgroup_for_groups_rows))
        )
        project_row = await self.db_session.scalar(stmt)
        if project_row is None:
            return ScalingGroupPermissionContext()
        scaling_groups = project_row.sgroup_for_groups_rows
        return ScalingGroupPermissionContext(
            object_id_to_additional_permission_map={
                row.scaling_group: project_permissions for row in scaling_groups
            }
        )

    @override
    async def build_ctx_in_user_scope(
        self,
        ctx: ClientContext,
        scope: UserScope,
    ) -> ScalingGroupPermissionContext:
        from ai.backend.manager.models.keypair import KeyPairRow
        from ai.backend.manager.models.user import UserRow

        user_permissions = await self.calculate_permission(ctx, scope)
        if not user_permissions:
            # User is not part of the domain.
            return ScalingGroupPermissionContext()

        stmt = (
            sa.select(UserRow)
            .where(UserRow.uuid == scope.user_id)
            .options(
                selectinload(UserRow.keypairs).options(
                    joinedload(KeyPairRow.sgroup_for_keypairs_rows)
                )
            )
        )
        user_row = cast(UserRow | None, await self.db_session.scalar(stmt))
        if user_row is None:
            return ScalingGroupPermissionContext()

        object_id_to_additional_permission_map: dict[str, frozenset[ScalingGroupPermission]] = {}
        for keypair in user_row.keypairs:
            scaling_groups = keypair.sgroup_for_keypairs_rows
            for sg in scaling_groups:
                if sg.scaling_group not in object_id_to_additional_permission_map:
                    object_id_to_additional_permission_map[sg.scaling_group] = user_permissions
        return ScalingGroupPermissionContext(
            object_id_to_additional_permission_map=object_id_to_additional_permission_map
        )

    @override
    @classmethod
    async def _permission_for_owner(
        cls,
    ) -> frozenset[ScalingGroupPermission]:
        return OWNER_PERMISSIONS

    @override
    @classmethod
    async def _permission_for_admin(
        cls,
    ) -> frozenset[ScalingGroupPermission]:
        return ADMIN_PERMISSIONS

    @override
    @classmethod
    async def _permission_for_monitor(
        cls,
    ) -> frozenset[ScalingGroupPermission]:
        return MONITOR_PERMISSIONS

    @override
    @classmethod
    async def _permission_for_privileged_member(
        cls,
    ) -> frozenset[ScalingGroupPermission]:
        return PRIVILEGED_MEMBER_PERMISSIONS

    @override
    @classmethod
    async def _permission_for_member(
        cls,
    ) -> frozenset[ScalingGroupPermission]:
        return MEMBER_PERMISSIONS


async def get_scaling_groups(
    target_scope: ScopeType,
    requested_permission: ScalingGroupPermission,
    sgroup_names: Iterable[str] | None = None,
    *,
    ctx: ClientContext,
    db_session: SASession,
) -> list[ScalingGroupModel]:
    ret: list[ScalingGroupModel] = []
    builder = ScalingGroupPermissionContextBuilder(db_session)
    permission_ctx = await builder.build(ctx, target_scope, requested_permission)
    cond = permission_ctx.query_condition
    if cond is None:
        return ret
    _stmt = sa.select(ScalingGroupRow).where(cond)
    if sgroup_names is not None:
        _stmt = _stmt.where(ScalingGroupRow.name.in_(sgroup_names))
    async for row in await db_session.stream_scalars(_stmt):
        permissions = await permission_ctx.calculate_final_permission(row)
        ret.append(ScalingGroupModel.from_row(row, permissions))
    return ret
