"""Slices Client."""

import base64
import json
from datetime import datetime, timedelta
from functools import cached_property
from pathlib import Path
from uuid import uuid4

import httpx
import jwt
import msgspec
from httpx import HTTPError
from pytz import UTC

from slices_clientlib_core.base.auth import (
    MINIMUM_LIFETIME_FOR_FRESH_TOKEN,
    ensure_auth_token_is_fresh,
    load_auth_token,
    parse_auth_token,
)
from slices_clientlib_core.base.cache import load_cache_entry, save_cache_entry
from slices_clientlib_core.base.config import (
    DEFAULT_HTTP_USER_AGENT,
    ApiEntry,
    SlicesCentralApiClientsConfig,
    load_default_central_api_clients_config,
)
from slices_clientlib_core.base.exceptions import ApiException, from_httpx_exception
from slices_clientlib_core.util.utils import json_encoder

from .base.const import BASE_HTTP_TIMEOUT_S
from .experiment import ExperimentClient
from .oauth2.token import AuthJWT, OAuthToken
from .oauth2.token_exchange import (
    ExchangedToken,
    get_access_token_for_audience,
    get_id_token_for_audience,
    token_decoder,
)
from .project import ProjectClient
from .pubkey import PublicKeyClient
from .user import UserClient
from .util.force_ip_version import get_httpx_transport_for_ip_version
from .util.logging_httpx_client import enable_httpx_client_logging


class NotLoggedInException(Exception):
    """Not logged in."""


def add_request_id_header(request):
    """Add Request ID header."""
    request.headers["X-Request-ID"] = str(uuid4())
    return request


class OAuthHttpxAuth(httpx.Auth):
    """OAuth HTTPX Auth."""

    _token_for_audience: ExchangedToken | None = None
    _token_expires_at: datetime | None = None

    def __init__(self, slices_client: "SlicesClient", target_audience: str, project_id: str | None = None) -> None:
        """Create HTTPX Auth provider that sets the appropriate Authorization-header."""
        self.slices_client = slices_client
        self.target_audience = target_audience
        self.project_id = project_id

    def auth_flow(self, request):
        """Auth flow."""

        aud_b64 = base64.b64encode(bytes(self.target_audience, "utf-8")).decode()
        cache_id = f"{aud_b64}___{self.project_id}" if self.project_id else f"{aud_b64}___"

        if not self._token_for_audience:
            # Try to find a cached token
            cache_entry = load_cache_entry("access_token", cache_id, timedelta(minutes=10))
            if cache_entry:
                try:
                    self._token_for_audience = token_decoder.decode(cache_entry.content)
                    self._token_expires_at = datetime.now(UTC) + timedelta(seconds=self._token_for_audience.expires_in)
                except:  # noqa E722
                    # Ignore failures processing cache entry.
                    # This will cause silent fall back to not caching.
                    self._token_for_audience = None
                    self._token_expires_at = None

        if not self._token_expires_at or (
            datetime.now(UTC) > self._token_expires_at + MINIMUM_LIFETIME_FOR_FRESH_TOKEN
        ):
            # This token is scoped to the CLI
            auth_token = self.slices_client._get_fresh_auth_token()

            # So we need to exchange it for a token that is scoped to the target audience
            self._token_for_audience = get_access_token_for_audience(auth_token, self.target_audience, self.project_id)
            self._token_expires_at = datetime.now(UTC) + timedelta(seconds=self._token_for_audience.expires_in)

            save_cache_entry("access_token", cache_id, json_encoder.encode(self._token_for_audience))

        assert self._token_for_audience

        request.headers["Authorization"] = f"Bearer {self._token_for_audience.access_token}"

        # Add a request ID header
        # While not directly related to authentication, we do it here because
        # it's an easy location to mutate the request.
        add_request_id_header(request)
        yield request


def _load_file_or_url(location: str, client: httpx.Client) -> str:
    if location.startswith("http"):
        response = client.get(location, headers={"Accept": "application/json"})
        try:
            response.raise_for_status()
            return response.content.decode()
        except HTTPError as err:
            raise from_httpx_exception(err) from err
    elif Path(location).exists():
        with Path(location).open() as f:
            return f.read()
    else:
        raise ValueError(f"{location!r} is not a valid path or URL")


