"""
Custom Microsoft SSO class that allows overriding default Microsoft endpoints.

This module provides a subclass of fastapi_sso's MicrosoftSSO that allows
custom authorization, token, and userinfo endpoints to be specified via environment
variables.

Environment Variables:
- MICROSOFT_AUTHORIZATION_ENDPOINT: Custom authorization endpoint URL
- MICROSOFT_TOKEN_ENDPOINT: Custom token endpoint URL
- MICROSOFT_USERINFO_ENDPOINT: Custom userinfo endpoint URL

If these are not set, the default Microsoft endpoints are used.
"""

import os
from typing import Final

import pydantic
from fastapi_sso.sso.base import DiscoveryDocument
from fastapi_sso.sso.microsoft import MicrosoftSSO

from litellm._logging import verbose_proxy_logger


class CustomMicrosoftSSO(MicrosoftSSO):
    """
    Microsoft SSO subclass that allows overriding default endpoints via environment variables.

    Supports:
    - MICROSOFT_AUTHORIZATION_ENDPOINT
    - MICROSOFT_TOKEN_ENDPOINT
    - MICROSOFT_USERINFO_ENDPOINT
    """

    def __init__(
        self,
        client_id: str,
        client_secret: str,
        redirect_uri: pydantic.AnyHttpUrl | str | None = None,
        allow_insecure_http: bool = False,
        scope: list[str] | None = None,
        tenant: str | None = None,
    ):
        super().__init__(
            client_id=client_id,
            client_secret=client_secret,
            redirect_uri=redirect_uri,
            allow_insecure_http=allow_insecure_http,
            scope=scope,
            tenant=tenant,
        )

    async def get_discovery_document(self) -> DiscoveryDocument:
        """
        Override to support custom endpoints via environment variables.
        Falls back to default Microsoft endpoints if not set.
        """
        custom_authorization_endpoint: Final = os.getenv("MICROSOFT_AUTHORIZATION_ENDPOINT", None)
        custom_token_endpoint: Final = os.getenv("MICROSOFT_TOKEN_ENDPOINT", None)
        custom_userinfo_endpoint: Final = os.getenv("MICROSOFT_USERINFO_ENDPOINT", None)

        # Use custom endpoints if set, otherwise use defaults
        authorization_endpoint: Final = (
            custom_authorization_endpoint or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/authorize"
        )
        token_endpoint = custom_token_endpoint or f"https://login.microsoftonline.com/{self.tenant}/oauth2/v2.0/token"
        userinfo_endpoint: Final = custom_userinfo_endpoint or f"https://graph.microsoft.com/{self.version}/me"

        if custom_authorization_endpoint or custom_token_endpoint or custom_userinfo_endpoint:
            verbose_proxy_logger.debug(
                "Using custom Microsoft SSO endpoints - authorization: %s, token: %s, userinfo: %s",
                authorization_endpoint,
                token_endpoint,
                userinfo_endpoint,
            )

        return DiscoveryDocument(
            authorization_endpoint=authorization_endpoint,
            token_endpoint=token_endpoint,
            userinfo_endpoint=userinfo_endpoint,
        )