"""Azure AD (Entra ID) authentication provider."""

import logging
from typing import Dict, Any, Optional
import jwt
from jwt import PyJWKClient
import requests
from .models import AuthenticationError

LOGGER = logging.getLogger(__name__)


class AzureADAuthProvider:
    """Authenticate requests using Azure AD JWT tokens."""
    
    def __init__(
        self,
        tenant_id: str,
        client_id: str,
        audience: str,
    ):
        """Initialize Azure AD auth provider.
        
        Args:
            tenant_id: Azure AD tenant ID
            client_id: Application (client) ID from app registration
            audience: Expected audience claim (usually api://your-app-id)
        """
        self.tenant_id = tenant_id
        self.client_id = client_id
        self.audience = audience
        
        # JWKS endpoint for token signature verification
        self.jwks_uri = (
            f"https://login.microsoftonline.com/{tenant_id}/discovery/v2.0/keys"
        )
        self.jwk_client: Optional[PyJWKClient] = None
        
        LOGGER.info(
            "Initialized Azure AD auth provider",
            extra={"context": {"tenant_id": tenant_id, "client_id": client_id}},
        )
    
    async def authenticate(self, token: str) -> Dict[str, Any]:
        """Validate JWT token against Azure AD.
        
        Args:
            token: JWT access token from Azure AD
            
        Returns:
            User context dictionary
            
        Raises:
            AuthenticationError: If token is invalid, expired, or has wrong audience
        """
        try:
            # Lazy initialize JWK client
            if not self.jwk_client:
                self.jwk_client = PyJWKClient(self.jwks_uri)
            
            # Get signing key from token header
            signing_key = self.jwk_client.get_signing_key_from_jwt(token)
            
            # Decode and validate token
            payload = jwt.decode(
                token,
                signing_key.key,
                algorithms=["RS256"],
                audience=self.audience,
                issuer=f"https://sts.windows.net/{self.tenant_id}/",
            )
            
            # Extract user context
            user_context = {
                "id": payload.get("oid") or payload.get("sub"),
                "email": payload.get("upn") or payload.get("email"),
                "name": payload.get("name"),
                "roles": payload.get("roles", []),
                "scopes": (payload.get("scp") or "").split() if "scp" in payload else [],
                "tenant_id": payload.get("tid"),
                "source": "azure_ad",
            }
            
            LOGGER.debug(
                "Successfully authenticated user",
                extra={"context": {"user_id": user_context["id"]}},
            )
            
            return user_context
        
        except jwt.ExpiredSignatureError:
            raise AuthenticationError("Token has expired")
        except jwt.InvalidAudienceError:
            raise AuthenticationError(f"Invalid audience. Expected: {self.audience}")
        except jwt.InvalidIssuerError:
            raise AuthenticationError(f"Invalid issuer. Expected tenant: {self.tenant_id}")
        except jwt.DecodeError as exc:
            raise AuthenticationError(f"Failed to decode token: {exc}")
        except Exception as exc:
            LOGGER.error(
                "Authentication failed",
                extra={"context": {"error": str(exc)}},
            )
            raise AuthenticationError(f"Authentication failed: {exc}")
