# -*- coding: utf-8 -*- ----------------------------------------------------===#
#
#  Copyright 2016-2026 Trovares Inc. dba Rocketgraph.  All rights reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#===------------------------------------------------------------------------===#

"""Public OIDC authentication API.

This module exposes the configuration dataclasses, token cache, and entry-point
functions that users interact with.  HTTP plumbing, discovery, and flow
implementations live in _oidc_flows.py.
"""

from __future__ import annotations

import json

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional

from .common import XgtConnectionError
from ._auth import BearerTokenAuth
from ._oidc_flows import (
  OidcHttpError,
  _now_epoch,
  _default_cache_path,
  _write_json_atomic,
  _token_needs_refresh,
  _interactive_oidc_token_payload,
  _login_oidc_token_payload,
  _refresh_oidc_token,
  _client_credentials_flow,
)

# Re-export OidcHttpError as part of the public surface so callers can import
# it from xgt.oidc without knowing about _oidc_flows.
__all__ = [
  'OidcHttpError',
  'OidcClientConfig',
  'OidcAuth',
  'OidcCredentialsAuth',
  'OidcAuthorizationCodeAuth',
  'OidcTokenCache',
  'cached_oidc_auth',
  'cached_credentials_auth',
]

# ---------------------------------------------------------------------------
# Configuration dataclasses
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class OidcClientConfig:
  """
  Base configuration for OIDC authentication.  Shared by both the interactive
  device-flow (OidcAuth), browser-based authorization code flow
  (OidcAuthorizationCodeAuth), and the headless client-credentials flow
  (OidcCredentialsAuth).

  .. experimental:: The API of this class may change in future releases.

  Parameters
  ----------
  issuer : str
    OIDC issuer URL (e.g. 'https://auth.example.com').  Used to locate the
    discovery document and to validate the 'iss' claim in access tokens.
    If empty, the value is fetched from the server via GetOidcIssuer.
  client_id : str
    OAuth2 client ID registered with the IdP.
    If empty, the value is fetched from the server via GetOidcIssuer.
  scopes : Optional[tuple[str, ...]]
    OAuth2 scopes to request. If None, use server-advertised scopes when
    available, otherwise fall back to ('openid', 'profile', 'email').
  audience : Optional[str]
    Audience value included in the token request.  Required by some IdPs
    (e.g. Auth0).  If None and use_server_audience is True, the value is
    fetched from the server via GetOidcIssuer.
  use_server_audience : bool
    When True and audience is None, retrieve the audience from the server's
    OIDC configuration rather than requiring the caller to supply it.
  ca_cert_path : Optional[str]
    CA bundle to trust for OIDC HTTPS calls (discovery, token, userinfo).
    Supports '~' expansion.
  cache_path : Optional[str]
    Override the default on-disk token cache location.  Supports '~' expansion.
  token_endpoint : Optional[str]
    Override the token endpoint URL.  Skips discovery if set together with
    device_authorization_endpoint.
  device_authorization_endpoint : Optional[str]
    Override the device authorization endpoint URL.  Only used by device flow.
  """
  #: OIDC issuer URL. If empty, fetched from the server via GetOidcIssuer.
  issuer: str = ''
  #: OAuth2 client ID registered with the IdP. If empty, fetched from the server.
  client_id: str = ''
  #: OAuth2 scopes to request. If None, uses server-advertised or default scopes.
  scopes: Optional[tuple[str, ...]] = None
  #: Audience value for the token request. Required by some IdPs (e.g. Auth0).
  audience: Optional[str] = None
  #: When True and audience is None, retrieve the audience from the server's OIDC configuration.
  use_server_audience: bool = False
  #: CA bundle to trust for OIDC HTTPS calls. Supports '~' expansion.
  ca_cert_path: Optional[str] = None
  #: Override the default on-disk token cache location. Supports '~' expansion.
  cache_path: Optional[str] = None
  #: Override the token endpoint URL.
  token_endpoint: Optional[str] = None
  #: Override the device authorization endpoint URL. Only used by device flow.
  device_authorization_endpoint: Optional[str] = None

  def resolved_cache_path(self) -> Path:
    if self.cache_path is not None:
      return Path(self.cache_path).expanduser()
    return _default_cache_path(self.issuer, self.client_id)

