from __future__ import annotations

from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any, cast, override

import sqlalchemy as sa
from sqlalchemy.dialects import postgresql as pgsql
from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
from sqlalchemy.ext.asyncio import AsyncSession as SASession
from sqlalchemy.orm import (
    Mapped,
    joinedload,
    load_only,
    mapped_column,
    relationship,
    selectinload,
)
from sqlalchemy.sql.expression import false, true

from ai.backend.common.auth import PublicKey
from ai.backend.common.types import AccessKey, AgentId, ResourceSlot, SlotName, SlotTypes
from ai.backend.manager.data.agent.types import (
    AgentData,
    AgentDataForHeartbeatUpdate,
    AgentStatus,
)
from ai.backend.manager.models.base import (
    Base,
    CurvePublicKeyColumn,
    EnumType,
    ResourceSlotColumn,
)
from ai.backend.manager.models.kernel import KernelRow
from ai.backend.manager.models.keypair import KeyPairRow
from ai.backend.manager.models.rbac import (
    AbstractPermissionContext,
    AbstractPermissionContextBuilder,
    DomainScope,
    ProjectScope,
    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 AgentPermission, ScalingGroupPermission
from ai.backend.manager.models.resource_slot import AgentResourceRow
from ai.backend.manager.models.types import QueryCondition
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine, execute_with_txn_retry

if TYPE_CHECKING:
    from ai.backend.manager.models.scaling_group import ScalingGroupRow

__all__: Sequence[str] = (
    "AgentRow",
    "agents",
    "list_schedulable_agents_by_sgroup",
)


class AgentRow(Base):  # type: ignore[misc]
    __tablename__ = "agents"

    id: Mapped[str] = mapped_column("id", sa.String(length=64), primary_key=True)
    status: Mapped[AgentStatus] = mapped_column(
        "status", EnumType(AgentStatus), nullable=False, index=True, default=AgentStatus.ALIVE
    )
    status_changed: Mapped[datetime | None] = mapped_column(
        "status_changed", sa.DateTime(timezone=True), nullable=True
    )
    region: Mapped[str] = mapped_column("region", sa.String(length=64), index=True, nullable=False)
    scaling_group: Mapped[str] = mapped_column(
        "scaling_group",
        sa.ForeignKey("scaling_groups.name"),
        index=True,
        nullable=False,
        server_default="default",
        default="default",
    )
    schedulable: Mapped[bool] = mapped_column(
        "schedulable", sa.Boolean(), nullable=False, server_default=true(), default=True
    )
    available_slots: Mapped[ResourceSlot] = mapped_column(
        "available_slots", ResourceSlotColumn(), nullable=False
    )
    occupied_slots: Mapped[ResourceSlot] = mapped_column(
        "occupied_slots", ResourceSlotColumn(), nullable=False
    )
    addr: Mapped[str] = mapped_column("addr", sa.String(length=128), nullable=False)
    public_host: Mapped[str | None] = mapped_column(
        "public_host", sa.String(length=256), nullable=True
    )
    public_key: Mapped[PublicKey | None] = mapped_column(
        "public_key", CurvePublicKeyColumn(), nullable=True
    )
    first_contact: Mapped[datetime | None] = mapped_column(
        "first_contact", sa.DateTime(timezone=True), server_default=sa.func.now()
    )
    lost_at: Mapped[datetime | None] = mapped_column(
        "lost_at", sa.DateTime(timezone=True), nullable=True
    )
    version: Mapped[str] = mapped_column("version", sa.String(length=64), nullable=False)
    architecture: Mapped[str] = mapped_column("architecture", sa.String(length=32), nullable=False)
    compute_plugins: Mapped[dict[str, Any]] = mapped_column(
        "compute_plugins", pgsql.JSONB(), nullable=False, default={}
    )
    auto_terminate_abusing_kernel: Mapped[bool] = mapped_column(
        "auto_terminate_abusing_kernel",
        sa.Boolean(),
        nullable=False,
        server_default=false(),
        default=False,
    )

    kernels: Mapped[list[KernelRow]] = relationship("KernelRow", back_populates="agent_row")
    agent_resource_rows: Mapped[list[AgentResourceRow]] = relationship("AgentResourceRow")
    scaling_group_row: Mapped[ScalingGroupRow] = relationship(
        "ScalingGroupRow", back_populates="agents"
    )

    def actual_occupied_slots(self) -> ResourceSlot:
        occupied = ResourceSlot()
        for resource_row in self.agent_resource_rows:
            occupied[resource_row.slot_name] = resource_row.used
        return occupied

    def to_data(self) -> AgentData:
        return AgentData(
            id=AgentId(self.id),
            status=self.status,
            status_changed=self.status_changed,
            region=self.region,
            scaling_group=self.scaling_group,
            schedulable=self.schedulable,
            available_slots=self.available_slots,
            cached_occupied_slots=self.occupied_slots,
            actual_occupied_slots=self.actual_occupied_slots(),
            addr=self.addr,
            public_host=self.public_host,
            first_contact=self.first_contact,
            lost_at=self.lost_at,
            version=self.version,
            architecture=self.architecture,
            compute_plugins=self.compute_plugins,
            public_key=self.public_key,
            auto_terminate_abusing_kernel=self.auto_terminate_abusing_kernel,
        )

    def to_heartbeat_update_data(self) -> AgentDataForHeartbeatUpdate:
        return AgentDataForHeartbeatUpdate(
            status=self.status,
            status_changed=self.status_changed,
            scaling_group=self.scaling_group,
            available_slots=self.available_slots,
            addr=self.addr,
            public_host=self.public_host,
            version=self.version,
            architecture=self.architecture,
            compute_plugins=self.compute_plugins,
            public_key=self.public_key,
            auto_terminate_abusing_kernel=self.auto_terminate_abusing_kernel,
        )

    @classmethod
    async def get_agents_by_condition(
        cls, conditions: Sequence[QueryCondition], *, db: ExtendedAsyncSAEngine
    ) -> Sequence[AgentRow]:
        query_stmt = sa.select(AgentRow)
        for cond in conditions:
            query_stmt = cond(query_stmt)

        async def fetch(db_session: SASession) -> Sequence[AgentRow]:
            return (await db_session.scalars(query_stmt)).all()

        async with db.connect() as db_conn:
            return await execute_with_txn_retry(fetch, db.begin_readonly_session, db_conn)

    @classmethod
    async def get_schedulable_agents_by_sgroup(
        cls, sgroup_name: str, *, db: ExtendedAsyncSAEngine
    ) -> Sequence[AgentRow]:
        return await cls.get_agents_by_condition(
            [
                by_scaling_group(sgroup_name),
                by_schedulable(True),
                by_status(AgentStatus.ALIVE),
            ],
            db=db,
        )

    @classmethod
    async def get_occupied_slots(
        cls,
        db: ExtendedAsyncSAEngine,
        agent_id: AgentId,
        known_slot_types: Mapping[SlotName, SlotTypes],
    ) -> ResourceSlot:
        async with db.begin_readonly_session() as db_session:
            query = sa.select(AgentResourceRow.slot_name, AgentResourceRow.used).where(
                AgentResourceRow.agent_id == agent_id,
            )
            result = await db_session.execute(query)
            occupied_slots = ResourceSlot.from_known_slots(known_slot_types)
            for row in result:
                occupied_slots[row.slot_name] = row.used
            return occupied_slots


# For compatibility
agents = AgentRow.__table__


def by_scaling_group(
    scaling_group: str,
) -> QueryCondition:
    def _by_scaling_group(
        query_stmt: sa.sql.Select[Any],
    ) -> sa.sql.Select[Any]:
        return query_stmt.where(AgentRow.scaling_group == scaling_group)

    return _by_scaling_group


def by_status(
    status: AgentStatus,
) -> QueryCondition:
    def _by_status(
        query_stmt: sa.sql.Select[Any],
    ) -> sa.sql.Select[Any]:
        return query_stmt.where(AgentRow.status == status)

    return _by_status


def by_schedulable(
    schedulable: bool,
) -> QueryCondition:
    def _by_schedulable(
        query_stmt: sa.sql.Select[Any],
    ) -> sa.sql.Select[Any]:
        schedulable_ = true() if schedulable else false()
        return query_stmt.where(AgentRow.schedulable == schedulable_)

    return _by_schedulable


async def list_schedulable_agents_by_sgroup(
    db_sess: SASession,
    sgroup_name: str,
) -> Sequence[AgentRow]:
    query = sa.select(AgentRow).where(
        (AgentRow.status == AgentStatus.ALIVE)
        & (AgentRow.scaling_group == sgroup_name)
        & (AgentRow.schedulable == true()),
    )

    result = await db_sess.execute(query)
    return result.scalars().all()


type WhereClauseType = sa.sql.expression.BinaryExpression[Any] | sa.sql.expression.BooleanClauseList
# TypeAlias is deprecated since 3.12 but mypy does not follow up yet

OWNER_PERMISSIONS: frozenset[AgentPermission] = frozenset([perm for perm in AgentPermission])
ADMIN_PERMISSIONS: frozenset[AgentPermission] = frozenset([perm for perm in AgentPermission])
MONITOR_PERMISSIONS: frozenset[AgentPermission] = frozenset([
    AgentPermission.READ_ATTRIBUTE,
    AgentPermission.UPDATE_ATTRIBUTE,
])
PRIVILEGED_MEMBER_PERMISSIONS: frozenset[AgentPermission] = frozenset([
    AgentPermission.CREATE_COMPUTE_SESSION,
    AgentPermission.CREATE_SERVICE,
])
MEMBER_PERMISSIONS: frozenset[AgentPermission] = frozenset([
    AgentPermission.CREATE_COMPUTE_SESSION,
    AgentPermission.CREATE_SERVICE,
])


@dataclass
class AgentPermissionContext(AbstractPermissionContext[AgentPermission, AgentRow, AgentId]):
    from ai.backend.manager.models.scaling_group import ScalingGroupPermissionContext

    sgroup_permission_ctx: ScalingGroupPermissionContext | None = None

    @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, AgentRow.id.in_(self.object_id_to_additional_permission_map.keys())
            )
        if self.object_id_to_overriding_permission_map:
            cond = _OR_coalesce(
                cond, AgentRow.id.in_(self.object_id_to_overriding_permission_map.keys())
            )

        if self.sgroup_permission_ctx is not None:
            if cond is not None:
                sgroup_names = self.sgroup_permission_ctx.sgroup_to_permissions_map.keys()
                cond = cond & AgentRow.scaling_group.in_(sgroup_names)
        return cond

    def apply_sgroup_permission_ctx(
        self, sgroup_permission_ctx: ScalingGroupPermissionContext
    ) -> None:
        self.sgroup_permission_ctx = sgroup_permission_ctx

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

    async def calculate_final_permission(self, rbac_obj: AgentRow) -> frozenset[AgentPermission]:
        agent_row = rbac_obj
        agent_id = cast(AgentId, agent_row.id)
        permissions: set[AgentPermission] = set()

        if (
            overriding_perm := self.object_id_to_overriding_permission_map.get(agent_id)
        ) is not None:
            permissions = set(overriding_perm)
        else:
            permissions |= self.object_id_to_additional_permission_map.get(agent_id, set())

        if self.sgroup_permission_ctx is not None:
            sgroup_permission_map = self.sgroup_permission_ctx.sgroup_to_permissions_map
            sgroup_perms = sgroup_permission_map.get(agent_row.scaling_group)
            if sgroup_perms is None or ScalingGroupPermission.AGENT_PERMISSIONS not in sgroup_perms:
                permissions = set()

        return frozenset(permissions)


