"""
enfonica_credentials.py — Enfonica service account authentication (reference implementation)

Exchanges an Enfonica service account key for an OAuth2 bearer access token
using the JWT Bearer Grant (RFC 7523), and caches the token in memory.
Copy this single file into your project.

Requires: Python 3.9+ and the `cryptography` package (pip install cryptography).
HTTP uses the standard library only.

Creating credentials (pick one):

    creds = EnfonicaCredentials.from_env()                         # path in ENFONICA_APPLICATION_CREDENTIALS
    creds = EnfonicaCredentials.from_service_account_file("key.json")
    creds = EnfonicaCredentials.from_service_account_info(json_str_or_dict)   # e.g. from a secret manager

Using them on each API request:

    headers = {"Authorization": creds.authorization_header()}     # "Bearer <token>"
    token = creds.get_access_token()                               # just the token

Notes:
  - Create ONE instance and reuse it for the life of your application. Tokens
    are cached in memory and refreshed automatically after 30 minutes (the
    server issues 60-minute tokens), so a new instance per request means a
    token request per API call.
  - Instances are thread-safe.
  - A 401 from the API is a real error (bad key, revoked key, missing
    permission). Do not retry it by fetching a new token.
  - Tokens are never written to disk.
"""

from __future__ import annotations

import base64
import json
import os
import re
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, Union

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding

SCOPE = "https://api.enfonica.com/auth/platform"
GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
ENV_VAR = "ENFONICA_APPLICATION_CREDENTIALS"

JWT_LIFETIME_SECONDS = 3600
TOKEN_REFRESH_SECONDS = 30 * 60
HTTP_TIMEOUT_SECONDS = 30

_REQUIRED_FIELDS = ("type", "project_id", "private_key_id", "private_key", "client_email", "token_uri")


class EnfonicaCredentialsError(Exception):
    """Raised when a key file is missing or invalid, or a token request fails."""


class EnfonicaCredentials:
    """Service account credentials that produce cached Enfonica access tokens."""

    @classmethod
    def from_env(cls) -> "EnfonicaCredentials":
        """Load the key file named by the ENFONICA_APPLICATION_CREDENTIALS environment variable."""
        path = os.environ.get(ENV_VAR)
        if not path:
            raise EnfonicaCredentialsError(
                f"{ENV_VAR} is not set. Set it to the path of a service account key file, "
                "or use from_service_account_file() / from_service_account_info()."
            )
        return cls.from_service_account_file(path)

    @classmethod
    def from_service_account_file(cls, path: str) -> "EnfonicaCredentials":
        """Load a service account key file from disk."""
        try:
            with open(path, encoding="utf-8") as f:
                return cls.from_service_account_info(f.read())
        except OSError as e:
            raise EnfonicaCredentialsError(f"Could not read key file '{path}': {e}") from e

    @classmethod
    def from_service_account_info(cls, info: Union[str, dict]) -> "EnfonicaCredentials":
        """Load a service account key from its JSON text or an already-parsed dict."""
        if isinstance(info, str):
            try:
                info = json.loads(info)
            except json.JSONDecodeError as e:
                raise EnfonicaCredentialsError(f"Invalid key file: not valid JSON: {e}") from e
        return cls(info)

    def __init__(self, info: dict[str, Any]) -> None:
        if not isinstance(info, dict):
            raise EnfonicaCredentialsError(f"Invalid key file: expected a JSON object, got {type(info).__name__}")
        if info.get("type") != "service_account":
            raise EnfonicaCredentialsError(
                f"Invalid key file: expected type 'service_account', got {info.get('type')!r}"
            )
        missing = [f for f in _REQUIRED_FIELDS if not isinstance(info.get(f), str) or not info[f]]
        if missing:
            raise EnfonicaCredentialsError(f"Invalid key file: missing field(s): {', '.join(missing)}")

        self.project_id: str = info["project_id"]
        self.client_email: str = info["client_email"]
        self.token_uri: str = info["token_uri"]  # always read from the key file, never hardcoded
        self._private_key_id: str = info["private_key_id"]
        self._private_key = _load_private_key(info["private_key"])

        self._lock = threading.Lock()
        self._token: str | None = None
        self._fetched_at = 0.0

    def get_access_token(self) -> str:
        """Return a valid access token, fetching a new one if none is cached or it is 30+ minutes old."""
        with self._lock:
            if self._token is None or time.monotonic() - self._fetched_at >= TOKEN_REFRESH_SECONDS:
                self._token = self._fetch_token()
                self._fetched_at = time.monotonic()
            return self._token

    def authorization_header(self) -> str:
        """Return the value for the HTTP Authorization header: 'Bearer <token>'."""
        return f"Bearer {self.get_access_token()}"

    def _fetch_token(self) -> str:
        body = urllib.parse.urlencode({"grant_type": GRANT_TYPE, "assertion": self._signed_assertion()}).encode()
        request = urllib.request.Request(
            self.token_uri,
            data=body,
            method="POST",
            headers={"Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json"},
        )
        try:
            with urllib.request.urlopen(request, timeout=HTTP_TIMEOUT_SECONDS) as response:
                text = response.read().decode("utf-8")
        except urllib.error.HTTPError as e:
            detail = e.read().decode("utf-8", errors="replace")
            raise EnfonicaCredentialsError(
                f"Token request failed: HTTP {e.code} from {self.token_uri}: {detail}"
            ) from e
        except urllib.error.URLError as e:
            raise EnfonicaCredentialsError(f"Token request to {self.token_uri} failed: {e.reason}") from e

        try:
            token = json.loads(text)["access_token"]
        except (json.JSONDecodeError, KeyError, TypeError) as e:
            raise EnfonicaCredentialsError(f"Token response was malformed: {text!r}") from e
        if not isinstance(token, str) or not token:
            raise EnfonicaCredentialsError(f"Token response had an empty access_token: {text!r}")
        return token

    def _signed_assertion(self) -> str:
        """Build the RS256-signed JWT assertion (RFC 7523 §3)."""
        now = int(time.time())
        header = {"alg": "RS256", "typ": "JWT", "kid": self._private_key_id}
        payload = {
            "iss": self.client_email,
            "aud": self.token_uri,
            "scope": SCOPE,
            "iat": now,
            "exp": now + JWT_LIFETIME_SECONDS,
        }
        signing_input = f"{_b64url_json(header)}.{_b64url_json(payload)}".encode("ascii")
        signature = self._private_key.sign(signing_input, padding.PKCS1v15(), hashes.SHA256())
        return f"{signing_input.decode('ascii')}.{_b64url(signature)}"


def _load_private_key(pem: str):
    # Decode the PEM body ourselves so keys wrapped at any line width (or not
    # at all) load correctly; some PEM parsers insist on 64-character lines.
    body = re.sub(r"-----(BEGIN|END) [A-Z ]+-----|\s+", "", pem)
    try:
        return serialization.load_der_private_key(base64.b64decode(body), password=None)
    except ValueError as e:
        raise EnfonicaCredentialsError(f"Invalid key file: private_key is not a valid PKCS#8 RSA key: {e}") from e


def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _b64url_json(obj: dict) -> str:
    return _b64url(json.dumps(obj, separators=(",", ":")).encode("utf-8"))