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,
        secret: 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.
        secret : str, optional
            A shared secret for the cluster to use to authenticate. If not
            specified, a random secret will be generated and returned.
        """
        message = license_pb2.UpdateClusterRequest(
            address=address,
            cluster_deployment_id=cluster_deployment_id,
            id=id,
            user_address=user_address,
            secret=secret,
        )
        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
