"""
JWT Token Creation and Validation

Uses PyJWT directly - no allauth dependency.
Tokens are tied to Django sessions for immediate revocation on logout.
"""

import time
from typing import NamedTuple

import jwt
from django.contrib.sessions.backends.base import SessionBase

from .settings import get_settings


class TokenPair(NamedTuple):
    """Access and refresh token pair."""
    access_token: str
    refresh_token: str
    expires_in: int


class TokenPayload(NamedTuple):
    """Decoded token payload."""
    user_id: int | str
    session_key: str
    token_type: str
    is_staff: bool
    is_superuser: bool
    exp: int
    iat: int


class JWTUser:
    """
    Minimal user object created from JWT claims.

    Used as request.user for JWT-authenticated requests.
    No database query required - all data comes from the token.

    If you need the full User object with all fields, query explicitly:
        user = User.objects.get(pk=request.user.id)
    """

    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  # Assumed active if they have a valid token

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

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


def create_access_token(
    user_id: int | str,
    session_key: str,
    *,
    is_staff: bool = False,
    is_superuser: bool = False,
) -> str:
    """
    Create a short-lived access token.

    The token contains:
    - sub: user ID
    - sid: session key (for revocation checking)
    - staff: is_staff flag
    - super: is_superuser flag
    - type: "access"
    - iat: issued at
    - exp: expiration
    """
    settings = get_settings()
    now = int(time.time())

    payload = {
        "sub": str(user_id),
        "sid": session_key,
        "staff": is_staff,
        "super": is_superuser,
        "type": "access",
        "iat": now,
        "exp": now + settings.access_token_expires_in,
    }

    return jwt.encode(
        payload,
        settings.private_key,
        algorithm=settings.algorithm,
    )


def create_refresh_token(
    user_id: int | str,
    session_key: str,
    *,
    is_staff: bool = False,
    is_superuser: bool = False,
) -> str:
    """
    Create a longer-lived refresh token.

    The token contains:
    - sub: user ID
    - sid: session key (for revocation checking)
    - staff: is_staff flag
    - super: is_superuser flag
    - type: "refresh"
    - iat: issued at
    - exp: expiration
    """
    settings = get_settings()
    now = int(time.time())

    payload = {
        "sub": str(user_id),
        "sid": session_key,
        "staff": is_staff,
        "super": is_superuser,
        "type": "refresh",
        "iat": now,
        "exp": now + settings.refresh_token_expires_in,
    }

    return jwt.encode(
        payload,
        settings.private_key,
        algorithm=settings.algorithm,
    )


def create_token_pair(
    user_id: int | str,
    session_key: str,
    *,
    is_staff: bool = False,
    is_superuser: bool = False,
) -> TokenPair:
    """Create both access and refresh tokens."""
    settings = get_settings()
    return TokenPair(
        access_token=create_access_token(
            user_id, session_key, is_staff=is_staff, is_superuser=is_superuser
        ),
        refresh_token=create_refresh_token(
            user_id, session_key, is_staff=is_staff, is_superuser=is_superuser
        ),
        expires_in=settings.access_token_expires_in,
    )


def decode_token(token: str, expected_type: str = None) -> TokenPayload | None:
    """
    Decode and validate a JWT token.

    Returns None if:
    - Token is invalid or expired
    - Token type doesn't match expected_type (if specified)
    """
    settings = get_settings()

    try:
        payload = jwt.decode(
            token,
            settings.public_key,
            algorithms=[settings.algorithm],
        )
    except jwt.PyJWTError:
        return None

    # Validate token type if specified
    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 validate_session(session_key: str) -> bool:
    """
    Check if a session is still valid (exists and not expired).

    This is the key to immediate logout revocation - if the session
    is destroyed, tokens tied to it become invalid.
    """
    from importlib import import_module

    from django.conf import settings as django_settings

    jwt_settings = get_settings()

    if not jwt_settings.validate_session:
        return True

    # Use the configured session engine
    engine = import_module(django_settings.SESSION_ENGINE)
    SessionStore = engine.SessionStore

    # Try to load the session
    session = SessionStore(session_key=session_key)

    # Check if session exists and is not empty
    # exists() is more reliable than checking load() result
    return session.exists(session_key)


def refresh_tokens(refresh_token: str) -> TokenPair | None:
    """
    Use a refresh token to obtain new tokens.

    Returns None if:
    - Refresh token is invalid or expired
    - Associated session no longer exists
    """
    payload = decode_token(refresh_token, expected_type="refresh")

    if payload is None:
        return None

    # Validate the session still exists
    if not validate_session(payload.session_key):
        return None

    # Issue new token pair with same claims
    return create_token_pair(
        payload.user_id,
        payload.session_key,
        is_staff=payload.is_staff,
        is_superuser=payload.is_superuser,
    )
