"""
JWT detection, validation and payload normalization for DTA and Identity tokens.

Supports:
- DTA JWT (current): HS256, issuer dta.totvs.ai (e.g. SAML flow).
- Identity JWT (new): RS256 with kid, validated via JWKS.
"""
import logging
from typing import Any, Dict, Literal, Optional

import jwt
import requests

from .constants import DTA_ISSUER
from .exceptions import AuthError

logger = logging.getLogger(__name__)

JWTType = Literal["dta", "identity"]

DEFAULT_JWKS_BASE_URL = "https://api-fluig.totvs.app/accounts/api/v1/jwks"

IDENTITY_JWT_AUDIENCE = "fluig_authenticator_resource"


def _strip_bearer(token: str) -> str:
    """Remove 'Bearer ' prefix if present."""
    if token.startswith("Bearer "):
        return token[7:].strip()
    return token.strip()


def _decode_jwt_or_raise(
    token: str,
    key: Any,
    algorithms: list,
    log_label: str = "JWT",
    **kwargs: Any,
) -> Dict[str, Any]:
    """Decode JWT and map PyJWT exceptions to AuthError."""
    try:
        return jwt.decode(token, key, algorithms=algorithms, **kwargs)
    except jwt.ExpiredSignatureError as e:
        logger.warning("%s expired", log_label)
        raise AuthError("Token expired", status_code=401) from e
    except jwt.InvalidTokenError as e:
        logger.warning("Invalid %s: %s", log_label, e)
        raise AuthError("Invalid token", status_code=401) from e


def detect_jwt_type(token: str) -> JWTType:
    """
    Detect JWT type from token header only (no signature validation or network).

    Args:
        token: Raw JWT string (optionally with "Bearer " prefix).

    Returns:
        'dta' for SAML/HS256, 'identity' for OIDC/RS256 with kid.

    Raises:
        AuthError: If header is missing, invalid, or algorithm unsupported.
    """
    token = _strip_bearer(token)
    try:
        header = jwt.get_unverified_header(token)
    except Exception as e:
        logger.debug("Failed to decode JWT header: %s", e)
        raise AuthError("Invalid JWT token structure", status_code=401) from e

    algorithm = header.get("alg")
    if algorithm == "HS256":
        return "dta"
    if algorithm == "RS256" and "kid" in header:
        return "identity"

    logger.warning("Unsupported JWT algorithm: %s", algorithm)
    raise AuthError(f"Unsupported JWT algorithm: {algorithm}", status_code=401)


def validate_dta_jwt(
    token: str,
    secret: str,
    issuer: str = DTA_ISSUER,
) -> Dict[str, Any]:
    """
    Validate a DTA-issued JWT (HS256).

    Args:
        token: Raw JWT string.
        secret: HMAC secret (e.g. JWT_SECRET_KEY).
        issuer: Expected issuer claim (default dta.totvs.ai).

    Returns:
        Decoded payload dict.

    Raises:
        AuthError: If token is invalid, expired, or issuer mismatch.
    """
    token = _strip_bearer(token)
    if not secret:
        raise AuthError("DTA JWT secret not configured", status_code=500)

    payload = _decode_jwt_or_raise(
        token, secret, ["HS256"], log_label="DTA JWT"
    )
    if payload.get("iss") != issuer:
        logger.warning("Invalid issuer for DTA JWT: %s", payload.get("iss"))
        raise AuthError("Invalid issuer", status_code=401)

    return payload


def get_jwks_public_key(
    kid: str,
    jwks_base_url: str,
    jwks_cache: Optional[Dict[str, Any]] = None,
    timeout: int = 5,
) -> Any:
    """
    Fetch and return the public key for the given kid from JWKS.

    Args:
        kid: Key ID from JWT header.
        jwks_base_url: Base URL for JWKS (e.g. .../jwks).
        jwks_cache: Optional dict to cache keys (mutated in place).
        timeout: Request timeout in seconds.

    Returns:
        Public key object for use with PyJWT decode.

    Raises:
        AuthError: On network error, missing key, or parse error.
    """
    cache = jwks_cache if jwks_cache is not None else {}
    if kid in cache:
        return cache[kid]

    url = f"{jwks_base_url.rstrip('/')}/{kid}"
    try:
        response = requests.get(url, timeout=timeout)
        response.raise_for_status()
        jwks = response.json()
    except requests.exceptions.RequestException as e:
        logger.error("JWKS request failed %s: %s", url, e)
        raise AuthError(
            "Identity Provider Connection Failed",
            status_code=500,
        ) from e

    keys = jwks.get("keys", [])
    if not keys:
        raise AuthError("No keys found in JWKS response", status_code=401)

    key_data = keys[0]
    if key_data.get("kid") != kid:
        logger.warning("JWKS response KID mismatch, expected %s", kid)
        raise AuthError("Public key not found", status_code=401)

    try:
        public_key = jwt.algorithms.RSAAlgorithm.from_jwk(key_data)
    except Exception as e:
        logger.error("Error parsing JWKS key: %s", e)
        raise AuthError("Token validation error", status_code=401) from e

    cache[kid] = public_key
    return public_key


