import json
import urllib.request
from functools import lru_cache
from typing import Any, List
from urllib.parse import urlparse
from urllib.error import HTTPError, URLError
from .api_jwk import PyJWK, PyJWKSet
from .api_jwt import decode_complete as decode_token
from .exceptions import PyJWKClientError, PyJWKClientConnectionError


class _NoRedirectHandler(urllib.request.HTTPRedirectHandler):
    def redirect_request(
        self,
        req: urllib.request.Request,
        fp: Any,
        code: int,
        msg: str,
        headers: Any,
        newurl: str,
    ) -> None:
        return None


class PyJWKClient:
    def __init__(self, uri: str, cache_keys: bool = True, max_cached_keys: int = 16):
        # urllib's default OpenerDirector also handles file://, ftp://, and
        # data: URIs. Reject anything that isn't http(s) eagerly so a caller
        # passing an attacker-influenced URL (e.g. taken from a `jku` token
        # header) can't read local files or reach other unintended schemes.
        scheme = urlparse(uri).scheme.lower()
        if scheme not in ("http", "https"):
            raise PyJWKClientError(
                f"Invalid JWKS URI scheme {scheme!r}: only 'http' and 'https' "
                f"are supported."
            )
        self.uri = uri
        if cache_keys:
            # Cache signing keys
            # Ignore mypy (https://github.com/python/mypy/issues/2427)
            self.get_signing_key = lru_cache(maxsize=max_cached_keys)(self.get_signing_key)  # type: ignore

    def fetch_data(self) -> Any:
        try:
            handlers: list[Any] = [_NoRedirectHandler()]
            opener = urllib.request.build_opener(*handlers)
            with opener.open(self.uri) as response:
                return json.load(response)
        except (URLError, TimeoutError) as e:
            if isinstance(e, HTTPError):
                e.close()
            raise PyJWKClientConnectionError(
                f'Fail to fetch data from the url, err: "{e}"'
            ) from e

    def get_jwk_set(self) -> PyJWKSet:
        data = self.fetch_data()
        return PyJWKSet.from_dict(data)

    def get_signing_keys(self) -> List[PyJWK]:
        jwk_set = self.get_jwk_set()
        signing_keys = [
            jwk_set_key
            for jwk_set_key in jwk_set.keys
            if jwk_set_key.public_key_use in ["sig", None] and jwk_set_key.key_id
        ]

        if not signing_keys:
            raise PyJWKClientError("The JWKS endpoint did not contain any signing keys")

        return signing_keys

    def get_signing_key(self, kid: str) -> PyJWK:
        signing_keys = self.get_signing_keys()
        signing_key = None

        for key in signing_keys:
            if key.key_id == kid:
                signing_key = key
                break

        if not signing_key:
            raise PyJWKClientError(
                f'Unable to find a signing key that matches: "{kid}"'
            )

        return signing_key

    def get_signing_key_from_jwt(self, token: str) -> PyJWK:
        unverified = decode_token(token, options={"verify_signature": False})
        header = unverified["header"]
        return self.get_signing_key(header.get("kid"))