# -*- 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.
#
#===------------------------------------------------------------------------===#

"""Internal OIDC flow implementations.

This module is private to the xgt package.  All public symbols live in
oidc.py.  The split keeps HTTP plumbing, discovery, and flow logic separate
from the dataclasses and public API that users interact with.
"""

from __future__ import annotations

import base64
import hashlib
import http.server
import json
import os
import secrets
import ssl
import time
import webbrowser

from dataclasses import replace
from pathlib import Path
from typing import Any, Optional
from urllib import error, parse, request

from .common import XgtConnectionError, XgtSecurityError

# Type-annotation-only references to avoid a circular import.  These classes
# are defined in oidc.py, which imports from this module.  At runtime the
# annotations are strings (PEP 563), so no actual import is needed here.
# fmt: off
# OidcClientConfig, OidcAuth, OidcCredentialsAuth  (referenced in signatures only)
# fmt: on

_DEFAULT_SCOPES = ('openid', 'profile', 'email')

# Treat a token as needing refresh this many seconds before it actually expires.
# Avoids presenting a token to the server that expires in transit.
_REFRESH_SKEW_SECONDS = 60

# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------

class OidcHttpError(XgtConnectionError):
  """Raised for HTTP-level OIDC errors.  Carries the raw OAuth2 error code
  (e.g. 'authorization_pending', 'slow_down') so callers can branch on it
  without parsing the error message string."""
  def __init__(self, message: str,
               oidc_error: Optional[str] = None) -> None:
    super().__init__(message)
    self.oidc_error = oidc_error
    self.oidc_description = None

# ---------------------------------------------------------------------------
# File utilities
# ---------------------------------------------------------------------------

def _now_epoch() -> int:
  return int(time.time())

def _default_cache_path(issuer: str, client_id: str) -> Path:
  # Hash the issuer + client_id pair so different IdP/client combinations
  # get separate cache files without exposing the raw values in the filename.
  digest = hashlib.sha256(f'{issuer}|{client_id}'.encode('utf-8')).hexdigest()
  return Path.home() / '.rocketgraph' / 'oidc' / f'{digest}.json'

def _ensure_parent_dir(path: Path) -> None:
  path.parent.mkdir(parents=True, exist_ok=True)

