{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeOperators #-}

module Servant.Auth.Firebase where

import Control.Lens hiding ((.=))
import Control.Monad
import Control.Monad.Except
import Data.Proxy
import Data.Word8 (isSpace, toLower)
import Servant.API
import qualified Servant.Swagger as Servant
import qualified Crypto.JOSE.JWK as JWK
import qualified Network.URI as URI
import qualified Crypto.JWT as JWT
import qualified Data.Aeson as Aeson
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.HashMap.Strict as HMS
import qualified Data.PEM as PEM
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.X509 as X509
import qualified Network.HTTP.Client as HTTP
import qualified Network.HTTP.Simple as HTTP
import qualified Network.Wai as Wai
import qualified Servant as Servant
import qualified Servant.Client.Core as Servant
import qualified Servant.Client.Core as Servant.Client
import qualified Servant.Server.Internal as Servant

data Protected

newtype ProjectId = ProjectId { unFirebaseProjectId :: T.Text }

data FirebaseLoginSettings = FirebaseLoginSettings
  { firebaseLoginProjectId :: ProjectId
  , firebaseLoginGetKeys :: IO (HMS.HashMap T.Text T.Text)
  }

defaultFirebaseLoginSettings
  :: HTTP.Manager -> ProjectId -> FirebaseLoginSettings
defaultFirebaseLoginSettings mgr pid = FirebaseLoginSettings
    { firebaseLoginProjectId = pid
    , firebaseLoginGetKeys = HTTP.getResponseBody <$> HTTP.httpJSON req
    }
  where
    -- TODO: proper error handling here
    req =
      HTTP.setRequestSecure True .
      HTTP.setRequestPort 443 .
      HTTP.setRequestHost "www.googleapis.com" .
      HTTP.setRequestPath "/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com" .
      HTTP.setRequestManager mgr $
      HTTP.defaultRequest {
          HTTP.responseTimeout = HTTP.responseTimeoutMicro (500 * 1000)
        }

newtype UserId = UserId { unUserId :: T.Text }
  deriving newtype (Aeson.FromJSON, Aeson.ToJSON, Show, Eq)

newtype UnverifiedJWT = UnverifiedJWT JWT.SignedJWT

-- TODO: MAKE SURE PATTERN MATCH FAILURES AREN'T PROPAGATED TO CLIENT!!!
verifyUser :: FirebaseLoginSettings -> UnverifiedJWT -> IO UserId
verifyUser settings (UnverifiedJWT jwt) = do
  jwkmap <- firebaseLoginGetKeys settings
  let projectId = unFirebaseProjectId $ firebaseLoginProjectId settings

  t <- case jwt ^.. JWT.signatures . JWT.header . JWT.kid of
    [Just (JWT.HeaderParam () t)] -> pure t
    xs -> error $ "Expected exactly one signature with 'kid', got: " <> show xs

  jwkct <- case HMS.lookup t jwkmap of
    Nothing -> error $ "Could not find key " <> show t <> " in response"
    Just ct -> pure ct

  -- TODO: get rid of 'error'
  pem <- case PEM.pemParseBS (T.encodeUtf8 jwkct) of
    Left e -> error $ show e
    Right [e] -> pure e
    Right xs -> error $ show xs

  cert <- case X509.decodeSignedCertificate (PEM.pemContent pem) of
    Left e -> error $ show e
    Right c -> pure c

  jwk <- runExceptT (JWK.fromX509Certificate cert) >>= \case
    Left (e :: JWT.JWTError) -> error $ show e
    Right jwk -> pure jwk

  issUri <- case URI.parseURI $ "https://securetoken.google.com/" <> T.unpack projectId of
    Just issUri -> pure issUri
    Nothing -> error $ "Could not use project ID in URI"

  let config =
        JWT.defaultJWTValidationSettings
          (\sou -> Just projectId == sou ^? JWT.string) & -- aud
          JWT.issuerPredicate .~ (\sou -> Just issUri == sou ^? JWT.uri) -- iss
  runExceptT (JWT.verifyClaims config jwk jwt) >>= \case
    Right cs -> do
      case cs ^. JWT.claimSub of
        Nothing -> error "Could not get a subject from claim set"
        Just sou -> case sou ^? JWT.string of
          Nothing -> error "Expected subject to be string"
          Just u -> pure (UserId u)
    Left (e :: JWT.JWTError) -> error (show e)

instance FromHttpApiData UnverifiedJWT where
  parseUrlPiece = const $ Left "No support for JWT"
  parseHeader bs = case JWT.decodeCompact (BL.fromStrict bs) of
    Left (e :: JWT.Error) -> Left $ T.pack $ show e
    Right jwt -> Right $ UnverifiedJWT jwt

instance
    ( Servant.HasClient m sub
    , Servant.RunClient m ) => Servant.HasClient m (Protected :> sub) where
  -- TODO: something better than just Text
  type Client m (Protected :> sub) = T.Text -> Servant.Client m sub
  clientWithRoute p1 Proxy req = \bs ->
    Servant.clientWithRoute
      p1 (Proxy :: Proxy sub)
      (Servant.Client.addHeader "Authorization" ("Bearer " <> bs) req)
  hoistClientMonad p1 Proxy hoist c = \bs ->
    Servant.Client.hoistClientMonad p1 (Proxy :: Proxy sub) hoist (c bs)

-- | Find and decode an 'Authorization' header from the request as JWT
decodeJWTHdr :: Wai.Request -> Either String UnverifiedJWT
decodeJWTHdr req = do
    ah <- case lookup "Authorization" (Wai.requestHeaders req) of
      Just x -> Right x
      Nothing -> Left "No authorization header"
    let (b, rest) = BS.break isSpace ah
    guard (BS.map toLower b == "bearer")
    tok <- case snd <$> BS.uncons rest of
      Nothing -> Left "No token"
      Just x -> Right x
    case JWT.decodeCompact (BL.fromStrict tok) of
      Left (e :: JWT.Error) -> Left $ show e <> ": " <> show rest
      Right jwt -> Right (UnverifiedJWT jwt)

runJWTAuth :: FirebaseLoginSettings -> Wai.Request -> Servant.DelayedIO UserId
runJWTAuth settings req = case decodeJWTHdr req of
    Left e -> error $ "bad auth: " <> e -- TODO: delayedFailFatal
    Right ujwt -> liftIO $ verifyUser settings ujwt

instance
    ( Servant.HasContextEntry context FirebaseLoginSettings
    , Servant.HasServer sub context
    ) => Servant.HasServer (Protected :> sub) context where
  type ServerT (Protected :> sub) m = UserId -> Servant.ServerT sub m

  route Proxy c subserver = do
      Servant.route (Proxy :: Proxy sub)
        c (subserver `Servant.addAuthCheck` authCheck)
    where
      authCheck :: Servant.DelayedIO UserId
      authCheck = Servant.withRequest $ runJWTAuth (Servant.getContextEntry c)

  hoistServerWithContext Proxy p hoist s = \uid ->
    Servant.hoistServerWithContext (Proxy :: Proxy sub) p hoist (s uid)

instance Servant.HasSwagger sub => Servant.HasSwagger (Protected :> sub) where
  toSwagger Proxy = Servant.toSwagger (Proxy :: Proxy sub)
