"""Sign-in with an OpenID Connect provider (Authelia, docs/accounts.md).

An extra way in next to username + password: an account is linked once in the settings
while signed in normally, after that the provider's login leads into it. Nobody gets a new
account from the provider — an unlinked identity is turned away.

Authorization code flow with PKCE (S256) and a confidential client (client_secret_basic).
The id_token is checked here (RS256 against the provider's JWKS, iss, aud, exp, nonce)
with `cryptography`, so no OIDC library is needed.

The book key is not touched: it is derived from the diora password on the device. A device
that signs in this way gets it from an unlocked device or by typing the password once
(accounts/keyshare.py, static/js/app/33-unlock.js).
"""
import base64
import hashlib
import json
import secrets
import time

from urllib.parse import urlencode

import requests
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from django.conf import settings

TIMEOUT = 10
_CACHE_SECONDS = 3600
_cache = {}  # url -> (fetched_at, json)


class OidcError(Exception):
    pass


def enabled() -> bool:
    return bool(settings.OIDC_ISSUER and settings.OIDC_CLIENT_ID and settings.OIDC_CLIENT_SECRET)


def _b64url_decode(s: str) -> bytes:
    return base64.urlsafe_b64decode(s + '=' * (-len(s) % 4))


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


def _get_json(url: str, fresh: bool = False) -> dict:
    hit = _cache.get(url)
    if hit and not fresh and time.time() - hit[0] < _CACHE_SECONDS:
        return hit[1]
    try:
        resp = requests.get(url, timeout=TIMEOUT)
        resp.raise_for_status()
        data = resp.json()
    except (requests.RequestException, ValueError) as exc:
        raise OidcError(f'{url} nicht erreichbar: {exc}') from exc
    _cache[url] = (time.time(), data)
    return data


def discovery() -> dict:
    issuer = settings.OIDC_ISSUER.rstrip('/')
    return _get_json(issuer + '/.well-known/openid-configuration')


def new_flow() -> dict:
    """Random state, nonce and PKCE verifier for one sign-in (kept in the session)."""
    return {
        'state': secrets.token_urlsafe(24),
        'nonce': secrets.token_urlsafe(24),
        'verifier': secrets.token_urlsafe(48),
    }


def authorize_url(flow: dict, redirect_uri: str) -> str:
    challenge = _b64url(hashlib.sha256(flow['verifier'].encode('ascii')).digest())
    params = {
        'response_type': 'code',
        'client_id': settings.OIDC_CLIENT_ID,
        'redirect_uri': redirect_uri,
        'scope': 'openid profile email groups',
        'state': flow['state'],
        'nonce': flow['nonce'],
        'code_challenge': challenge,
        'code_challenge_method': 'S256',
    }
    return discovery()['authorization_endpoint'] + '?' + urlencode(params)


def exchange_code(code: str, flow: dict, redirect_uri: str) -> dict:
    """Trade the code for tokens and return the checked id_token claims."""
    try:
        resp = requests.post(
            discovery()['token_endpoint'],
            data={
                'grant_type': 'authorization_code',
                'code': code,
                'redirect_uri': redirect_uri,
                'code_verifier': flow['verifier'],
            },
            auth=(settings.OIDC_CLIENT_ID, settings.OIDC_CLIENT_SECRET),
            timeout=TIMEOUT,
        )
    except requests.RequestException as exc:
        raise OidcError(f'Token-Abruf fehlgeschlagen: {exc}') from exc
    if resp.status_code != 200:
        raise OidcError(f'Token-Abruf abgelehnt ({resp.status_code})')
    try:
        id_token = resp.json()['id_token']
    except (ValueError, KeyError) as exc:
        raise OidcError('Antwort ohne id_token') from exc
    return verify_id_token(id_token, flow['nonce'])


def _public_key(kid: str):
    jwks_uri = discovery()['jwks_uri']
    for fresh in (False, True):  # a key rotation shows up as an unknown kid: fetch again once
        for jwk in _get_json(jwks_uri, fresh=fresh).get('keys', []):
            if jwk.get('kty') == 'RSA' and (not kid or jwk.get('kid') == kid):
                n = int.from_bytes(_b64url_decode(jwk['n']), 'big')
                e = int.from_bytes(_b64url_decode(jwk['e']), 'big')
                return rsa.RSAPublicNumbers(e, n).public_key()
    raise OidcError('Signaturschlüssel des Anbieters nicht gefunden')


def verify_id_token(token: str, nonce: str, now: float = None) -> dict:
    try:
        header_b64, payload_b64, sig_b64 = token.split('.')
        header = json.loads(_b64url_decode(header_b64))
        claims = json.loads(_b64url_decode(payload_b64))
        signature = _b64url_decode(sig_b64)
    except (ValueError, TypeError) as exc:
        raise OidcError('id_token unlesbar') from exc
    if header.get('alg') != 'RS256':
        raise OidcError(f'nicht unterstützte Signatur {header.get("alg")!r} (erwartet RS256)')
    try:
        _public_key(header.get('kid', '')).verify(
            signature, f'{header_b64}.{payload_b64}'.encode('ascii'),
            padding.PKCS1v15(), hashes.SHA256())
    except InvalidSignature as exc:
        raise OidcError('Signatur des id_token ungültig') from exc

    now = time.time() if now is None else now
    if claims.get('iss', '').rstrip('/') != settings.OIDC_ISSUER.rstrip('/'):
        raise OidcError('falscher Aussteller')
    aud = claims.get('aud')
    if settings.OIDC_CLIENT_ID not in (aud if isinstance(aud, list) else [aud]):
        raise OidcError('id_token ist nicht für diora')
    if not isinstance(claims.get('exp'), (int, float)) or claims['exp'] < now - 60:
        raise OidcError('id_token abgelaufen')
    if not nonce or claims.get('nonce') != nonce:
        raise OidcError('nonce passt nicht')
    if not claims.get('sub'):
        raise OidcError('id_token ohne sub')
    return claims