"""Cached version of PyJWKClient.

Caches in ~/.slices/cache.
"""

import base64
import json
from datetime import timedelta

import httpx
from httpx import HTTPError
from jwt import PyJWKClient, PyJWKClientConnectionError

from slices_clientlib_core.base.cache import clear_cache_entry, load_cache_entry, save_cache_entry
from slices_clientlib_core.base.config import DEFAULT_HTTP_USER_AGENT
from slices_clientlib_core.base.const import BASE_HTTP_TIMEOUT_S
from slices_clientlib_core.util.force_ip_version import get_httpx_transport_for_ip_version
from slices_clientlib_core.util.logging_httpx_client import enable_httpx_client_logging


class CachingPyJWKClient(PyJWKClient):
    """Cached version of PyJWKClient.

    Caches in ~/.slices/cache.
    """

    def __init__(self, *args, **kwargs):
        """Create CachingPyJWKClient."""
        super().__init__(*args, **kwargs)

        self.client = enable_httpx_client_logging(
            httpx.Client(
                headers={"User-Agent": DEFAULT_HTTP_USER_AGENT, "Accept": "application/json"},
                follow_redirects=True,
                max_redirects=5,
                timeout=httpx.Timeout(timeout=BASE_HTTP_TIMEOUT_S),
                transport=get_httpx_transport_for_ip_version(),
            )
        )

    def fetch_data(self):
        """Fetch the JWK data."""

        uri_b64 = base64.b64encode(bytes(self.uri, "utf-8")).decode()

        cache_entry = load_cache_entry("jwk", uri_b64, timedelta(hours=6))
        if cache_entry:
            try:
                return json.loads(cache_entry.content)
            except json.JSONDecodeError:
                clear_cache_entry("jwk", uri_b64)
                raise

        try:
            response = self.client.get(self.uri, headers=self.headers, timeout=self.timeout)
            response.raise_for_status()
            fresh_data = response.json()
            save_cache_entry("jwk", uri_b64, json.dumps(fresh_data))
        except (HTTPError, TimeoutError) as e:
            clear_cache_entry("jwk", uri_b64)
            raise PyJWKClientConnectionError(f'Fail to fetch data from the url, err: "{e}"') from e
        else:
            return fresh_data

        # see jwks_client.py for original code
        # r = urllib.request.Request(url=self.uri, headers=self.headers)
        # with urllib.request.urlopen(
        #         r, timeout=self.timeout, context=self.ssl_context
        # ) as response:
        #     return json.load(response)
