"""Auth0 JWKS-based JWT validation for access tokens.

This module provides the shared JWT helpers used by both the auth middleware
(tenant isolation) and the FastAPI dependency layer (``get_current_user``).
Keeping them here avoids a circular import between ``middleware.py`` and
``deps.py``.
"""

from __future__ import annotations

import asyncio
import logging
from functools import partial

import jwt as pyjwt
from jwt import PyJWKClient

from ..settings import Settings

logger = logging.getLogger(__name__)

# Module-level JWKS client cache (per domain).
_jwks_clients: dict[str, PyJWKClient] = {}


def _get_jwks_client(domain: str) -> PyJWKClient:
    """Get or create a cached PyJWKClient for the Auth0 domain."""
    if domain not in _jwks_clients:
        jwks_url = f"https://{domain}/.well-known/jwks.json"
        _jwks_clients[domain] = PyJWKClient(jwks_url, cache_keys=True)
    return _jwks_clients[domain]


async def validate_access_token(token: str, settings: Settings) -> dict:
    """Validate an Auth0 access token and return its claims.

    Raises ``Exception`` on any validation failure.
    """
    if not settings.auth0_domain or not settings.auth0_audience:
        raise ValueError("Auth0 domain and audience must be configured for JWT validation")

    jwks_client = _get_jwks_client(settings.auth0_domain)
    # JWKS fetch is synchronous (HTTP call on cache miss) — run in executor
    loop = asyncio.get_running_loop()
    signing_key = await loop.run_in_executor(
        None, partial(jwks_client.get_signing_key_from_jwt, token)
    )

    claims = pyjwt.decode(
        token,
        signing_key.key,
        algorithms=["RS256"],
        audience=settings.auth0_audience,
        issuer=f"https://{settings.auth0_domain}/",
    )

    return claims


async def resolve_org_login(claims: dict, registry: object | None) -> str:
    """Resolve org_login from JWT claims and the installation registry.

    Strategies (in order):
    1. ``org_id`` claim → registry lookup (Auth0 Organizations flow)
    2. Single-org fallback → if the registry has exactly one active
       installation, use that org.  Safe because Auth0 login is already
       gated to authorized org members via the ``restrict_signups`` Action.

    Returns the org_login string, or ``""`` if resolution fails.
    """
    if not registry:
        return ""

    # Strategy 1: org_id claim (web auth with Auth0 Organizations)
    org_id = claims.get("org_id", "")
    if org_id:
        try:
            installation = await registry.get_installation_by_auth0_org(org_id)
            if installation:
                return installation.org_login
        except Exception:
            logger.debug("Failed to resolve org_id %s", org_id, exc_info=True)

    # Strategy 2: single-org fallback (device auth without org_id)
    try:
        orgs = await registry.list_orgs()
        if len(orgs) == 1:
            return orgs[0]
    except Exception:
        logger.debug("Failed to list orgs for single-org fallback", exc_info=True)

    return ""


async def resolve_jwt_org(token: str, settings: Settings, registry) -> str | None:
    """Validate a JWT and resolve its org_id to an org_login.

    Returns the ``org_login`` string, or ``None`` if the token is invalid
    or the org cannot be resolved.  This is the shared helper used by both
    the auth middleware and the dependency layer.
    """
    try:
        claims = await validate_access_token(token, settings)
    except Exception:
        logger.debug("JWT validation failed in resolve_jwt_org", exc_info=True)
        return None

    org_login = await resolve_org_login(claims, registry)
    return org_login or None
