import typing as t
from typing import overload
from abc import ABCMeta, abstractmethod
from .types import KeyDict, KeyAny, KeyParameters
from ..registry import (
    KeyParameterRegistryDict,
    JWK_PARAMETER_REGISTRY,
    KeyOperationRegistryDict,
    JWK_OPERATION_REGISTRY,
)
from ..util import to_bytes
from ..rfc7638 import thumbprint
from ..errors import (
    UnsupportedKeyUseError,
    UnsupportedKeyAlgorithmError,
    UnsupportedKeyOperationError,
)


NativePrivateKey = t.TypeVar("NativePrivateKey")
NativePublicKey = t.TypeVar("NativePublicKey")


class NativeKeyBinding(metaclass=ABCMeta):
    use_key_ops_registry: t.ClassVar[t.Dict[str, t.List[str]]] = {
        "sig": ["sign", "verify"],
        "enc": ["encrypt", "decrypt", "wrapKey", "unwrapKey", "deriveKey", "deriveBits"],
    }

    @classmethod
    @abstractmethod
    def convert_raw_key_to_dict(cls, raw_key: t.Any, private: bool) -> KeyDict:
        pass

    @classmethod
    @abstractmethod
    def import_from_dict(cls, value: KeyDict):
        pass

    @classmethod
    @abstractmethod
    def import_from_bytes(cls, value: bytes, password: t.Optional[t.Any] = None):
        pass

    @staticmethod
    def as_bytes(key: "BaseKey", encoding=None, private=None, password=None) -> bytes:  # pragma: no cover
        return key.raw_value

    @classmethod
    def validate_dict_key_registry(cls, dict_key: KeyDict, registry: KeyParameterRegistryDict):
        for k in registry:
            if registry[k].required and k not in dict_key:
                raise ValueError(f'"{k}" is required')

            if k in dict_key:
                try:
                    registry[k].validate(dict_key[k])
                except ValueError as error:
                    raise ValueError(f'"{k}" {error}')

    @classmethod
    def validate_dict_key_use_operations(cls, dict_key: KeyDict):
        if "use" in dict_key and "key_ops" in dict_key:
            _use: str = dict_key["use"]  # type: ignore
            operations = cls.use_key_ops_registry[_use]
            for op in dict_key["key_ops"]:
                if op not in operations:
                    raise ValueError('"use" and "key_ops" does not match')