@dataclass(frozen=True)
class OidcAuth(OidcClientConfig):
  """
  Interactive OIDC authentication.

  .. experimental:: The API of this class may change in future releases.

  The default flow is 'auto': prefer browser-based authorization code + PKCE
  when the discovery document exposes an authorization endpoint, otherwise fall
  back to device flow when a device authorization endpoint is available.

  Parameters
  ----------
  flow : str
    One of 'auto', 'auth_code', or 'device'.
  force_login : bool
    Ignore any cached token and always start a new login flow.
  open_browser : bool
    Automatically open the verification URL or authorization URL in the default
    browser.
    Set to False in headless environments where a browser is unavailable.
  authorization_endpoint : Optional[str]
    Override the authorization endpoint URL. Used by authorization-code flow.
  redirect_uri : Optional[str]
    Explicit redirect URI registered with the IdP. If omitted, a loopback
    redirect of the form http://127.0.0.1:{callback_port}/callback is used.
  callback_host : str
    Host/interface for the temporary loopback listener.
  callback_port : int
    Port for the temporary loopback listener.
  callback_timeout_seconds : int
    Maximum time to wait for the browser redirect before failing.
  pkce : bool
    When True, use PKCE for the authorization code flow. Set to False for
    providers that support auth-code flow but reject PKCE parameters.
  client_secret : Optional[str]
    Optional client secret for the authorization code token exchange.
    Some providers require this field to be present even for localhost-style
    OAuth clients; OpenShift is one example.
  """
  #: Login flow to use: 'auto', 'auth_code', or 'device'.
  flow: str = 'auto'
  #: When True, ignore any cached token and always start a new login flow.
  force_login: bool = False
  #: When True, automatically open the browser for the login URL.
  open_browser: bool = True
  #: Override the authorization endpoint URL. Used by authorization-code flow.
  authorization_endpoint: Optional[str] = None
  #: Explicit redirect URI. If omitted, a loopback URI on callback_port is used.
  redirect_uri: Optional[str] = None
  #: Host/interface for the temporary loopback listener.
  callback_host: str = '127.0.0.1'
  #: Port for the temporary loopback listener.
  callback_port: int = 8765
  #: Maximum seconds to wait for the browser redirect before failing.
  callback_timeout_seconds: int = 300
  #: When True, use PKCE for the authorization code flow.
  pkce: bool = True
  #: Optional client secret for the authorization code token exchange.
  client_secret: Optional[str] = None

@dataclass(frozen=True)
class OidcCredentialsAuth(OidcClientConfig):
  """
  Client-credentials authentication for headless/scripted OIDC login.

  .. experimental:: The API of this class may change in future releases.

  The IdP issues an access token directly in exchange for the client_id and
  client_secret — no browser, no user interaction.  Use this for CI/CD
  pipelines, cron jobs, or any automated context.

  Parameters
  ----------
  client_secret : str
    Secret registered with the IdP for this client application.
    Should be supplied via an environment variable or secrets manager,
    not hardcoded.
  """
  #: Client secret registered with the IdP. Supply via environment variable, not hardcoded.
  client_secret: str = ''

@dataclass(frozen=True)
class OidcAuthorizationCodeAuth(OidcAuth):
  """Compatibility wrapper for OidcAuth(flow='auth_code').

  .. experimental:: The API of this class may change in future releases.
  """
  #: Fixed to 'auth_code'.
  flow: str = 'auth_code'

# ---------------------------------------------------------------------------
# Token cache
# ---------------------------------------------------------------------------

class OidcTokenCache:
  """Manages a single JSON token file on disk for a given OidcClientConfig.

  .. experimental:: The API of this class may change in future releases.
  """

  def __init__(self, config: OidcClientConfig):
    self._config = config
    self._path = config.resolved_cache_path()

  @property
  def path(self) -> Path:
    return self._path

  def load(self) -> Optional[dict[str, Any]]:
    """Load the cached token payload, or None if unavailable/unusable.

    Returns the payload in one of two states:

    - Fresh access token: full payload, ready to use directly.
    - Expired access token with a refresh token: payload with access_token stripped, signalling the caller to attempt a refresh.

    Returns None (and clears the cache) if the file is missing, corrupt, or
    contains an expired token with no refresh path.
    """
    if not self._path.exists():
      return None
    try:
      with self._path.open('r', encoding='utf-8') as handle:
        token_payload = json.load(handle)
    except (OSError, json.JSONDecodeError):
      self.clear()
      return None

    if not isinstance(token_payload, dict):
      self.clear()
      return None

    access_token = token_payload.get('access_token')
    refresh_token = token_payload.get('refresh_token')
    expires_at = token_payload.get('expires_at')

    if access_token and not _token_needs_refresh(token_payload):
      # Token is still fresh — return as-is.
      return token_payload

    if refresh_token:
      # Access token is expired but we have a refresh token.  Strip the stale
      # access_token so the caller doesn't accidentally use it.
      token_payload = dict(token_payload)
      token_payload.pop('access_token', None)
      return token_payload

    # Expired with no refresh path — cache is useless.
    if expires_at is not None or access_token is not None:
      self.clear()
    return None

  def save(self, token_payload: dict[str, Any]) -> dict[str, Any]:
    """Persist the token payload and annotate it with timing fields.

    Adds 'obtained_at' (epoch) and, if 'expires_in' is present, 'expires_at'
    so that future loads can evaluate freshness without re-decoding the JWT.
    Returns the augmented payload.
    """
    now = _now_epoch()
    stored = dict(token_payload)
    stored['obtained_at'] = now
    expires_in = int(stored.get('expires_in', 0))
    if expires_in > 0:
      stored['expires_at'] = now + expires_in
    _write_json_atomic(self._path, stored)
    return stored

  def clear(self) -> None:
    if self._path.exists():
      self._path.unlink()

