from typing import List

import grpc
from google.protobuf import timestamp_pb2

from python_pachyderm.errors import AuthServiceNotActivated
from python_pachyderm.proto.v2.enterprise import enterprise_pb2
from python_pachyderm.proto.v2.license import license_pb2, license_pb2_grpc


class LicenseMixin:
    """A mixin for license-related functionality."""

    _channel: grpc.Channel

    def __init__(self):
        self.__stub = license_pb2_grpc.APIStub(self._channel)
        super().__init__()

    def activate_license(
        self, activation_code: str, expires: timestamp_pb2.Timestamp = None
    ) -> enterprise_pb2.TokenInfo:
        """Activates the license service.

        Parameters
        ----------
        activation_code : str
            A Pachyderm enterprise activation code. New users can obtain trial
            activation codes.
        expires : timestamp_pb2.Timestamp, optional
            A protobuf object indicating when this activation code will expire.
            This should generally not be set and is only applied if it is
            earlier than the signed expiration time of `activation_code`.

        Returns
        -------
        enterprise_pb2.TokenInfo
            A protobuf object that has the expiration for the current token.
            Field will be unset if there is no current token.
        """
        message = license_pb2.ActivateRequest(
            activation_code=activation_code, expires=expires
        )
        return self.__stub.Activate(message).info

    def add_cluster(
        self,
        id: str,
        address: str,
        secret: str = None,
        user_address: str = None,
        cluster_deployment_id: str = None,
        enterprise_server: bool = False,
    ) -> license_pb2.AddClusterResponse:
        """Register a cluster with the license service.

        Parameters
        ----------
        id : str
            The unique ID to identify the cluster.
        address : str
            The public GRPC address for the license server to reach the cluster.
        secret : str, optional
            A shared secret for the cluster to use to authenticate. If not
            specified, a random secret will be generated and returned.
        user_address : str, optional
            The pachd address for users to connect to.
        cluster_deployment_id : str, optional
            The deployment ID value generated by each cluster.
        enterprise_server : bool, optional
            Indicates whether the address points to an enterprise server.

        Returns
        -------
        license_pb2.AddClusterResponse
            A protobuf object that returns the `secret`.
        """
        message = license_pb2.AddClusterRequest(
            address=address,
            cluster_deployment_id=cluster_deployment_id,
            enterprise_server=enterprise_server,
            id=id,
            secret=secret,
            user_address=user_address,
        )
        return self.__stub.AddCluster(message)

    def update_cluster(
        self,
        id: str,
        address: str,
        user_address: str = None,
        cluster_deployment_id: str = None,
    ) -> None:
        """Update a cluster registered with the license service.

        Parameters
        ----------
        id : str
            The unique ID to identify the cluster.
        address : str
            The public GRPC address for the license server to reach the cluster.
        user_address : str, optional
            The pachd address for users to connect to.
        cluster_deployment_id : str, optional
            The deployment ID value generated by each cluster.
        """
        message = license_pb2.UpdateClusterRequest(
            address=address,
            cluster_deployment_id=cluster_deployment_id,
            id=id,
            user_address=user_address,
        )
        self.__stub.UpdateCluster(message)

    def delete_cluster(self, id: str) -> None:
        """Delete a cluster registered with the license service.

        Parameters
        ----------
        id : str
            The unique ID to identify the cluster.
        """
        message = license_pb2.DeleteClusterRequest(id=id)
        self.__stub.DeleteCluster(message)

    def list_clusters(self) -> List[license_pb2.ClusterStatus]:
        """List clusters registered with the license service.

        Returns
        -------
        List[license_pb2.ClusterStatus]
            A list of protobuf objects that return info on a cluster.
        """
        message = license_pb2.ListClustersRequest()
        return self.__stub.ListClusters(message).clusters

    def get_activation_code(self) -> license_pb2.GetActivationCodeResponse:
        """Gets the enterprise code used to activate the server.

        Returns
        -------
        license_pb2.GetActivationCodeResponse
            A protobuf object that returns a state enum, info on the token,
            and the activation code.
        """
        message = license_pb2.GetActivationCodeRequest()
        return self.__stub.GetActivationCode(message)

    def delete_all_license(self) -> None:
        """Remove all clusters and deactivate the license service.

        Raises
        ------
        AuthServiceNotActivated
        """
        message = license_pb2.DeleteAllRequest()
        try:
            self.__stub.DeleteAll(message)
        except grpc.RpcError as err:
            raise AuthServiceNotActivated.try_from(err)

    def list_user_clusters(self) -> List[license_pb2.UserClusterInfo]:
        """Lists all clusters available to user.

        Returns
        -------
        List[license_pb2.UserClusterInfo]
            A list of protobuf objects that return info on clusters
            available to the users.
        """
        message = license_pb2.ListUserClustersRequest()
        return self.__stub.ListUserClusters(message).clusters

