from __future__ import annotations

import logging
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any
from uuid import UUID

import graphene
import sqlalchemy as sa
from graphql import Undefined
from sqlalchemy.engine.row import Row

from ai.backend.common.data.user.types import UserRole
from ai.backend.common.types import BinarySize, ResourceSlot
from ai.backend.logging import BraceStyleAdapter
from ai.backend.manager.data.resource_preset.types import ResourcePresetData
from ai.backend.manager.models.minilang.ordering import ColumnMapType, QueryOrderParser
from ai.backend.manager.models.minilang.queryfilter import FieldSpecType, QueryFilterParser
from ai.backend.manager.models.resource_preset import ResourcePresetRow, resource_presets
from ai.backend.manager.repositories.base.creator import Creator
from ai.backend.manager.repositories.base.updater import Updater
from ai.backend.manager.repositories.resource_preset.creators import ResourcePresetCreatorSpec
from ai.backend.manager.repositories.resource_preset.updaters import ResourcePresetUpdaterSpec
from ai.backend.manager.types import OptionalState, TriState

from .base import (
    BigInt,
    batch_result,
    batch_result_in_scalar_stream,
)

if TYPE_CHECKING:
    from .schema import GraphQueryContext

log = BraceStyleAdapter(logging.getLogger("ai.backend.manager.models"))

__all__: Sequence[str] = (
    "CreateResourcePreset",
    "DeleteResourcePreset",
    "ModifyResourcePreset",
    "ResourcePreset",
)


class ResourcePreset(graphene.ObjectType):  # type: ignore[misc]
    id = graphene.UUID(description="Added in 25.4.0. ID of the resource preset.")
    name = graphene.String()
    resource_slots = graphene.JSONString()
    shared_memory = BigInt()
    scaling_group_name = graphene.String(
        description=(
            "Added in 25.4.0. A name of scaling group(=resource group) of the resource preset associated with."
        ),
    )

    _queryfilter_fieldspec: FieldSpecType = {
        "id": ("id", None),
        "name": ("name", None),
        "scaling_group_name": ("scaling_group_name", None),
    }

    _queryorder_colmap: ColumnMapType = {
        "id": ("id", None),
        "name": ("name", None),
        "scaling_group_name": ("scaling_group_name", None),
    }

    @classmethod
    def from_row(
        cls,
        ctx: GraphQueryContext,
        row: ResourcePresetRow | ResourcePresetData | Row[Any] | None,
    ) -> ResourcePreset | None:
        match row:
            case ResourcePresetRow():
                shared_memory = str(row.shared_memory) if row.shared_memory is not None else None
                return cls(
                    id=row.id,
                    name=row.name,
                    resource_slots=row.resource_slots.to_json(),
                    shared_memory=shared_memory,
                    scaling_group_name=row.scaling_group_name,
                )
            case ResourcePresetData():
                shared_memory = str(row.shared_memory) if row.shared_memory is not None else None
                return cls(
                    id=row.id,
                    name=row.name,
                    resource_slots=row.resource_slots.to_json(),
                    shared_memory=shared_memory,
                    scaling_group_name=row.scaling_group_name,
                )
            case Row():
                shared_memory = str(row.shared_memory) if row.shared_memory else None
                return cls(
                    id=row.id,
                    name=row.name,
                    resource_slots=row.resource_slots.to_json(),
                    shared_memory=shared_memory,
                    scaling_group_name=row.scaling_group_name,
                )
            case _:
                return None

    @classmethod
    async def load_all(
        cls,
        ctx: GraphQueryContext,
        *,
        filter: str | None = None,
        order: str | None = None,
    ) -> Sequence[ResourcePreset]:
        query = sa.select(ResourcePresetRow)
        if filter is not None:
            filter_parser = QueryFilterParser(cls._queryfilter_fieldspec)
            query = filter_parser.append_filter(query, filter)
        if order is not None:
            order_parser = QueryOrderParser(cls._queryorder_colmap)
            query = order_parser.append_ordering(query, order)
        else:
            query = query.order_by(ResourcePresetRow.name)
        async with ctx.db.begin_readonly_session() as db_session:
            return [
                obj
                async for r in (await db_session.stream_scalars(query))
                if (obj := cls.from_row(ctx, r)) is not None
            ]

    @classmethod
    async def batch_load_by_id(
        cls,
        ctx: GraphQueryContext,
        ids: Sequence[UUID],
    ) -> Sequence[ResourcePreset | None]:
        query = (
            sa.select(ResourcePresetRow)
            .where(ResourcePresetRow.id.in_(ids))
            .order_by(ResourcePresetRow.id)
        )
        async with ctx.db.begin_readonly_session() as db_session:
            return await batch_result_in_scalar_stream(
                ctx,
                db_session,
                query,
                cls,
                ids,
                lambda row: row.id,
            )

    @classmethod
    async def batch_load_by_name(
        cls,
        ctx: GraphQueryContext,
        names: Sequence[str],
    ) -> Sequence[ResourcePreset | None]:
        query = (
            sa.select(resource_presets)
            .select_from(resource_presets)
            .where(resource_presets.c.name.in_(names))
            .order_by(resource_presets.c.name)
        )
        async with ctx.db.begin_readonly() as conn:
            return await batch_result(
                ctx,
                conn,
                query,
                cls,
                names,
                lambda row: row.name,
            )