class AgentPermissionContextBuilder(
    AbstractPermissionContextBuilder[AgentPermission, AgentPermissionContext]
):
    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[AgentPermission]:
        roles = await get_predefined_roles_in_scope(ctx, target_scope, self.db_session)
        return await self._calculate_permission_by_predefined_roles(roles)

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

        perm_ctx = AgentPermissionContext()
        _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,
    ) -> AgentPermissionContext:
        from ai.backend.manager.models.scaling_group import (
            ScalingGroupForDomainRow,
            ScalingGroupRow,
        )

        permissions = await self.calculate_permission(ctx, scope)
        aid_permission_map: dict[AgentId, frozenset[AgentPermission]] = {}

        _stmt = (
            sa.select(ScalingGroupForDomainRow)
            .where(ScalingGroupForDomainRow.domain == scope.domain_name)
            .options(
                joinedload(ScalingGroupForDomainRow.sgroup_row).options(
                    selectinload(ScalingGroupRow.agents)
                )
            )
        )
        for row in await self.db_session.scalars(_stmt):
            sg_row = row.sgroup_row
            for ag in sg_row.agents:
                aid_permission_map[AgentId(ag.id)] = permissions
        return AgentPermissionContext(object_id_to_additional_permission_map=aid_permission_map)

    @override
    async def build_ctx_in_project_scope(
        self,
        ctx: ClientContext,
        scope: ProjectScope,
    ) -> AgentPermissionContext:
        from ai.backend.manager.models.scaling_group import (
            ScalingGroupForProjectRow,
            ScalingGroupRow,
        )

        permissions = await self.calculate_permission(ctx, scope)
        aid_permission_map: dict[AgentId, frozenset[AgentPermission]] = {}

        _stmt = (
            sa.select(ScalingGroupForProjectRow)
            .where(ScalingGroupForProjectRow.group == scope.project_id)
            .options(
                joinedload(ScalingGroupForProjectRow.sgroup_row).options(
                    selectinload(ScalingGroupRow.agents)
                )
            )
        )
        for row in await self.db_session.scalars(_stmt):
            sg_row = row.sgroup_row
            for ag in sg_row.agents:
                aid_permission_map[AgentId(ag.id)] = permissions
        return AgentPermissionContext(object_id_to_additional_permission_map=aid_permission_map)

    @override
    async def build_ctx_in_user_scope(
        self,
        ctx: ClientContext,
        scope: UserScope,
    ) -> AgentPermissionContext:
        from ai.backend.manager.models.scaling_group import (
            ScalingGroupForKeypairsRow,
            ScalingGroupRow,
        )

        permissions = await self.calculate_permission(ctx, scope)
        aid_permission_map: dict[AgentId, frozenset[AgentPermission]] = {}

        _kp_stmt = (
            sa.select(KeyPairRow)
            .where(KeyPairRow.user == scope.user_id)
            .options(load_only(KeyPairRow.access_key))
        )
        kp_rows = (await self.db_session.scalars(_kp_stmt)).all()
        access_keys = cast(list[AccessKey], [r.access_key for r in kp_rows])

        _stmt = (
            sa.select(ScalingGroupForKeypairsRow)
            .where(ScalingGroupForKeypairsRow.access_key.in_(access_keys))
            .options(
                joinedload(ScalingGroupForKeypairsRow.sgroup_row).options(
                    selectinload(ScalingGroupRow.agents)
                )
            )
        )
        for row in await self.db_session.scalars(_stmt):
            sg_row = row.sgroup_row
            for ag in sg_row.agents:
                aid_permission_map[AgentId(ag.id)] = permissions
        return AgentPermissionContext(object_id_to_additional_permission_map=aid_permission_map)

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

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

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

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

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


async def get_permission_ctx(
    db_conn: SAConnection,
    ctx: ClientContext,
    target_scope: ScopeType,
    requested_permission: AgentPermission,
) -> AgentPermissionContext:
    from ai.backend.manager.models.scaling_group import ScalingGroupPermissionContextBuilder

    async with ctx.db.begin_readonly_session(db_conn) as db_session:
        sgroup_perm_ctx = await ScalingGroupPermissionContextBuilder(db_session).build(
            ctx, target_scope, ScalingGroupPermission.AGENT_PERMISSIONS
        )

        builder = AgentPermissionContextBuilder(db_session)
        permission_ctx = await builder.build(ctx, target_scope, requested_permission)
        permission_ctx.apply_sgroup_permission_ctx(sgroup_perm_ctx)
    return permission_ctx