from typing import List

import grpc
from google.protobuf import timestamp_pb2

from python_pachyderm.errors import AuthServiceNotActivated
from python_pachyderm.proto.v2.enterprise import enterprise_pb2
from python_pachyderm.proto.v2.license import license_pb2, license_pb2_grpc


class LicenseMixin:
    """A mixin for license-related functionality."""

    _channel: grpc.Channel

    def __init__(self):
        self.__stub = license_pb2_grpc.APIStub(self._channel)
        super().__init__()

    def activate_license(
        self, activation_code: str, expires: timestamp_pb2.Timestamp = None
    ) -> enterprise_pb2.TokenInfo:
        """Activates the license service.

        Parameters
        ----------
        activation_code : str
            A Pachyderm enterprise activation code. New users can obtain trial
            activation codes.
        expires : timestamp_pb2.Timestamp, optional
            A protobuf object indicating when this activation code will expire.
            This should generally not be set and is only applied if it is
            earlier than the signed expiration time of `activation_code`.

        Returns
        -------
        enterprise_pb2.TokenInfo
            A protobuf object that has the expiration for the current token.
            Field will be unset if there is no current token.
        """
        message = license_pb2.ActivateRequest(
            activation_code=activation_code, expires=expires
        )
        return self.__stub.Activate(message).info

    def add_cluster(
        self,
        id: str,
        address: str,
        secret: str = None,
        user_address: str = None,
        cluster_deployment_id: str = None,
        enterprise_server: bool = False,
    ) -> license_pb2.AddClusterResponse:
        """Register a cluster with the license service.

        Parameters
        ----------
        id : str
            The unique ID to identify the cluster.
        address : str
            The public GRPC address for the license server to reach the cluster.
        secret : str, optional
            A shared secret for the cluster to use to authenticate. If not
            specified, a random secret will be generated and returned.
        user_address : str, optional
            The pachd address for users to connect to.
        cluster_deployment_id : str, optional
            The deployment ID value generated by each cluster.
        enterprise_server : bool, optional
            Indicates whether the address points to an enterprise server.

        Returns
        -------
        license_pb2.AddClusterResponse
            A protobuf object that returns the `secret`.
        """
        message = license_pb2.AddClusterRequest(
            address=address,
            cluster_deployment_id=cluster_deployment_id,
            enterprise_server=enterprise_server,
            id=id,
            secret=secret,
            user_address=user_address,
        )
        return self.__stub.AddCluster(message)

    def update_cluster(
        self,
        id: str,
        address: str,
        user_address: str = None,
        cluster_deployment_id: str = None,
    ) -> None:
        """Update a cluster registered with the license service.

        Parameters
        ----------
        id : str
            The unique ID to identify the cluster.
        address : str
            The public GRPC address for the license server to reach the cluster.
        user_address : str, optional
            The pachd address for users to connect to.
        cluster_deployment_id : str, optional
            The deployment ID value generated by each cluster.
        """
        message = license_pb2.UpdateClusterRequest(
            address=address,
            cluster_deployment_id=cluster_deployment_id,
            id=id,
            user_address=user_address,
        )
        self.__stub.UpdateCluster(message)

    def delete_cluster(self, id: str) -> None:
        """Delete a cluster registered with the license service.

        Parameters
        ----------
        id : str
            The unique ID to identify the cluster.
        """
        message = license_pb2.DeleteClusterRequest(id=id)
        self.__stub.DeleteCluster(message)

    def list_clusters(self) -> List[license_pb2.ClusterStatus]:
        """List clusters registered with the license service.

        Returns
        -------
        List[license_pb2.ClusterStatus]
            A list of protobuf objects that return info on a cluster.
        """
        message = license_pb2.ListClustersRequest()
        return self.__stub.ListClusters(message).clusters

    def get_activation_code(self) -> license_pb2.GetActivationCodeResponse:
        """Gets the enterprise code used to activate the server.

        Returns
        -------
        license_pb2.GetActivationCodeResponse
            A protobuf object that returns a state enum, info on the token,
            and the activation code.
        """
        message = license_pb2.GetActivationCodeRequest()
        return self.__stub.GetActivationCode(message)

    def delete_all_license(self) -> None:
        """Remove all clusters and deactivate the license service.

        Raises
        ------
        AuthServiceNotActivated
        """
        message = license_pb2.DeleteAllRequest()
        try:
            self.__stub.DeleteAll(message)
        except grpc.RpcError as err:
            raise AuthServiceNotActivated.try_from(err)

    def list_user_clusters(self) -> List[license_pb2.UserClusterInfo]:
        """Lists all clusters available to user.

        Returns
        -------
        List[license_pb2.UserClusterInfo]
            A list of protobuf objects that return info on clusters
            available to the users.
        """
        message = license_pb2.ListUserClustersRequest()
        return self.__stub.ListUserClusters(message).clusters