class SlicesClient:
    """Slices Client."""

    _auth_token: OAuthToken | None
    _central_api_clients_config: SlicesCentralApiClientsConfig

    def __init__(
        self,
        *,
        central_config: SlicesCentralApiClientsConfig | str | None = None,
        ignore_config_cache: bool = False,
        http_user_agent: str = DEFAULT_HTTP_USER_AGENT,
    ) -> None:
        """Create Slices Client.

        You can optionally provide a central slices config to override the default.
        This accepts flexible types, one of:
           - A SlicesCentralApiClientsConfig config object (= directly specify config)
           - A string with a URL (from which the config will be downloaded)
           - A string with a filename (from which the config will be read)
        """

        # _no_auth_client is a client without any authentication, which can be used to download the central config
        self.http_user_agent = http_user_agent

        self._no_auth_client = enable_httpx_client_logging(
            httpx.Client(
                headers={"User-Agent": 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(),
            )
        )

        self._auth_token = load_auth_token()

        self._init_param_central_config = central_config

        # Load central config
        if central_config and isinstance(central_config, str):
            self._central_api_clients_config = msgspec.json.decode(
                _load_file_or_url(central_config, self._no_auth_client), type=SlicesCentralApiClientsConfig
            )
        elif central_config and isinstance(central_config, SlicesCentralApiClientsConfig):
            self._central_api_clients_config = central_config
        elif central_config is None:
            self._central_api_clients_config = load_default_central_api_clients_config(
                self._no_auth_client, ignore_cache=ignore_config_cache
            )
        else:
            raise TypeError(f"Unsupported central_config {central_config!r}")

    def reload_central_config(self) -> bool:
        """Reload the central slices config from the default online source.

        This is useful because it is locally cached.

        This does nothing if the central slices config was not loaded from
        the default online source or its cached version.

        Returns: True if reload was done, False if reload was skipped.
        """
        if self._init_param_central_config is not None:
            return False
        self._central_api_clients_config = load_default_central_api_clients_config(
            self._no_auth_client, ignore_cache=True
        )
        return True

    def _get_auth(self, audience: str, project_id: str | None = None, experiment_id: str | None = None):
        if experiment_id:
            if not project_id:
                raise ValueError("project_id must be set when experiment_id is set")
            return ExperimentJwtHttpxAuth(self, experiment_id, project_id)

        return OAuthHttpxAuth(self, audience, project_id)

    def _create_httpx_client(
        self, *, base_url: str, audience: str, project_id: str | None = None, experiment_id: str | None = None
    ) -> httpx.Client:
        return enable_httpx_client_logging(
            httpx.Client(
                base_url=base_url,
                auth=self._get_auth(audience, project_id, experiment_id),
                headers={"User-Agent": self.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 _create_httpx_client_for_api_entry(
        self,
        api_entry: ApiEntry,
        project_id: str | None = None,
        experiment_id: str | None = None,
    ) -> httpx.Client:
        return self._create_httpx_client(
            base_url=api_entry.base_url, audience=api_entry.audience, project_id=project_id, experiment_id=experiment_id
        )

    def _get_fresh_auth_token(self):
        if not self._auth_token:
            raise NotLoggedInException("Not logged in")

        self._auth_token = ensure_auth_token_is_fresh(self._auth_token)
        return self._auth_token

    @property
    def user_info(self) -> AuthJWT | None:
        """Auth Token."""
        return self._auth_token and parse_auth_token(self._auth_token)

    @cached_property
    def project(self):
        """Project API Client."""
        return ProjectClient(
            self._create_httpx_client_for_api_entry(
                self._central_api_clients_config.central_hub_services.project_authority
            )
        )

    @cached_property
    def user(self):
        """User API Client."""
        return UserClient(
            self._create_httpx_client_for_api_entry(
                self._central_api_clients_config.central_hub_services.user_authority
            )
        )

    @cached_property
    def experiment(self):
        """Experiment API Client."""
        return ExperimentClient(
            self._create_httpx_client_for_api_entry(
                self._central_api_clients_config.central_hub_services.experiment_authority
            ),
            project_client=self.project,
        )

    @cached_property
    def public_key(self):
        """Public Key API Client."""
        return PublicKeyClient(
            self._create_httpx_client_for_api_entry(
                self._central_api_clients_config.central_hub_services.pubkey_authority
            )
        )

    def for_project(self, project_id: str) -> "SlicesProjectClient":
        """Get Slices Project Client."""
        return SlicesProjectClient(project_id, self)

    def for_experiment(self, experiment_id: str, project_id: str) -> "SlicesExperimentClient":
        """Get Slices Experiment Client."""
        return SlicesExperimentClient(experiment_id, project_id, self)

    def get_id_token(self, audience: str | list[str], project_id: str) -> str:
        """Get ID Token."""
        if not self._auth_token:
            raise NotLoggedInException("Not logged in")

        return get_id_token_for_audience(self._auth_token, audience, project_id)


class SlicesProjectClient:
    """Slices client for a specific project."""

    def __init__(self, project_id: str, generic_client: SlicesClient) -> None:
        """Create Slices Project Client."""
        super().__init__()
        self.project_id = project_id
        self._generic_client = generic_client

    def for_experiment(self, experiment_id: str) -> "SlicesExperimentClient":
        """Get Slices Experiment Client."""
        return SlicesExperimentClient(experiment_id, self.project_id, self._generic_client)


class ExperimentJwtHttpxAuth(httpx.Auth):
    """HTTPX Auth implementation which fetches a JWT for a specific experiment."""

    _token: str | None = None
    _token_expires_at: datetime | None = None

    def __init__(self, slices_client: SlicesClient, experiment_id: str, project_id: str) -> None:
        """Experiment JWT Authentication for HTTPX."""
        super().__init__()
        self.slices_client = slices_client
        self.experiment_id = experiment_id
        self.project_id = project_id

    # Since this is not used, it would be nice to remove it
    # @deprecated("This is not used anywhere in the client itself or CLI. Will be removed at some point.")
    def get_experiment_jwt(self, experiment_id: str) -> str:
        """Get Experiment JWT."""
        return self.slices_client.experiment.get_jwt(experiment_id)

    def auth_flow(self, request):
        """Auth flow."""

        if not self._token:
            # Try to find a cached token
            cache_entry = load_cache_entry("experiment_token", self.experiment_id, timedelta(minutes=10))
            if cache_entry:
                try:
                    self._token = cache_entry.content
                    decoded_token = jwt.decode(self._token, options={"verify_signature": False})
                    self._token_expires_at = datetime.fromtimestamp(decoded_token["exp"], UTC)
                except:  # noqa E722
                    # Ignore failures processing cache entry.
                    # This will cause silent fall back to not caching.
                    self._token = None
                    self._token_expires_at = None

        if not self._token_expires_at or (
            datetime.now(UTC) > self._token_expires_at - MINIMUM_LIFETIME_FOR_FRESH_TOKEN
        ):
            # We request a new token for the experiment
            self._token = self.slices_client.experiment.get_jwt(self.experiment_id)

            decoded_token = jwt.decode(self._token, options={"verify_signature": False})

            if (exp_proj := decoded_token.get("proj")) != self.project_id:
                raise ApiException(
                    f"'{self.experiment_id}' is an experiment in project '{exp_proj}', "
                    f"but was expecting it to be from project '{self.project_id}'"
                )

            self._token_expires_at = datetime.fromtimestamp(decoded_token["exp"], UTC)
            # Check if new token expire is OK.
            if not self._token_expires_at:
                raise ApiException(
                    f"Refreshed experiment token for '{self.experiment_id}' has no exp: {json.dumps(decoded_token)}"
                )
            if datetime.now(UTC) > self._token_expires_at + MINIMUM_LIFETIME_FOR_FRESH_TOKEN:
                raise ApiException(
                    f"Refreshed experiment token for '{self.experiment_id}' is not valid long enough. "
                    f"exp='{self._token_expires_at.isoformat()}' now='{datetime.now(UTC).isoformat()}' "
                    f"token={json.dumps(decoded_token)}"
                )

            save_cache_entry("experiment_token", self.experiment_id, self._token)

        assert self._token

        request.headers["Authorization"] = f"Bearer {self._token}"

        yield request


class SlicesExperimentClient:
    """Slices client for a specific experiment."""

    def __init__(self, experiment_id: str, project_id: str, generic_client: SlicesClient) -> None:
        """Create Slices Project Client."""
        super().__init__()
        self.experiment_id = experiment_id
        self.project_id = project_id
        self._generic_client = generic_client
