module Utils.Tokens
  ( bearerToken
  , encodeBearer, BearerTokenException(..), decodeBearer
  , bearerParseJSON'
  , askBearer
  , formEmbedBearerPost, formEmbedBearerGet
  ) where

import Import.NoModel

import Utils.Lens

import Model
import Model.Tokens

import Jose.Jwk (JwkSet(..))
import qualified Jose.Jwt as Jose

import Data.Aeson.Types (Parser)
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Parser as JSON
import qualified Data.Aeson.Parser.Internal as JSON (jsonEOF')
import qualified Data.Aeson.Internal as JSON (iparse, formatError)

import qualified Data.HashMap.Strict as HashMap

import Settings

import CryptoID

import Text.Blaze (Markup)


bearerParseJSON' :: forall m.
                  ( Hashable (AuthId (HandlerSite m)), Eq (AuthId (HandlerSite m))
                  , HasCryptoUUID (AuthId (HandlerSite m)) (ReaderT CryptoIDKey Parser)
                  , ParseRoute (HandlerSite m)
                  , Hashable (Route (HandlerSite m))
                  , MonadHandler m
                  , MonadCrypto m
                  , MonadCryptoKey m ~ CryptoIDKey
                  )
               => m (Value -> Parser (BearerToken (HandlerSite m)))
-- ^ Read `CryptoIDKey` from monadic context and construct a `Parser` for `BearerToken`s
bearerParseJSON' = do
  cidKey <- cryptoIDKey return
  return $ flip runReaderT cidKey . bearerParseJSON


bearerToken :: forall m.
               ( MonadHandler m
               , HasInstanceID (HandlerSite m) InstanceId
               , HasClusterID (HandlerSite m) ClusterId
               , HasAppSettings (HandlerSite m)
               )
            => HashSet (Either Value (AuthId (HandlerSite m)))
            -> Maybe (HashSet (Route (HandlerSite m)))
            -> Maybe AuthDNF
            -> Maybe (Maybe UTCTime) -- ^ @Nothing@ determines default expiry time automatically
            -> Maybe UTCTime -- ^ @Nothing@ means token starts to be valid immediately
            -> m (BearerToken (HandlerSite m))
-- ^ Smart constructor for `BearerToken`, does not set route restrictions (due to polymorphism), use `tokenRestrict`
bearerToken bearerAuthority bearerRoutes bearerAddAuth mBearerExpiresAt bearerStartsAt = do
  bearerIdentifier <- liftIO getRandom
  bearerIssuedAt <- liftIO getCurrentTime
  bearerIssuedBy <- getsYesod $ view instanceID
  bearerIssuedFor <- getsYesod $ view clusterID

  defaultExpiration <- getsYesod $ view _appBearerExpiration

  let bearerExpiresAt
        | Just t <- mBearerExpiresAt
        = t
        | Just tDiff <- defaultExpiration
        = Just $ tDiff `addUTCTime` fromMaybe bearerIssuedAt bearerStartsAt
        | otherwise
        = Nothing
      bearerRestrictions = HashMap.empty

  return BearerToken{..}


encodeBearer :: forall m.
               ( MonadHandler m
               , HasJSONWebKeySet (HandlerSite m) JwkSet
               , HasAppSettings (HandlerSite m)
               , HasCryptoUUID (AuthId (HandlerSite m)) m
               , RenderRoute (HandlerSite m)
               )
            => BearerToken (HandlerSite m) -> m Jwt
-- ^ Call `bearerToJSON` and encode the result as a `Jwt` according to `appJwtEncoding`
encodeBearer token = do
  payload <- Jose.Claims . toStrict . JSON.encode <$> bearerToJSON token
  JwkSet jwks <- getsYesod $ view jsonWebKeySet
  jwtEncoding <- getsYesod $ view _appBearerEncoding
  either throwM return =<< liftIO (Jose.encode jwks jwtEncoding payload)


data BearerTokenException
  = BearerTokenJwtError Jose.JwtError -- ^ An Error occurred in the underlying `Jwt`-Implementation
  | BearerTokenUnsecured -- ^ `Jwt` is insufficiently secured (unsigned and not encrypted)
  | BearerTokenInvalidFormat String -- ^ Content of the `Jwt` could not be parsed as a `BearerToken`
  | BearerTokenExpired | BearerTokenNotStarted
  deriving (Eq, Show, Generic, Typeable)

instance Exception BearerTokenException

decodeBearer :: forall m.
               ( MonadHandler m
               , HasJSONWebKeySet (HandlerSite m) JwkSet
               , Hashable (AuthId (HandlerSite m)), Eq (AuthId (HandlerSite m))
               , HasCryptoUUID (AuthId (HandlerSite m)) (ReaderT CryptoIDKey Parser)
               , MonadCryptoKey m ~ CryptoIDKey
               , MonadCrypto m
               , ParseRoute (HandlerSite m)
               , Hashable (Route (HandlerSite m))
               )
            => Jwt -> m (BearerToken (HandlerSite m))
-- ^ Decode a `Jwt` and call `bearerParseJSON`
--
-- Throws `bearerTokenException`s
decodeBearer (Jwt bs) = do
  JwkSet jwks <- getsYesod $ view jsonWebKeySet
  content <- either (throwM . BearerTokenJwtError) return =<< liftIO (Jose.decode jwks Nothing bs)
  content' <- case content of
    Jose.Unsecured _ -> throwM BearerTokenUnsecured
    Jose.Jws (_header, payload) -> return payload
    Jose.Jwe (_header, payload) -> return payload
  parser <- bearerParseJSON'
  bearer@BearerToken{..} <- either (throwM . BearerTokenInvalidFormat . uncurry JSON.formatError) return $ JSON.eitherDecodeStrictWith JSON.jsonEOF' (JSON.iparse parser) content'
  now <- liftIO getCurrentTime
  unless (NTop bearerExpiresAt > NTop (Just now)) $
    throwM BearerTokenExpired
  unless (bearerStartsAt <= Just now) $
    throwM BearerTokenNotStarted
  return bearer


askBearer :: forall m. ( MonadHandler m )
          => m (Maybe Jwt)
-- ^ Retrieve current `Jwt` from HTTP-Header, POST-Parameter, or GET-Parameter
askBearer = runMaybeT $ asum
  [ MaybeT lookupBearerAuth >>= hoistMaybe . fromPathPiece
  , MaybeT $ lookupGlobalPostParam PostBearer
  , MaybeT $ lookupGlobalGetParam GetBearer
  , fmap Jwt . MaybeT $ lookupSessionBS (toPathPiece SessionBearer)
  ]

formEmbedBearerPost, formEmbedBearerGet :: MonadHandler m => (Markup -> m a) -> (Markup -> m a)
formEmbedBearerPost f fragment = do
  mBearer <- askBearer
  f [shamlet|
       $newline never
       $maybe bearer <- mBearer
         <input type=hidden name=#{toPathPiece PostBearer} value=#{toPathPiece bearer}>
       #{fragment}
    |]
formEmbedBearerGet f fragment = do
  mBearer <- askBearer
  f [shamlet|
       $newline never
       $maybe bearer <- mBearer
         <input type=hidden name=#{toPathPiece GetBearer} value=#{toPathPiece bearer}>
       #{fragment}
    |]