# ---------------------------------------------------------------------------
# Public auth entry points
# ---------------------------------------------------------------------------

def cached_oidc_auth(config: OidcClientConfig,
                     force_login: bool = False,
                     open_browser: bool = True) -> BearerTokenAuth:
  """Return a valid BearerTokenAuth, using the cache where possible.

  .. experimental:: The API of this function may change in future releases.

  Resolution order:

  1. Fresh cached access token: return immediately.
  2. Cached refresh token: attempt a silent refresh; on failure, run a new login flow.
  3. No usable cache: run a new login flow.

  If the IdP does not rotate refresh tokens, the existing refresh token is
  preserved in the refreshed payload so subsequent calls can still refresh.
  """
  cache = OidcTokenCache(config)
  token_payload = None if force_login else cache.load()
  if token_payload is None:
    if isinstance(config, OidcAuth):
      token_payload = _interactive_oidc_token_payload(
          config, open_browser=open_browser)
    else:
      token_payload = _login_oidc_token_payload(config,
                                                open_browser=open_browser)
    token_payload = cache.save(token_payload)
    access_token = token_payload.get('access_token')
    if not access_token:
      raise XgtConnectionError('OIDC token response is missing access_token')
    return BearerTokenAuth(access_token)

  access_token = token_payload.get('access_token')
  if access_token and not _token_needs_refresh(token_payload):
    return BearerTokenAuth(access_token)

  refresh_token = token_payload.get('refresh_token')
  if not refresh_token:
    if isinstance(config, OidcAuth):
      token_payload = _interactive_oidc_token_payload(
          config, open_browser=open_browser)
    else:
      token_payload = _login_oidc_token_payload(config,
                                                open_browser=open_browser)
    token_payload = cache.save(token_payload)
    access_token = token_payload.get('access_token')
    if not access_token:
      raise XgtConnectionError('OIDC token response is missing access_token')
    return BearerTokenAuth(access_token)

  try:
    # Only the token endpoint is needed for refresh; use the lighter helper
    # rather than _discover_oidc_endpoints, which requires a device endpoint.
    refreshed = _refresh_oidc_token(config, refresh_token)
  except XgtConnectionError:
    # Refresh token is expired or revoked — start over.
    cache.clear()
    if isinstance(config, OidcAuth):
      token_payload = _interactive_oidc_token_payload(
          config, open_browser=open_browser)
    else:
      token_payload = _login_oidc_token_payload(config,
                                                open_browser=open_browser)
    token_payload = cache.save(token_payload)
    access_token = token_payload.get('access_token')
    if not access_token:
      raise XgtConnectionError('OIDC token response is missing access_token')
    return BearerTokenAuth(access_token)

  # Some IdPs do not return a new refresh token on every refresh; preserve the
  # existing one so the next call can still refresh without a full login.
  if 'refresh_token' not in refreshed:
    refreshed['refresh_token'] = refresh_token
  refreshed = cache.save(refreshed)
  access_token = refreshed.get('access_token')
  if not access_token:
    raise XgtConnectionError('OIDC refresh response is missing access_token')
  return BearerTokenAuth(access_token)

def cached_credentials_auth(config: OidcCredentialsAuth) -> BearerTokenAuth:
  """Return a valid BearerTokenAuth for a service client, using the cache.

  .. experimental:: The API of this function may change in future releases.

  Unlike cached_oidc_auth there is no interactive fallback: when the cached
  token is expired the client credentials flow is run again automatically.
  No refresh token handling is needed since the client can always re-authenticate.
  """
  cache = OidcTokenCache(config)
  token_payload = cache.load()
  if token_payload is not None:
    access_token = token_payload.get('access_token')
    if access_token and not _token_needs_refresh(token_payload):
      return BearerTokenAuth(access_token)

  token_payload = _client_credentials_flow(config)
  token_payload = cache.save(token_payload)
  access_token = token_payload.get('access_token')
  if not access_token:
    raise XgtConnectionError(
        'OIDC client credentials response is missing access_token')
  return BearerTokenAuth(access_token)
