from .... import types
from ....compiler.registry import register_action

# API imports
from ...base import Action
from ...datatypes import common as common
from ...datatypes import compute as params

# Type imports
from ...datatypes.common import Response


@register_action
class AddDatasinks(Action):
    """Add datasinks to consumer list.

    Takes a list of datasinks and adds them to the current consumer list.

    Attributes:
        datasinks (List[params.DatasinkInfo]): name of target datasinks
    """

    datasinks: list[params.DatasinkInfo]

    async def execute(self) -> Response:
        """Execute the AddDatasinks action.

        Returns:
            Response: response generated by the RPC call.
        """
        # TODO: convert the common response from grpc to dataypes.common.Response
        return await types.COMPUTE.add_datasinks(**self.model_dump())


@register_action
class SetDatasinks(Action):
    """Set the datasink consumer list.

    Takes a list of datasinks and replaces the current consumer list with them.

    Attributes:
        datasinks (List[params.DatasinkInfo]): name of target datasinks
    """

    datasinks: list[params.DatasinkInfo]

    async def execute(self) -> Response:
        """Execute the SetDatasinks action.

        Returns:
            Response: response generated by the RPC call.
        """
        return await types.COMPUTE.set_datasinks(**self.model_dump())


@register_action
class RemoveDatasinks(Action):
    """Remove datasinks from consumer list.

    Takes a list of datasinks and removes them from the current consumer list.

    Attributes:
        datasinks (List[params.DatasinkInfo]): name of target datasinks
    """

    datasinks: list[params.DatasinkInfo]

    async def execute(self) -> Response:
        """Execute the RemoveDatasinks action.

        Returns:
            Response: response generated by the RPC call.
        """
        return await types.COMPUTE.remove_datasinks(**self.model_dump())
