"""
JWT access/refresh tokens — adapter-agnostic (PyJWT).

Config is injected (`JWTConfig`) rather than read from any framework's settings.
`validate_session` (the immediate-logout-revocation check) is Django-session-bound
and stays in the Django adapter; `refresh_tokens` takes a `session_validator`
callable so the core stays framework-free.
"""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Callable, NamedTuple

import jwt


@dataclass(frozen=True)
class JWTConfig:
    private_key: str
    public_key: str
    algorithm: str = "HS256"
    access_token_expires_in: int = 300
    refresh_token_expires_in: int = 604800


class TokenPair(NamedTuple):
    access_token: str
    refresh_token: str
    expires_in: int


class TokenPayload(NamedTuple):
    user_id: int | str
    session_key: str
    token_type: str
    is_staff: bool
    is_superuser: bool
    exp: int
    iat: int


class JWTUser:
    """Minimal `Identity` built from JWT claims — no DB query."""

    def __init__(self, payload: TokenPayload):
        self.id = int(payload.user_id) if isinstance(payload.user_id, str) else payload.user_id
        self.pk = self.id
        self.is_staff = payload.is_staff
        self.is_superuser = payload.is_superuser
        self.is_authenticated = True
        self.is_anonymous = False
        self.is_active = True

    def __str__(self) -> str:
        return f"JWTUser(id={self.id})"

    def __repr__(self) -> str:
        return f"JWTUser(id={self.id}, is_staff={self.is_staff}, is_superuser={self.is_superuser})"


def _mint(user_id: int | str, session_key: str, token_type: str, ttl: int,
          config: JWTConfig, is_staff: bool, is_superuser: bool) -> str:
    now = int(time.time())
    payload = {
        "sub": str(user_id),
        "sid": session_key,
        "staff": is_staff,
        "super": is_superuser,
        "type": token_type,
        "iat": now,
        "exp": now + ttl,
    }
    return jwt.encode(payload, config.private_key, algorithm=config.algorithm)


def create_access_token(user_id, session_key, config: JWTConfig, *,
                        is_staff: bool = False, is_superuser: bool = False) -> str:
    return _mint(user_id, session_key, "access", config.access_token_expires_in,
                 config, is_staff, is_superuser)


def create_refresh_token(user_id, session_key, config: JWTConfig, *,
                         is_staff: bool = False, is_superuser: bool = False) -> str:
    return _mint(user_id, session_key, "refresh", config.refresh_token_expires_in,
                 config, is_staff, is_superuser)


def create_token_pair(user_id, session_key, config: JWTConfig, *,
                      is_staff: bool = False, is_superuser: bool = False) -> TokenPair:
    return TokenPair(
        access_token=create_access_token(user_id, session_key, config,
                                         is_staff=is_staff, is_superuser=is_superuser),
        refresh_token=create_refresh_token(user_id, session_key, config,
                                           is_staff=is_staff, is_superuser=is_superuser),
        expires_in=config.access_token_expires_in,
    )


def decode_token(token: str, config: JWTConfig, expected_type: str | None = None) -> TokenPayload | None:
    """Decode + validate. None on invalid/expired token, or type mismatch."""
    try:
        payload = jwt.decode(token, config.public_key, algorithms=[config.algorithm])
    except jwt.PyJWTError:
        return None
    if expected_type and payload.get("type") != expected_type:
        return None
    return TokenPayload(
        user_id=payload["sub"],
        session_key=payload["sid"],
        token_type=payload["type"],
        is_staff=payload.get("staff", False),
        is_superuser=payload.get("super", False),
        exp=payload["exp"],
        iat=payload["iat"],
    )


def refresh_tokens(
    refresh_token: str,
    config: JWTConfig,
    session_validator: Callable[[str], bool] | None = None,
) -> TokenPair | None:
    """Exchange a refresh token for a new pair. None if invalid or the session is gone.

    `session_validator(session_key) -> bool` lets the Django adapter enforce
    immediate-logout revocation; omit it (or pass a always-True) where there is
    no session store.
    """
    payload = decode_token(refresh_token, config, expected_type="refresh")
    if payload is None:
        return None
    if session_validator is not None and not session_validator(payload.session_key):
        return None
    return create_token_pair(payload.user_id, payload.session_key, config,
                             is_staff=payload.is_staff, is_superuser=payload.is_superuser)