def _write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
  # Write to a sibling .tmp file first, then rename into place.  os.replace()
  # is atomic on POSIX, so a concurrent reader never sees a partial write.
  # Mode 0o600 ensures the token file is readable only by the owning user.
  _ensure_parent_dir(path)
  tmp_path = path.with_suffix(path.suffix + '.tmp')
  fd = os.open(tmp_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
  with os.fdopen(fd, 'w', encoding='utf-8') as handle:
    json.dump(payload, handle, indent=2, sort_keys=True)
    handle.write('\n')
  os.replace(tmp_path, path)

# ---------------------------------------------------------------------------
# HTTP transport
# ---------------------------------------------------------------------------

def _ssl_context(config: Optional[OidcClientConfig] = None
                 ) -> Optional[ssl.SSLContext]:
  if config is None or config.ca_cert_path is None:
    return None
  return ssl.create_default_context(
      cafile=str(Path(config.ca_cert_path).expanduser()))

def _http_post_form(url: str, form_data: dict[str, Any],
                    config: Optional[OidcClientConfig] = None
                    ) -> dict[str, Any]:
  encoded = parse.urlencode(form_data).encode('utf-8')
  req = request.Request(url, data=encoded,
                        headers={'Content-Type':
                                 'application/x-www-form-urlencoded'})
  try:
    with request.urlopen(req, context=_ssl_context(config)) as response:
      return json.loads(response.read().decode('utf-8'))
  except error.HTTPError as ex:
    details = ex.read().decode('utf-8')
    try:
      payload = json.loads(details)
      error_value = payload.get('error', ex.reason)
      description = payload.get('error_description', details)
      # Raise OidcHttpError so callers can inspect the structured error code.
      oidc_error = OidcHttpError(
          f'OIDC HTTP request failed: {error_value}: {description}',
          oidc_error=error_value)
      oidc_error.oidc_description = description
      raise oidc_error from ex
    except json.JSONDecodeError:
      raise XgtConnectionError(
          f'OIDC HTTP request failed: {ex.reason}: {details}') from ex
  except error.URLError as ex:
    raise XgtConnectionError(f'OIDC HTTP request failed: {ex.reason}') from ex

def _http_get_json(url: str,
                   config: Optional[OidcClientConfig] = None) -> dict[str, Any]:
  req = request.Request(url)
  try:
    with request.urlopen(req, context=_ssl_context(config)) as response:
      return json.loads(response.read().decode('utf-8'))
  except error.HTTPError as ex:
    details = ex.read().decode('utf-8')
    raise XgtConnectionError(
        f'OIDC metadata request failed: {ex.reason}: {details}') from ex
  except error.URLError as ex:
    raise XgtConnectionError(
        f'OIDC metadata request failed: {ex.reason}') from ex

# ---------------------------------------------------------------------------
# Discovery helpers
# ---------------------------------------------------------------------------

def _oidc_metadata_url(issuer: str) -> str:
  return issuer.rstrip('/') + '/.well-known/openid-configuration'

def _openshift_oauth_metadata_url(issuer: str) -> str:
  return issuer.rstrip('/') + '/.well-known/oauth-authorization-server'

def _openshift_authorization_endpoint(issuer: str) -> str:
  return issuer.rstrip('/') + '/oauth/authorize'

def _openshift_token_endpoint(issuer: str) -> str:
  return issuer.rstrip('/') + '/oauth/token'

def _discover_metadata(config: OidcClientConfig) -> dict[str, Any]:
  """Fetch provider metadata.

  Prefer standard OIDC discovery. If that is missing, fall back to the
  OpenShift OAuth metadata document.
  """
  try:
    return _http_get_json(_oidc_metadata_url(config.issuer), config=config)
  except XgtConnectionError as ex:
    # urllib surfaces the HTTP reason phrase in the exception message.
    # Only swallow 404s (standard OIDC discovery missing); re-raise anything
    # else (network errors, auth failures, etc.) so they aren't masked.
    if 'Not Found' not in str(ex):
      raise
  return _http_get_json(_openshift_oauth_metadata_url(config.issuer),
                        config=config)

def _discover_oidc_endpoints(config: OidcClientConfig) -> dict[str, str]:
  """Resolve token and device authorization endpoints for the device flow.

  If both endpoints are supplied directly in config, discovery is skipped.
  Otherwise the OIDC discovery document is fetched from
  {issuer}/.well-known/openid-configuration.
  """
  if not config.issuer:
    raise XgtConnectionError('OIDC issuer is not configured')
  if not config.client_id:
    raise XgtConnectionError('OIDC client_id is not configured')
  if (config.token_endpoint is not None and
      config.device_authorization_endpoint is not None):
    return {
      'token_endpoint': config.token_endpoint,
      'device_authorization_endpoint': config.device_authorization_endpoint,
    }
  try:
    metadata = _discover_metadata(config)
  except XgtConnectionError:
    metadata = {}
  token_endpoint = config.token_endpoint or metadata.get('token_endpoint')
  device_endpoint = (config.device_authorization_endpoint or
                     metadata.get('device_authorization_endpoint'))
  if not token_endpoint:
    raise XgtConnectionError('OIDC discovery document is missing token_endpoint')
  if not device_endpoint:
    raise XgtConnectionError(
        'OIDC discovery document is missing device_authorization_endpoint')
  return {
    'token_endpoint': token_endpoint,
    'device_authorization_endpoint': device_endpoint,
  }

def _discover_authorization_code_endpoints(
    config: OidcAuth) -> dict[str, str]:
  """Resolve token and authorization endpoints for auth-code flow."""
  if not config.issuer:
    raise XgtConnectionError('OIDC issuer is not configured')
  if not config.client_id:
    raise XgtConnectionError('OIDC client_id is not configured')
  if (config.token_endpoint is not None and
      config.authorization_endpoint is not None):
    return {
      'token_endpoint': config.token_endpoint,
      'authorization_endpoint': config.authorization_endpoint,
    }
  # Swallow discovery failures: providers like OpenShift do not publish a
  # standard OIDC discovery document, so we fall back to their well-known
  # endpoint paths below.
  try:
    metadata = _discover_metadata(config)
  except XgtConnectionError:
    metadata = {}
  token_endpoint = config.token_endpoint or metadata.get('token_endpoint')
  authorization_endpoint = (config.authorization_endpoint or
                            metadata.get('authorization_endpoint'))
  if token_endpoint is None:
    token_endpoint = _openshift_token_endpoint(config.issuer)
  if authorization_endpoint is None:
    authorization_endpoint = _openshift_authorization_endpoint(config.issuer)
  return {
    'token_endpoint': token_endpoint,
    'authorization_endpoint': authorization_endpoint,
  }

def _discover_interactive_endpoints(config: OidcAuth) -> dict[str, Optional[str]]:
  """Resolve interactive-login endpoints for flow selection."""
  if not config.issuer:
    raise XgtConnectionError('OIDC issuer is not configured')
  if not config.client_id:
    raise XgtConnectionError('OIDC client_id is not configured')
  # Same OpenShift fallback as _discover_authorization_code_endpoints.
  try:
    metadata = _discover_metadata(config)
  except XgtConnectionError:
    metadata = {}
  token_endpoint = config.token_endpoint or metadata.get('token_endpoint')
  authorization_endpoint = (config.authorization_endpoint or
                            metadata.get('authorization_endpoint'))
  device_endpoint = (config.device_authorization_endpoint or
                     metadata.get('device_authorization_endpoint'))
  if token_endpoint is None:
    token_endpoint = _openshift_token_endpoint(config.issuer)
  if authorization_endpoint is None:
    authorization_endpoint = _openshift_authorization_endpoint(config.issuer)
  return {
    'token_endpoint': token_endpoint,
    'authorization_endpoint': authorization_endpoint,
    'device_authorization_endpoint': device_endpoint,
  }

# ---------------------------------------------------------------------------
# Device flow
# ---------------------------------------------------------------------------

def _scope_string(config: OidcClientConfig) -> str:
  return ' '.join(config.scopes or _DEFAULT_SCOPES)

def _prefer_device_flow_for_environment() -> bool:
  # SSH sessions have no local browser, so device flow (user_code + URL) is
  # more practical than trying to open a browser on the remote host.
  return bool(os.environ.get('SSH_CONNECTION') or
              os.environ.get('SSH_TTY'))

def _device_authorization_request(config: OidcClientConfig,
                                  endpoints: Optional[dict[str, str]] = None
                                  ) -> dict[str, Any]:
  """Send the device authorization request and return the IdP's response.

  The response contains device_code, user_code, verification_uri, and
  expires_in / interval for the polling loop.
  """
  endpoints = endpoints or _discover_oidc_endpoints(config)
  request_payload = {
    'client_id': config.client_id,
    'scope': _scope_string(config),
  }
  if config.audience is not None:
    request_payload['audience'] = config.audience
  return _http_post_form(endpoints['device_authorization_endpoint'],
                         request_payload, config=config)

def _poll_for_tokens(config: OidcClientConfig,
                     device_code_payload: dict[str, Any],
                     endpoints: Optional[dict[str, str]] = None
                     ) -> dict[str, Any]:
  """Poll the token endpoint until the user approves or the code expires.

  Implements the RFC 8628 polling protocol:
  - 'authorization_pending': user hasn't approved yet — wait and retry.
  - 'slow_down': IdP is throttling — increase the interval by 5 s and retry.
  - Any other error: propagate immediately.
  """
  endpoints = endpoints or _discover_oidc_endpoints(config)
  deadline = _now_epoch() + int(device_code_payload.get('expires_in', 600))
  interval = int(device_code_payload.get('interval', 5))
  request_payload = {
    'grant_type': 'urn:ietf:params:oauth:grant-type:device_code',
    'client_id': config.client_id,
    'device_code': device_code_payload['device_code'],
  }
  while _now_epoch() < deadline:
    try:
      return _http_post_form(endpoints['token_endpoint'], request_payload,
                             config=config)
    except OidcHttpError as ex:
      if ex.oidc_error == 'authorization_pending':
        time.sleep(interval)
        continue
      if ex.oidc_error == 'slow_down':
        interval += 5
        time.sleep(interval)
        continue
      raise
  raise XgtSecurityError('OIDC device login timed out before authorization')

def _login_oidc_token_payload(config: OidcClientConfig,
                              open_browser: bool = True) -> dict[str, Any]:
  """Run a complete device flow and return the raw token payload."""
  endpoints = _discover_oidc_endpoints(config)
  device_payload = _device_authorization_request(config, endpoints=endpoints)
  verification_url = (device_payload.get('verification_uri_complete') or
                      device_payload.get('verification_uri'))
  if not verification_url:
    raise XgtConnectionError(
        'OIDC device authorization response is missing verification URL')
  user_code = device_payload.get('user_code')
  print(f'Open this URL to authenticate: {verification_url}')
  if user_code:
    print(f'User code: {user_code}')
  if open_browser:
    webbrowser.open(verification_url)
  return _poll_for_tokens(config, device_payload, endpoints=endpoints)

# ---------------------------------------------------------------------------
# Token refresh
# ---------------------------------------------------------------------------

def _token_needs_refresh(token_payload: dict[str, Any]) -> bool:
  """Return True if the access token is expired or expiry is unknown."""
  expires_at = token_payload.get('expires_at')
  if expires_at is None:
    return True
  # Treat the token as expired slightly early to avoid presenting it to the
  # server when it may expire in transit.
  return int(expires_at) <= (_now_epoch() + _REFRESH_SKEW_SECONDS)

def _refresh_oidc_token(config: OidcClientConfig,
                        refresh_token: str,
                        endpoints: Optional[dict[str, str]] = None
                        ) -> dict[str, Any]:
  if endpoints is None:
    endpoints = {'token_endpoint': _get_token_endpoint(config)}
  request_payload = {
    'grant_type': 'refresh_token',
    'client_id': config.client_id,
    'refresh_token': refresh_token,
  }
  if config.audience is not None:
    request_payload['audience'] = config.audience
  return _http_post_form(endpoints['token_endpoint'], request_payload,
                         config=config)

# ---------------------------------------------------------------------------
# Authorization code flow
# ---------------------------------------------------------------------------

def _pkce_code_challenge(code_verifier: str) -> str:
  digest = hashlib.sha256(code_verifier.encode('utf-8')).digest()
  return base64.urlsafe_b64encode(digest).rstrip(b'=').decode('ascii')

def _authorization_code_redirect_uri(config: OidcAuth) -> str:
  if config.redirect_uri is not None:
    return config.redirect_uri
  return f'http://{config.callback_host}:{config.callback_port}/callback'

class _AuthorizationCodeCallbackServer(http.server.HTTPServer):
  allow_reuse_address = True

class _AuthorizationCodeCallbackHandler(http.server.BaseHTTPRequestHandler):
  def do_GET(self) -> None:
    parsed = parse.urlparse(self.path)
    if parsed.path != '/callback':
      self.send_response(404)
      self.end_headers()
      return

    self.server.query = parse.parse_qs(parsed.query)
    self.send_response(200)
    self.send_header('Content-Type', 'text/html; charset=utf-8')
    self.end_headers()
    self.wfile.write(
        b'<html><body><p>Authentication complete. You can close this '
        b'window.</p></body></html>')

  def log_message(self, format: str, *args: Any) -> None:
    del format, args

def _authorization_code_flow(
    config: OidcAuth,
    open_browser: Optional[bool] = None) -> dict[str, Any]:
  endpoints = _discover_authorization_code_endpoints(config)
  redirect_uri = _authorization_code_redirect_uri(config)
  state = secrets.token_urlsafe(24)
  code_verifier = secrets.token_urlsafe(64)
  authorization_params = {
    'response_type': 'code',
    'client_id': config.client_id,
    'redirect_uri': redirect_uri,
    'scope': _scope_string(config),
    'state': state,
  }
  if config.pkce:
    authorization_params['code_challenge'] = \
        _pkce_code_challenge(code_verifier)
    authorization_params['code_challenge_method'] = 'S256'
  if config.audience is not None:
    authorization_params['audience'] = config.audience

  authorization_url = (
      endpoints['authorization_endpoint'] + '?' +
      parse.urlencode(authorization_params))

  parsed_redirect = parse.urlparse(redirect_uri)
  if parsed_redirect.scheme != 'http':
    raise XgtConnectionError(
        'OidcAuthorizationCodeAuth requires an http:// redirect_uri')
  if parsed_redirect.hostname != config.callback_host:
    raise XgtConnectionError(
        'redirect_uri host must match callback_host for loopback auth')
  if parsed_redirect.port != config.callback_port:
    raise XgtConnectionError(
        'redirect_uri port must match callback_port for loopback auth')

  server = _AuthorizationCodeCallbackServer(
      (config.callback_host, config.callback_port),
      _AuthorizationCodeCallbackHandler)
  server.timeout = 1.0
  server.query = None
  deadline = time.time() + config.callback_timeout_seconds
  print(f'Open this URL to authenticate: {authorization_url}')
  should_open_browser = (open_browser if open_browser is not None
                         else config.open_browser)
  if should_open_browser:
    webbrowser.open(authorization_url)
  try:
    while server.query is None and time.time() < deadline:
      server.handle_request()
  finally:
    server.server_close()

  if server.query is None:
    raise XgtSecurityError(
        'OIDC authorization code login timed out before redirect')
  if 'error' in server.query:
    error_value = server.query['error'][0]
    description = server.query.get('error_description', [''])[0]
    raise XgtConnectionError(
        f'OIDC authorization failed: {error_value}: {description}')
  # State parameter is a CSRF token; mismatch means the redirect did not
  # originate from the authorization request we sent.
  if server.query.get('state', [''])[0] != state:
    raise XgtSecurityError('OIDC authorization state mismatch')
  code = server.query.get('code', [''])[0]
  if not code:
    raise XgtConnectionError('OIDC redirect is missing authorization code')

  request_payload = {
    'grant_type': 'authorization_code',
    'client_id': config.client_id,
    'client_secret': (config.client_secret
                      if config.client_secret is not None else ''),
    'code': code,
    'redirect_uri': redirect_uri,
  }
  if config.pkce:
    request_payload['code_verifier'] = code_verifier
  try:
    return _http_post_form(endpoints['token_endpoint'], request_payload,
                           config=config)
  except OidcHttpError as ex:
    raise XgtConnectionError(
        'OIDC authorization-code token exchange failed. '
        f'client_id={config.client_id!r}, redirect_uri={redirect_uri!r}, '
        f'error={ex.oidc_error!r}, '
        f'description={ex.oidc_description!r}') from ex

# ---------------------------------------------------------------------------
# Interactive flow selector
# ---------------------------------------------------------------------------

def _interactive_oidc_token_payload(config: OidcAuth,
                                    open_browser: bool = True
                                    ) -> dict[str, Any]:
  if config.flow == 'auth_code':
    return _authorization_code_flow(config, open_browser=open_browser)
  if config.flow == 'device':
    return _login_oidc_token_payload(config, open_browser=open_browser)
  if config.flow != 'auto':
    raise XgtConnectionError(
        "OidcAuth.flow must be 'auto', 'auth_code', or 'device'")

  endpoints = _discover_interactive_endpoints(config)
  if (_prefer_device_flow_for_environment() and
      endpoints.get('device_authorization_endpoint')):
    device_config = replace(
        config,
        token_endpoint=endpoints['token_endpoint'],
        device_authorization_endpoint=endpoints['device_authorization_endpoint'])
    return _login_oidc_token_payload(device_config, open_browser=open_browser)
  if endpoints.get('authorization_endpoint'):
    auth_code_config = replace(
        config,
        token_endpoint=endpoints['token_endpoint'],
        authorization_endpoint=endpoints['authorization_endpoint'])
    return _authorization_code_flow(auth_code_config,
                                    open_browser=open_browser)
  if endpoints.get('device_authorization_endpoint'):
    device_config = replace(
        config,
        token_endpoint=endpoints['token_endpoint'],
        device_authorization_endpoint=endpoints['device_authorization_endpoint'])
    return _login_oidc_token_payload(device_config, open_browser=open_browser)
  raise XgtConnectionError(
      'OIDC discovery document supports neither authorization_code nor '
      'device flow')

# ---------------------------------------------------------------------------
# Client credentials flow
# ---------------------------------------------------------------------------

def _get_token_endpoint(config: OidcClientConfig) -> str:
  """Resolve the token endpoint for the client credentials flow.

  Only the token endpoint is needed here — device_authorization_endpoint is
  not required, so this is lighter than _discover_oidc_endpoints.
  """
  if config.token_endpoint is not None:
    return config.token_endpoint
  if not config.issuer:
    raise XgtConnectionError('OIDC issuer is not configured')
  metadata = _discover_metadata(config)
  token_endpoint = metadata.get('token_endpoint')
  if not token_endpoint:
    raise XgtConnectionError('OIDC discovery document is missing token_endpoint')
  return token_endpoint

def _client_credentials_flow(config: OidcCredentialsAuth) -> dict[str, Any]:
  """Exchange client_id + client_secret for an access token directly."""
  if not config.client_id:
    raise XgtConnectionError('OIDC client_id is not configured')
  if not config.client_secret:
    raise XgtConnectionError('OidcCredentialsAuth requires a client_secret')
  token_endpoint = _get_token_endpoint(config)
  request_payload = {
    'grant_type': 'client_credentials',
    'client_id': config.client_id,
    'client_secret': config.client_secret,
    'scope': _scope_string(config),
  }
  if config.audience is not None:
    request_payload['audience'] = config.audience
  return _http_post_form(token_endpoint, request_payload, config=config)
