from __future__ import annotations

import base64
import binascii
import hashlib
import secrets
import time
import uuid
from dataclasses import dataclass
from typing import Any

from django.conf import settings
from django.contrib.auth import get_user_model
from django.contrib.auth.base_user import AbstractBaseUser
from django.contrib.sessions.backends.base import SessionBase
from django.utils.functional import SimpleLazyObject

import jwt
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

from allauth.account.internal.userkit import str_to_user_id, user_id_to_str
from allauth.core.internal import jwkkit
from allauth.core.internal.sessionkit import get_session_user
from allauth.headless import app_settings
from allauth.headless.internal.sessionkit import lookup_session


@dataclass
class JWTConfig:
    signing_key: Any
    verifying_key: Any
    algorithm: str
    jwk_dict: dict[str, Any] | None = None


def validate_access_token(token: str) -> tuple[Any, dict[str, Any]] | None:
    payload = decode_token(token, "access")
    if payload is None:
        return None
    if app_settings.JWT_STATEFUL_VALIDATION_ENABLED:
        session = get_token_session(payload)
        if session is None:
            return None
    sub = payload["sub"]
    pk = str_to_user_id(sub)
    lazy_user = SimpleLazyObject(lambda: get_user_model().objects.get(pk=pk))
    return lazy_user, payload


def get_session_key_cipher(initialization_vector: bytes) -> Cipher:
    secret_key = settings.SECRET_KEY
    key = hashlib.sha256(secret_key.encode()).digest()
    algorithm = algorithms.AES(key)
    mode = modes.CTR(initialization_vector)
    cipher = Cipher(algorithm, mode)
    return cipher


def session_key_to_sid(session_key: str) -> str:
    """
    In case an access token leaks, we do not want the session encoded in
    that token to be put to use as an X-Session-Token or session
    cookie. Therefore, we encrypt the sesison key. Unauthenticated symmetric
    encryption is fine, as the JWT is signed.
    """
    initialization_vector = secrets.token_bytes(16)
    cipher = get_session_key_cipher(initialization_vector)
    encryptor = cipher.encryptor()
    encrypted_message = encryptor.update(session_key.encode()) + encryptor.finalize()
    sid = base64.b64encode(initialization_vector + encrypted_message).decode()
    return sid


def session_key_from_sid(sid: str) -> str | None:
    try:
        encrypted_data = base64.b64decode(sid)
    except binascii.Error:
        return None
    initialization_vector = encrypted_data[:16]
    encrypted_message = encrypted_data[16:]
    if len(encrypted_message) == 0:
        return None
    cipher = get_session_key_cipher(initialization_vector)
    decryptor = cipher.decryptor()
    message_decrypted = decryptor.update(encrypted_message) + decryptor.finalize()
    try:
        return message_decrypted.decode()
    except UnicodeDecodeError:
        return None


def validate_token_user(
    token: dict[str, Any], session: SessionBase
) -> AbstractBaseUser | None:
    user = get_session_user(session)
    if user is None:
        return None
    sub = user_id_to_str(user)
    if sub != token["sub"]:
        return None
    return user


def validate_refresh_token(
    token: str,
) -> tuple[Any, SessionBase, dict[str, Any]] | None:
    payload = decode_token(token, "refresh")
    if payload is None:
        return None
    jti = payload["jti"]
    session = get_token_session(payload)
    if session is None:
        return None
    refresh_token_jti_to_exp = get_refresh_token_state(session)
    exp = refresh_token_jti_to_exp.get(jti)
    now = time.time()
    if exp is None or exp <= now:
        return None
    user = validate_token_user(payload, session)
    return user, session, payload


def get_token_session(payload: dict[str, Any]) -> SessionBase | None:
    sid = payload.get("sid")
    if not isinstance(sid, str):
        return None
    session_key = session_key_from_sid(sid)
    if not session_key:
        return None
    session = lookup_session(session_key)
    if session is None:
        return None
    return session


def _get_jwt_config() -> JWTConfig:
    algorithm = app_settings.JWT_ALGORITHM
    jwk_dict = None
    if algorithm.startswith("HS"):
        key = app_settings.JWT_PRIVATE_KEY or settings.SECRET_KEY
        signing_key = key
        verifying_key = key
    elif algorithm.startswith("RS"):
        jwk_dict, signing_key = jwkkit.load_jwk_from_pem(app_settings.JWT_PRIVATE_KEY)
        verifying_key = signing_key.public_key()
    else:
        raise ValueError(f"Unsupported JWT algorithm: {algorithm}")
    return JWTConfig(
        signing_key=signing_key,
        verifying_key=verifying_key,
        algorithm=algorithm,
        jwk_dict=jwk_dict,
    )


def get_jwt_headers(config: JWTConfig) -> dict[str, Any]:
    headers = {}
    if config.jwk_dict:
        kid = config.jwk_dict.get("kid")
        if kid:
            headers["kid"] = kid
    return headers


def decode_token(token: str, use: str) -> dict[str, Any] | None:
    config = _get_jwt_config()
    try:
        payload = jwt.decode(
            token,
            key=config.verifying_key,
            algorithms=[config.algorithm],
            options={
                "verify_signature": True,
                "verify_iss": False,
                "verify_aud": False,
                "verify_exp": True,
            },
        )
    except jwt.PyJWTError:
        return None
    else:
        if payload.get("token_use") != use:
            return None
        for key in ("jti", "sub", "sid"):
            value = payload.get(key)
            if not isinstance(value, str):
                return None
        return payload


def create_token(
    token_use: str,
    *,
    sub: str,
    sid: str,
    claims: dict[str, Any] | None = None,
    expires_in: int,
) -> tuple[str, dict[str, Any]]:
    config = _get_jwt_config()
    now = int(time.time())
    payload = {}
    if claims is not None:
        payload.update(claims)
    payload.update(
        {
            "iat": now,
            "exp": now + expires_in,
            "sid": sid,
            "jti": str(uuid.uuid4()),
            "token_use": token_use,
            "sub": sub,
        }
    )
    headers = get_jwt_headers(config)
    return (
        jwt.encode(
            payload,
            config.signing_key,
            algorithm=config.algorithm,
            headers=headers,
        ),
        payload,
    )


def get_refresh_token_state(session: SessionBase) -> dict[str, int]:
    return session.setdefault("headless_refresh_tokens", {})


def create_refresh_token(user: AbstractBaseUser, session: SessionBase) -> str:
    assert user.is_authenticated  # nosec
    assert session.session_key  # nosec
    sub = user_id_to_str(user)
    sid = session_key_to_sid(session.session_key)
    token, payload = create_token(
        "refresh",
        sub=sub,
        sid=sid,
        expires_in=app_settings.JWT_REFRESH_TOKEN_EXPIRES_IN,
    )
    refresh_token_jti_to_exp = get_refresh_token_state(session)
    refresh_token_jti_to_exp[payload["jti"]] = payload["exp"]
    return token


def create_access_token(
    user: AbstractBaseUser, session: SessionBase, claims: dict[str, Any]
) -> str:
    assert user.is_authenticated  # nosec
    assert session.session_key  # nosec
    sid = session_key_to_sid(session.session_key)
    sub = user_id_to_str(user)
    return create_token(
        "access",
        sub=sub,
        sid=sid,
        claims=claims,
        expires_in=app_settings.JWT_ACCESS_TOKEN_EXPIRES_IN,
    )[0]


def invalidate_refresh_token(session: SessionBase, token: dict[str, Any]) -> None:
    refresh_token_jti_to_exp = get_refresh_token_state(session)
    jti = token["jti"]
    refresh_token_jti_to_exp.pop(jti, None)
