from __future__ import annotations
import typing as t
from ._keys import (
    JWKRegistry,
    KeySet,
    Key,
    KeySetSerialization,
)
from ._rfc7517.types import AnyKey, DictKey, KeyParameters
from ._rfc7518.oct_key import OctKey as OctKey
from ._rfc7518.rsa_key import RSAKey as RSAKey
from ._rfc7518.ec_key import ECKey as ECKey
from ._rfc8037.okp_key import OKPKey as OKPKey
from ._rfc8812 import register_secp256k1
from ._rfc7638 import calculate_thumbprint as thumbprint
from ._rfc9278 import calculate_thumbprint_uri as thumbprint_uri
from .registry import Header


__all__ = [
    # types
    "Key",
    "DictKey",
    "KeyParameters",
    "KeyCallable",
    "KeyFlexible",
    "KeySetSerialization",
    "KeyBase",
    "GuestProtocol",
    # modules
    "JWKRegistry",
    "OctKey",
    "RSAKey",
    "ECKey",
    "OKPKey",
    "KeySet",
    # methods
    "guess_key",
    "import_key",
    "generate_key",
    "thumbprint",
    "thumbprint_uri",
]

register_secp256k1()


class GuestProtocol(t.Protocol):  # pragma: no cover
    def headers(self) -> Header: ...

    def set_kid(self, kid: str) -> None: ...


KeyBase = t.Union[Key, KeySet]
KeyCallable = t.Callable[[GuestProtocol], KeyBase]
KeyFlexible = t.Union[KeyBase, KeyCallable]


def guess_key(
    key: KeyFlexible,
    obj: GuestProtocol,
    random: bool = False,
    use: t.Literal["sig", "enc"] | None = None,
) -> Key:
    """Guess key from a various sources.

    :param key: a very flexible key
    :param obj: a protocol that has ``headers`` and ``set_kid`` methods
    :param random: pick a random key from key set
    :param use: optional "use" value
    """
    resolved_key: KeyBase
    if callable(key):
        resolved_key = key(obj)
    else:
        resolved_key = key

    if isinstance(resolved_key, (OctKey, RSAKey, ECKey, OKPKey)):
        return resolved_key
    elif isinstance(resolved_key, KeySet):
        headers = obj.headers()
        kid: str | None = headers.get("kid")

        parameters: KeyParameters = {"alg": headers["alg"]}
        if use:
            parameters["use"] = use

        if not kid and random:
            # choose one key by random
            return_key = resolved_key.pick_random_key(headers["alg"], parameters)
            if return_key is None:
                raise ValueError("Invalid key")
            return_key.ensure_kid()
            obj.set_kid(t.cast(str, return_key.kid))
        else:
            return_key = resolved_key.get_by_kid(kid, parameters)
        return return_key
    else:
        raise ValueError("Invalid key")


@t.overload
def import_key(data: AnyKey, key_type: t.Literal["oct"], parameters: KeyParameters | None = None) -> OctKey: ...


@t.overload
def import_key(data: AnyKey, key_type: t.Literal["RSA"], parameters: KeyParameters | None = None) -> RSAKey: ...


@t.overload
def import_key(data: AnyKey, key_type: t.Literal["EC"], parameters: KeyParameters | None = None) -> ECKey: ...


@t.overload
def import_key(data: AnyKey, key_type: t.Literal["OKP"], parameters: KeyParameters | None = None) -> OKPKey: ...


@t.overload
def import_key(data: DictKey, key_type: None = None, parameters: KeyParameters | None = None) -> Key: ...


def import_key(
    data: AnyKey,
    key_type: t.Literal["oct", "RSA", "EC", "OKP"] | None = None,
    parameters: KeyParameters | None = None,
) -> Key:
    """Importing a key from bytes, string, and dict. When ``value`` is a dict,
    this method can tell the key type automatically, otherwise, developers
    SHOULD pass the ``key_type`` themselves.

    :param data: the key data in bytes, string, or dict.
    :param key_type: an optional key type in string.
    :param parameters: extra key parameters
    :return: OctKey, RSAKey, ECKey, or OKPKey
    """
    return JWKRegistry.import_key(data, key_type, parameters)


@t.overload
def generate_key(
    key_type: t.Literal["oct"],
    crv_or_size: int | None = None,
    parameters: KeyParameters | None = None,
    private: bool = True,
    auto_kid: bool = False,
) -> OctKey: ...


@t.overload
def generate_key(
    key_type: t.Literal["RSA"],
    crv_or_size: int | None = None,
    parameters: KeyParameters | None = None,
    private: bool = True,
    auto_kid: bool = False,
) -> RSAKey: ...


@t.overload
def generate_key(
    key_type: t.Literal["EC"],
    crv_or_size: t.Literal["P-256", "P-384", "P-521", "secp256k1"] | None = None,
    parameters: KeyParameters | None = None,
    private: bool = True,
    auto_kid: bool = False,
) -> ECKey: ...


@t.overload
def generate_key(
    key_type: t.Literal["OKP"],
    crv_or_size: t.Literal["Ed25519", "Ed448", "X25519", "X448"] | None = None,
    parameters: KeyParameters | None = None,
    private: bool = True,
    auto_kid: bool = False,
) -> OKPKey: ...


def generate_key(
    key_type: t.Literal["oct", "RSA", "EC", "OKP"],
    crv_or_size: str | int | None = None,
    parameters: KeyParameters | None = None,
    private: bool = True,
    auto_kid: bool = False,
) -> Key:
    """Generating key according to the given key type. When ``key_type`` is
    "oct" and "RSA", the second parameter SHOULD be a key size in bits.
    When ``key_type`` is "EC" and "OKP", the second
    parameter SHOULD be a "crv" string.
    """
    return JWKRegistry.generate_key(key_type, crv_or_size, parameters, private, auto_kid)