A mixin for license-related functionality. Activates the license service.          Parameters         ----------         activation_code : str             A Pachyderm enterprise activation code. New users can obtain trial             activation codes.         expires : timestamp_pb2.Timestamp, optional             A protobuf object indicating when this activation code will expire.             This should generally not be set and is only applied if it is             earlier than the signed expiration time of `activation_code`.          Returns         -------         enterprise_pb2.TokenInfo             A protobuf object that has the expiration for the current token.             Field will be unset if there is no current token. Register a cluster with the license service.          Parameters         ----------         id : str             The unique ID to identify the cluster.         address : str             The public GRPC address for the license server to reach the cluster.         secret : str, optional             A shared secret for the cluster to use to authenticate. If not             specified, a random secret will be generated and returned.         user_address : str, optional             The pachd address for users to connect to.         cluster_deployment_id : str, optional             The deployment ID value generated by each cluster.         enterprise_server : bool, optional             Indicates whether the address points to an enterprise server.          Returns         -------         license_pb2.AddClusterResponse             A protobuf object that returns the `secret`. Update a cluster registered with the license service.          Parameters         ----------         id : str             The unique ID to identify the cluster.         address : str             The public GRPC address for the license server to reach the cluster.         user_address : str, optional             The pachd address for users to connect to.         cluster_deployment_id : str, optional             The deployment ID value generated by each cluster. Delete a cluster registered with the license service.          Parameters         ----------         id : str             The unique ID to identify the cluster. List clusters registered with the license service.          Returns         -------         List[license_pb2.ClusterStatus]             A list of protobuf objects that return info on a cluster. Gets the enterprise code used to activate the server.          Returns         -------         license_pb2.GetActivationCodeResponse             A protobuf object that returns a state enum, info on the token,             and the activation code. Remove all clusters and deactivate the license service.          Raises         ------         AuthServiceNotActivated Lists all clusters available to user.          Returns         -------         List[license_pb2.UserClusterInfo]             A list of protobuf objects that return info on clusters             available to the users.