class BaseKey(t.Generic[NativePrivateKey, NativePublicKey]):
    key_type: t.ClassVar[str]
    binding: t.ClassVar[t.Type[NativeKeyBinding]]
    value_registry: t.ClassVar[KeyParameterRegistryDict]
    param_registry: t.ClassVar[KeyParameterRegistryDict] = JWK_PARAMETER_REGISTRY
    operation_registry: t.ClassVar[KeyOperationRegistryDict] = JWK_OPERATION_REGISTRY

    def __init__(
            self,
            raw_value: t.Union[NativePrivateKey, NativePublicKey],
            original_value: t.Any,
            parameters: t.Optional[KeyParameters] = None):
        self._raw_value = raw_value
        self.original_value = original_value
        self.extra_parameters = parameters
        self._dict_value: KeyDict = {}
        if isinstance(original_value, dict):
            if parameters is not None:
                data = {**original_value, **parameters, "kty": self.key_type}
            else:
                data = {**original_value, "kty": self.key_type}
            self.validate_dict_key(data)
            self._dict_value = data

    def keys(self):
        return self.dict_value.keys()

    def __getitem__(self, k: str):
        return self.dict_value[k]

    def get(self, k: str, default=None):
        return self.dict_value.get(k, default)

    @property
    def kid(self) -> str:
        kid = self.get("kid")
        if not kid:
            kid = self.thumbprint()
            self._dict_value["kid"] = kid
        return kid

    @property
    def raw_value(self):
        raise NotImplementedError()

    @property
    def is_private(self) -> bool:
        raise NotImplementedError()

    @property
    def dict_value(self) -> KeyDict:
        """Property of the Key in Dict (JSON)."""
        if self._dict_value:
            return self._dict_value

        data = self.binding.convert_raw_key_to_dict(self.raw_value, self.is_private)
        if self.extra_parameters is not None:
            data.update(self.extra_parameters)  # type: ignore
        data["kty"] = self.key_type
        self.validate_dict_key(data)
        self._dict_value = data
        return data

    @property
    def public_key(self) -> NativePublicKey:
        raise NotImplementedError()

    @property
    def private_key(self) -> t.Optional[NativePrivateKey]:
        raise NotImplementedError()

    def thumbprint(self) -> str:
        """Call this method will generate the thumbprint with algorithm
        defined in RFC7638."""
        fields = [k for k in self.value_registry if self.value_registry[k].required]
        fields.append("kty")
        return thumbprint(self.dict_value, fields)

    def as_dict(self, private: t.Optional[bool] = None, **params: t.Any) -> KeyDict:
        """Output this key to a JWK format (in dict). By default, it will return
        the ``dict_value`` of this key.

        :param private: determine whether this method should output private key or not
        :param params: other parameters added into this key
        :raise: ValueError
        """
        # check private conflicts
        if private and not self.is_private:
            raise ValueError("This key is not a private key.")

        data = self.dict_value.copy()
        if private is not False:
            data.update(params)
            return data

        # clear private fields
        for k in self.dict_value:
            if k in self.value_registry and self.value_registry[k].private:
                del data[k]

        data.update(params)
        return data

    def check_use(self, use: str):
        """Check if this key supports the given "use".

        :param use: this key is used for, e.g. "sig", "enc"
        :raise: UnsupportedKeyUseError
        """
        designed_use = self.get("use")
        if designed_use and designed_use != use:
            raise UnsupportedKeyUseError(f'This key is designed to be used for "{designed_use}"')

    def check_alg(self, alg: str):
        """Check if this key supports the given "alg".

        :param alg: the algorithm this key is intended to be used, e.g. "HS256", "ECDH-EC"
        :raise: UnsupportedKeyAlgorithmError
        """
        designed_alg = self.get("alg")
        if designed_alg and designed_alg != alg:
            raise UnsupportedKeyAlgorithmError(f'This key is designed for algorithm "{designed_alg}"')

    def check_key_op(self, operation: str):
        """Check if the given key_op is supported by this key.

        :param operation: key operation value, such as "sign", "encrypt".
        :raise: UnsupportedKeyOperationError
        """
        key_ops = self.get("key_ops")
        if key_ops is not None and operation not in key_ops:
            raise UnsupportedKeyOperationError(f'Unsupported key_op "{operation}"')

        assert operation in self.operation_registry
        reg = self.operation_registry[operation]
        if reg.private and not self.is_private:
            raise UnsupportedKeyOperationError(f'Invalid key_op "{operation}" for public key')

    @overload
    def get_op_key(self, operation: t.Literal["verify", "encrypt", "wrapKey", "deriveKey"]) -> NativePublicKey: ...

    @overload
    def get_op_key(self, operation: t.Literal["sign", "decrypt", "unwrapKey"]) -> NativePrivateKey: ...

    def get_op_key(self, operation: str) -> t.Union[NativePublicKey, NativePrivateKey]:
        self.check_key_op(operation)
        reg = self.operation_registry[operation]
        if reg.private:
            assert self.private_key is not None
            return self.private_key
        return self.public_key

    @classmethod
    def validate_dict_key(cls, data: KeyDict):
        cls.binding.validate_dict_key_registry(data, cls.param_registry)
        cls.binding.validate_dict_key_registry(data, cls.value_registry)
        cls.binding.validate_dict_key_use_operations(data)

    @classmethod
    def import_key(
            cls,
            value: KeyAny,
            parameters: t.Optional[KeyParameters] = None,
            password: t.Optional[t.Any] = None) -> "BaseKey":
        if isinstance(value, dict):
            cls.validate_dict_key(value)
            raw_key = cls.binding.import_from_dict(value)
            return cls(raw_key, value, parameters)

        raw_key = cls.binding.import_from_bytes(to_bytes(value), password)
        return cls(raw_key, value, parameters)

    @classmethod
    def generate_key(
            cls,
            size_or_crv,
            parameters: t.Optional[KeyParameters] = None,
            private: bool = True) -> "BaseKey":
        raise NotImplementedError()


class SymmetricKey(BaseKey[bytes, bytes], metaclass=ABCMeta):
    @property
    def raw_value(self) -> bytes:
        """The raw key in bytes."""
        return self._raw_value

    @property
    def is_private(self) -> bool:
        """A symmetric key will always be private."""
        return True

    @property
    def public_key(self) -> bytes:
        """Returns the ``raw_value`` as the public key."""
        return self.raw_value

    @property
    def private_key(self) -> bytes:
        """Returns the ``raw_value`` as the private key."""
        return self.raw_value


class AsymmetricKey(BaseKey[NativePrivateKey, NativePublicKey], metaclass=ABCMeta):
    @property
    def raw_value(self) -> t.Union[NativePublicKey, NativePrivateKey]:
        return self._raw_value

    def as_bytes(
            self,
            encoding: t.Optional[str] = None,
            private: t.Optional[bool] = None,
            password: t.Optional[str] = None) -> bytes:
        return self.binding.as_bytes(self, encoding, private, password)

    def as_pem(self, private=None, password=None) -> bytes:
        return self.as_bytes(private=private, password=password)

    def as_der(self, private=None, password=None) -> bytes:
        return self.as_bytes(encoding="DER", private=private, password=password)


class CurveKey(AsymmetricKey[NativePrivateKey, NativePublicKey]):
    @property
    @abstractmethod
    def curve_name(self) -> str:
        pass

    @abstractmethod
    def exchange_derive_key(self, key) -> bytes:
        pass
