content: {-# LANGUAGE FlexibleContexts #-}
  
  -- | Module implementing various OAuth2 flow types and their request/response handling.
  -- Provides support for:
  --
  --   * Authorization Code Grant
  --   * Device Authorization Grant
  --   * PKCE Extension
  --   * Token Refresh
  --   * User Info Endpoints
  module Network.OAuth2.Experiment.Flows where
  
  import Control.Concurrent
  import Control.Monad.IO.Class (MonadIO (..))
  import Control.Monad.Trans.Except
  import Data.Aeson (FromJSON)
  import Data.Bifunctor
  import Data.ByteString.Lazy.Char8 qualified as BSL
  import Data.Map.Strict qualified as Map
  import Data.Maybe
  import Network.HTTP.Client.Contrib
  import Network.HTTP.Conduit
  import Network.OAuth2
  import Network.OAuth2 qualified as OAuth2
  import Network.OAuth2.Experiment.Flows.DeviceAuthorizationRequest
  import Network.OAuth2.Experiment.Flows.RefreshTokenRequest
  import Network.OAuth2.Experiment.Flows.TokenRequest
  import Network.OAuth2.Experiment.Flows.UserInfoRequest
  import Network.OAuth2.Experiment.Grants.AuthorizationCode
  import Network.OAuth2.Experiment.Grants.DeviceAuthorization
  import Network.OAuth2.Experiment.Pkce
  import Network.OAuth2.Experiment.Types
  import Network.OAuth2.Experiment.Utils
  import URI.ByteString hiding (UserInfo)
  
  -------------------------------------------------------------------------------
  --                           Authorization Requests                          --
  -------------------------------------------------------------------------------
  
  -- | Constructs an Authorization Code request URI according to
  -- <https://www.rfc-editor.org/rfc/rfc6749#section-4.1.1 RFC 6749 Section 4.1.1>.
  --
  -- The generated URI includes:
  --   * client_id
  --   * response_type (always "code")
  --   * redirect_uri
  --   * state (if provided)
  --   * scope (if provided)
  mkAuthorizationRequest :: IdpApplication i AuthorizationCodeApplication -> URI
  mkAuthorizationRequest idpApp =
    let req = mkAuthorizationRequestParam (application idpApp)
        allParams =
          map (bimap tlToBS tlToBS) $
            Map.toList $
              toQueryParam req
     in appendQueryParams allParams $
          idpAuthorizeEndpoint (idp idpApp)
  
  -- | Constructs an Authorization Code request URI with PKCE support according to
  -- <https://datatracker.ietf.org/doc/html/rfc7636 RFC 7636>.
  --
  -- Returns both the authorization URI and the generated code verifier.
  -- The code verifier must be stored securely for later use in the token request.
  mkPkceAuthorizeRequest ::
    MonadIO m =>
    IdpApplication i AuthorizationCodeApplication ->
    m (URI, CodeVerifier)
  mkPkceAuthorizeRequest IdpApplication {..} = do
    (req, codeVerifier) <- mkPkceAuthorizeRequestParam application
    let allParams = map (bimap tlToBS tlToBS) $ Map.toList $ toQueryParam req
    let url =
          appendQueryParams allParams $
            idpAuthorizeEndpoint idp
    pure (url, codeVerifier)
  
  -------------------------------------------------------------------------------
  --                                Device Auth                                --
  -------------------------------------------------------------------------------
  
  -- | Makes Device Authorization Request
  -- https://www.rfc-editor.org/rfc/rfc8628#section-3.1
  conduitDeviceAuthorizationRequest ::
    MonadIO m =>
    IdpApplication i DeviceAuthorizationApplication ->
    Manager ->
    ExceptT BSL.ByteString m DeviceAuthorizationResponse
  conduitDeviceAuthorizationRequest IdpApplication {..} mgr = do
    case idpDeviceAuthorizationEndpoint idp of
      Nothing -> throwE "[conduitDeviceAuthorizationRequest] Device Authorization Flow is not supported: missing device_authorization_endpoint."
      Just deviceAuthEndpoint -> do
        let deviceAuthReq = mkDeviceAuthorizationRequestParam application
            body = unionMapsToQueryParams [toQueryParam deviceAuthReq]
        ExceptT . liftIO $ do
          req <- addDefaultRequestHeaders <$> uriToRequest deviceAuthEndpoint
          let req' =
                if daAuthorizationRequestAuthenticationMethod application == ClientSecretBasic
                  then addSecretToHeader (daClientId application) (daClientSecret application) req
                  else req
          resp <- httpLbs (urlEncodedBody body req') mgr
          pure $ first ("[conduitDeviceAuthorizationRequest] " <>) $ handleResponseJSON resp
  
  -- | Polls for a token using the device authorization flow.
  --
  -- This implements the polling mechanism described in
  -- <https://www.rfc-editor.org/rfc/rfc8628#section-3.5 RFC 8628 Section 3.5>.
  -- Handles automatic retries and interval adjustments based on IdP responses.
  pollDeviceTokenRequest ::
    MonadIO m =>
    IdpApplication i DeviceAuthorizationApplication ->
    Manager ->
    DeviceAuthorizationResponse ->
    ExceptT TokenResponseError m TokenResponse
  pollDeviceTokenRequest idpApp mgr deviceAuthResp = do
    pollDeviceTokenRequestInternal
      idpApp
      mgr
      (deviceCode deviceAuthResp)
      (fromMaybe 5 $ interval deviceAuthResp)
  
  pollDeviceTokenRequestInternal ::
    MonadIO m =>
    IdpApplication i DeviceAuthorizationApplication ->
    Manager ->
    DeviceCode ->
    Int ->
    -- | Polling Interval
    ExceptT TokenResponseError m TokenResponse
  pollDeviceTokenRequestInternal idpApp mgr deviceCode intervalSeconds = do
    resp <- runExceptT (conduitTokenRequest idpApp mgr deviceCode)
    case resp of
      Left trRespError -> do
        case tokenResponseError trRespError of
          -- TODO: Didn't have a good idea to expand the error code
          -- specifically for device token request flow
          -- Device Token Response additional error code: https://www.rfc-editor.org/rfc/rfc8628#section-3.5
          UnknownErrorCode "authorization_pending" -> do
            liftIO $ threadDelay $ intervalSeconds * 1000000
            pollDeviceTokenRequestInternal idpApp mgr deviceCode intervalSeconds
          UnknownErrorCode "slow_down" -> do
            let newIntervalSeconds = intervalSeconds + 5
            liftIO $ threadDelay $ newIntervalSeconds * 1000000
            pollDeviceTokenRequestInternal idpApp mgr deviceCode newIntervalSeconds
          _ -> throwE trRespError
      Right v -> pure v
  
  -------------------------------------------------------------------------------
  --                               Token Request                               --
  -------------------------------------------------------------------------------
  
  -- | Sends a token request according to
  -- <https://www.rfc-editor.org/rfc/rfc6749#section-4.1.3 RFC 6749 Section 4.1.3>.
  --
  -- This is used for exchanging authorization codes, device codes, or other
  -- grant types for access tokens.
  conduitTokenRequest ::
    (HasTokenRequest a, ToQueryParam (TokenRequest a), MonadIO m) =>
    IdpApplication i a ->
    Manager ->
    ExchangeTokenInfo a ->
    ExceptT TokenResponseError m TokenResponse
  conduitTokenRequest idpApp mgr exchangeToken = do
    let req = mkTokenRequestParam (application idpApp) exchangeToken
        body =
          unionMapsToQueryParams
            [ toQueryParam req
            ]
     in conduitTokenRequestInternal idpApp mgr body
  
  -------------------------------------------------------------------------------
  --                             PKCE Token Request                            --
  -------------------------------------------------------------------------------
  
  -- | https://datatracker.ietf.org/doc/html/rfc7636#section-4.5
  conduitPkceTokenRequest ::
    (HasTokenRequest a, ToQueryParam (TokenRequest a), MonadIO m) =>
    IdpApplication i a ->
    Manager ->
    (ExchangeTokenInfo a, CodeVerifier) ->
    ExceptT TokenResponseError m TokenResponse
  conduitPkceTokenRequest idpApp mgr (exchangeToken, codeVerifier) =
    let req = mkTokenRequestParam (application idpApp) exchangeToken
        body =
          unionMapsToQueryParams
            [ toQueryParam req
            , toQueryParam codeVerifier
            ]
     in conduitTokenRequestInternal idpApp mgr body
  
  -------------------------------------------------------------------------------
  --                              Refresh Token                               --
  -------------------------------------------------------------------------------
  
  -- | Makes a Refresh Token Request according to
  -- <https://www.rfc-editor.org/rfc/rfc6749#section-6 RFC 6749 Section 6>.
  --
  -- Used to obtain a new access token using a refresh token.
  conduitRefreshTokenRequest ::
    (MonadIO m, HasRefreshTokenRequest a) =>
    IdpApplication i a ->
    Manager ->
    OAuth2.RefreshToken ->
    ExceptT TokenResponseError m TokenResponse
  conduitRefreshTokenRequest ia mgr rt =
    let tokenReq = mkRefreshTokenRequestParam (application ia) rt
        body = unionMapsToQueryParams [toQueryParam tokenReq]
     in conduitTokenRequestInternal ia mgr body
  
  -------------------------------------------------------------------------------
  --                                 User Info                                 --
  -------------------------------------------------------------------------------
  
  -- | Makes a standard request to the userinfo endpoint using GET method.
  --
  -- This is commonly used with OpenID Connect providers to fetch
  -- user profile information using an access token.
  conduitUserInfoRequest ::
    (MonadIO m, HasUserInfoRequest a, FromJSON b) =>
    IdpApplication i a ->
    Manager ->
    AccessToken ->
    ExceptT BSL.ByteString m b
  conduitUserInfoRequest = conduitUserInfoRequestWithCustomMethod authGetJSON
  
  -- | Makes a request to the userinfo endpoint using a custom HTTP method.
  --
  -- Some IdPs may require different HTTP methods (instead of GET) or custom headers
  -- for fetching user information. This function provides that flexibility.
  conduitUserInfoRequestWithCustomMethod ::
    (MonadIO m, HasUserInfoRequest a, FromJSON b) =>
    ( Manager ->
      AccessToken ->
      URI ->
      ExceptT BSL.ByteString m b
    ) ->
    IdpApplication i a ->
    Manager ->
    AccessToken ->
    ExceptT BSL.ByteString m b
  conduitUserInfoRequestWithCustomMethod fetchMethod IdpApplication {..} mgr at =
    fetchMethod mgr at (idpUserInfoEndpoint idp)
  
  -------------------------------------------------------------------------------
  --                              Internal helpers                             --
  -------------------------------------------------------------------------------
  
  conduitTokenRequestInternal ::
    ( MonadIO m
    , HasClientAuthenticationMethod a
    , FromJSON b
    ) =>
    IdpApplication i a ->
    -- | HTTP connection manager.
    Manager ->
    -- | Request body.
    PostBody ->
    -- | Response as ByteString
    ExceptT TokenResponseError m b
  conduitTokenRequestInternal IdpApplication {..} manager body = do
    let clientAuthMethod = getClientAuthenticationMethod application
        url = idpTokenEndpoint idp
        updateAuthHeader =
          case clientAuthMethod of
            ClientSecretBasic -> addClientAuthToHeader application
            ClientSecretPost -> id
            ClientAssertionJwt -> id
        go = do
          req <- uriToRequest url
          let req' = (updateAuthHeader . addDefaultRequestHeaders) req
          httpLbs (urlEncodedBody body req') manager
    resp <- ExceptT . liftIO $ fmap handleOAuth2TokenResponse go
    case parseResponseFlexible resp of
      Right obj -> return obj
      Left e -> throwE e
  