def validate_identity_jwt(
    token: str,
    jwks_base_url: str = DEFAULT_JWKS_BASE_URL,
    jwks_cache: Optional[Dict[str, Any]] = None,
    audience: str = IDENTITY_JWT_AUDIENCE,
    timeout: int = 5,
) -> Dict[str, Any]:
    """
    Validate an Identity/OIDC JWT (RS256) using JWKS.

    Args:
        token: Raw JWT string.
        jwks_base_url: Base URL for JWKS endpoint.
        jwks_cache: Optional dict to cache public keys.
        audience: Expected audience claim.
        timeout: JWKS request timeout.

    Returns:
        Decoded payload dict.

    Raises:
        AuthError: If token is invalid or expired.
    """
    token = _strip_bearer(token)
    try:
        header = jwt.get_unverified_header(token)
        unverified = jwt.decode(
            token,
            options={"verify_signature": False, "verify_exp": False},
        )
    except Exception as e:
        logger.warning("Identity JWT header/decode error: %s", e)
        raise AuthError("Invalid JWT token structure", status_code=401) from e

    kid = header.get("kid")
    if not kid:
        raise AuthError("Kid not found in header", status_code=401)

    issuer = unverified.get("iss")
    public_key = get_jwks_public_key(
        kid,
        jwks_base_url,
        jwks_cache=jwks_cache,
        timeout=timeout,
    )

    payload = _decode_jwt_or_raise(
        token,
        public_key,
        ["RS256"],
        log_label="Identity JWT",
        audience=audience,
    )
    if issuer and "fluig.io" not in str(issuer):
        logger.warning("Suspicious Identity issuer: %s", issuer)

    return payload


def normalize_user_data(
    payload: Dict[str, Any],
    jwt_type: JWTType,
) -> Dict[str, Any]:
    """
    Normalize raw JWT payload to a common user-data shape for DTA services.

    Args:
        payload: Decoded JWT payload.
        jwt_type: 'dta' or 'identity'.

    Returns:
        Dict with keys: name, identifier, email, organization, super,
        org_admin, tenants, roles, exp, jwt_type, and optionally
        generated_at (dta) or authorities, apps, company_id, tenant_code (identity).
    """
    if jwt_type == "dta":
        return {
            "name": payload.get("name"),
            "identifier": payload.get("identifier"),
            "email": payload.get("identifier"),
            "organization": payload.get("organization"),
            "super": payload.get("super", False),
            "org_admin": payload.get("org_admin", False),
            "tenants": payload.get("tenants", []),
            "roles": payload.get("roles", []),
            "exp": payload.get("exp"),
            "jwt_type": "dta",
            "generated_at": payload.get("generated_at"),
        }

    if jwt_type == "identity":
        authorities = payload.get("authorities", [])
        return {
            "name": payload.get("fullName"),
            "identifier": payload.get("email"),
            "email": payload.get("email"),
            "organization": "totvs",
            "super": "admin" in authorities,
            "org_admin": "admin" in authorities,
            "tenants": [],
            "roles": [],
            "exp": payload.get("exp"),
            "jwt_type": "identity",
            "authorities": authorities,
            "apps": payload.get("apps", []),
            "company_id": payload.get("companyId"),
            "tenant_code": payload.get("tenantCode"),
        }

    raise ValueError(f"Unsupported JWT type: {jwt_type}")


def validate_and_normalize(
    token: str,
    dta_secret: str,
    jwks_base_url: str = DEFAULT_JWKS_BASE_URL,
    jwks_cache: Optional[Dict[str, Any]] = None,
    dta_issuer: str = DTA_ISSUER,
    identity_audience: str = IDENTITY_JWT_AUDIENCE,
) -> Dict[str, Any]:
    """
    Detect JWT type, validate, and return normalized user data.

    Convenience function for services that have both DTA and Identity JWTs.

    Args:
        token: Raw JWT (optionally with "Bearer " prefix).
        dta_secret: Secret for DTA HS256 tokens.
        jwks_base_url: JWKS URL for Identity tokens.
        jwks_cache: Optional cache for JWKS keys.
        dta_issuer: Expected issuer for DTA tokens.
        identity_audience: Expected audience for Identity tokens.

    Returns:
        Normalized user data dict.

    Raises:
        AuthError: If token is invalid or unsupported.
    """
    token = _strip_bearer(token)
    jwt_type = detect_jwt_type(token)

    if jwt_type == "dta":
        payload = validate_dta_jwt(token, dta_secret, issuer=dta_issuer)
    elif jwt_type == "identity":
        payload = validate_identity_jwt(
            token,
            jwks_base_url=jwks_base_url,
            jwks_cache=jwks_cache,
            audience=identity_audience,
        )
    else:
        raise AuthError("Unsupported JWT type", status_code=401)

    return normalize_user_data(payload, jwt_type)