class CreateResourcePresetInput(graphene.InputObjectType):  # type: ignore[misc]
    resource_slots = graphene.JSONString(required=True)
    shared_memory = graphene.String(required=False)
    scaling_group_name = graphene.String(
        required=False,
        description=(
            "Added in 25.4.0. A name of scaling group(=resource group) of the resource preset associated with."
        ),
    )

    def to_creator(self, name: str) -> Creator[ResourcePresetRow]:
        return Creator(
            spec=ResourcePresetCreatorSpec(
                name=name,
                resource_slots=ResourceSlot.from_user_input(self.resource_slots, None),
                shared_memory=self.shared_memory if self.shared_memory else None,
                scaling_group_name=self.scaling_group_name if self.scaling_group_name else None,
            )
        )


class ModifyResourcePresetInput(graphene.InputObjectType):  # type: ignore[misc]
    name = graphene.String(
        required=False,
        description=("Added in 25.4.0. A name of resource preset."),
    )
    resource_slots = graphene.JSONString(required=False)
    shared_memory = graphene.String(required=False)
    scaling_group_name = graphene.String(
        required=False,
        description=(
            "Added in 25.4.0. A name of scaling group(=resource group) of the resource preset associated with."
        ),
    )

    def to_updater(self, id: UUID | None, name: str | None) -> Updater[ResourcePresetRow]:
        resource_slots = (
            ResourceSlot.from_json(self.resource_slots) if self.resource_slots else Undefined
        )

        pk_value = id if id is not None else (name if name is not None else "")
        return Updater(
            spec=ResourcePresetUpdaterSpec(
                resource_slots=OptionalState[ResourceSlot].from_graphql(resource_slots),
                name=OptionalState[str].from_graphql(self.name),
                shared_memory=TriState[BinarySize].from_graphql(
                    BinarySize.finite_from_str(self.shared_memory)
                    if self.shared_memory is not Undefined and self.shared_memory is not None
                    else self.shared_memory
                ),
                scaling_group_name=TriState[str].from_graphql(self.scaling_group_name),
            ),
            pk_value=pk_value,
        )


class CreateResourcePreset(graphene.Mutation):  # type: ignore[misc]
    allowed_roles = (UserRole.SUPERADMIN,)

    class Arguments:
        name = graphene.String(required=True)
        props = CreateResourcePresetInput(required=True)

    ok = graphene.Boolean()
    msg = graphene.String()
    resource_preset = graphene.Field(lambda: ResourcePreset, required=False)

    @classmethod
    async def mutate(
        cls,
        root: Any,
        info: graphene.ResolveInfo,
        name: str,
        props: CreateResourcePresetInput,
    ) -> CreateResourcePreset:
        from ai.backend.manager.services.resource_preset.actions.create_preset import (
            CreateResourcePresetAction,
        )

        graph_ctx: GraphQueryContext = info.context

        result = await graph_ctx.processors.resource_preset.create_preset.wait_for_complete(
            CreateResourcePresetAction(creator=props.to_creator(name))
        )

        return cls(True, "success", ResourcePreset.from_row(graph_ctx, result.resource_preset))


class ModifyResourcePreset(graphene.Mutation):  # type: ignore[misc]
    allowed_roles = (UserRole.SUPERADMIN,)

    class Arguments:
        id = graphene.UUID(
            required=False,
            default_value=None,
            description=("Added in 25.4.0. ID of the resource preset."),
        )
        name = graphene.String(
            required=False, default_value=None, deprecation_reason="Deprecated since 25.4.0."
        )
        props = ModifyResourcePresetInput(required=True)

    ok = graphene.Boolean()
    msg = graphene.String()

    @classmethod
    async def mutate(
        cls,
        root: Any,
        info: graphene.ResolveInfo,
        id: UUID | None,
        name: str | None,
        props: ModifyResourcePresetInput,
    ) -> ModifyResourcePreset:
        from ai.backend.manager.services.resource_preset.actions.modify_preset import (
            ModifyResourcePresetAction,
        )

        graph_ctx: GraphQueryContext = info.context

        await graph_ctx.processors.resource_preset.modify_preset.wait_for_complete(
            ModifyResourcePresetAction(id=id, name=name, updater=props.to_updater(id, name))
        )

        return cls(True, "success")


class DeleteResourcePreset(graphene.Mutation):  # type: ignore[misc]
    allowed_roles = (UserRole.SUPERADMIN,)

    class Arguments:
        id = graphene.UUID(
            required=False,
            default_value=None,
            description=("Added in 25.4.0. ID of the resource preset."),
        )
        name = graphene.String(
            required=False, default_value=None, deprecation_reason="Deprecated since 25.4.0."
        )

    ok = graphene.Boolean()
    msg = graphene.String()

    @classmethod
    async def mutate(
        cls,
        root: Any,
        info: graphene.ResolveInfo,
        id: UUID | None,
        name: str | None,
    ) -> DeleteResourcePreset:
        from ai.backend.manager.services.resource_preset.actions.delete_preset import (
            DeleteResourcePresetAction,
        )

        graph_ctx: GraphQueryContext = info.context

        await graph_ctx.processors.resource_preset.delete_preset.wait_for_complete(
            DeleteResourcePresetAction(id=id, name=name)
        )

        return cls(True, "success")
