"""
Custom JWT Token Generation for ThrillWiki

This module provides custom JWT token generation that includes authentication
method claims for enhanced MFA satisfaction logic.

Claims added:
- auth_method: How the user authenticated (password, passkey, totp, google, discord)
- mfa_verified: Whether MFA was verified during this login
- provider_mfa: Whether the OAuth provider (Discord) has MFA enabled
"""

from typing import Literal, TypedDict

from rest_framework_simplejwt.tokens import RefreshToken

# Type definitions for auth methods
AuthMethod = Literal["password", "passkey", "totp", "google", "discord"]


class TokenClaims(TypedDict, total=False):
    """Type definition for custom JWT claims."""

    auth_method: AuthMethod
    mfa_verified: bool
    provider_mfa: bool


def create_tokens_for_user(
    user,
    auth_method: AuthMethod = "password",
    mfa_verified: bool = False,
    provider_mfa: bool = False,
) -> dict[str, str]:
    """
    Generate JWT tokens with custom authentication claims.

    Args:
        user: The Django user object
        auth_method: How the user authenticated
        mfa_verified: True if MFA (TOTP/passkey) was verified at login
        provider_mfa: True if OAuth provider (Discord) has MFA enabled

    Returns:
        Dictionary with 'access' and 'refresh' token strings
    """
    refresh = RefreshToken.for_user(user)

    # Add custom claims to both refresh and access tokens
    refresh["auth_method"] = auth_method
    refresh["mfa_verified"] = mfa_verified
    refresh["provider_mfa"] = provider_mfa

    access = refresh.access_token

    return {
        "access": str(access),
        "refresh": str(refresh),
    }


def get_auth_method_for_provider(provider: str) -> AuthMethod:
    """
    Map OAuth provider name to AuthMethod type.

    Args:
        provider: The provider name (e.g., 'google', 'discord')

    Returns:
        The corresponding AuthMethod
    """
    provider_map: dict[str, AuthMethod] = {
        "google": "google",
        "discord": "discord",
    }
    return provider_map.get(provider, "password")


def get_provider_mfa_status(provider: str, extra_data: dict) -> bool:
    """
    Extract MFA status from OAuth provider extra_data.

    Only Discord exposes mfa_enabled. Google does not share this info.

    Args:
        provider: The OAuth provider name
        extra_data: The extra_data dict from SocialAccount

    Returns:
        True if provider has MFA enabled, False otherwise
    """
    if provider == "discord":
        return extra_data.get("mfa_enabled", False)

    # Google and other providers don't expose MFA status
    return False