{-# LANGUAGE UndecidableInstances #-}

module Model.Token
  ( BearerToken(..)
  , bearerToken
  , encodeToken, BearerTokenException(..), decodeToken
  , tokenToJSON, tokenParseJSON, tokenParseJSON'
  , askJwt
  ) where

import ClassyPrelude.Yesod
import Model
import Settings
import Utils (NTop(..), hoistMaybe)
import Utils.Lens hiding ((.=))
import Utils.Parameters

import Yesod.Auth (AuthId)

import qualified Jose.Jwa as Jose
import Jose.Jwk (JwkSet(..))
import Jose.Jwt (Jwt(..), IntDate(..))
import qualified Jose.Jwt as Jose

import Jose.Jwt.Instances ()
import Data.Aeson.Types.Instances ()

import Data.HashSet (HashSet)

import qualified Data.HashMap.Strict as HashMap

import Data.Aeson.Types (Parser, (.:?), (.:))
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Types 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 CryptoID

import Data.Time.Clock
import Data.Time.Clock.POSIX

import Control.Monad.Random (MonadRandom(..))
import Control.Monad.Trans.Maybe (MaybeT(..))


data BearerToken site = BearerToken
  { tokenIdentifier :: TokenId -- ^ Unique identifier for each token; maybe useful for tracing usage of tokens
  , tokenAuthority  :: AuthId site -- ^ Tokens only grant rights the `tokenAuthority` has (i.e. `AuthTag`s are evaluated with the user set to `tokenAuthority`)
  , tokenRoutes     :: Maybe (HashSet (Route site)) -- ^ Tokens can optionally be restricted to only be usable on a subset of routes
  , tokenAddAuth    :: Maybe AuthDNF -- ^ Tokens can specify an additional predicate logic formula of `AuthTag`s that needs to evaluate to `Authorized` in order for the token to be valid.
  , tokenIssuedAt   :: UTCTime
  , tokenIssuedBy   :: InstanceId
  , tokenExpiresAt
  , tokenStartsAt   :: Maybe UTCTime
  } deriving (Generic, Typeable)

deriving instance (Eq (AuthId site), Eq (Route site)) => Eq (BearerToken site)
deriving instance (Read (AuthId site), Eq (Route site), Hashable (Route site), Read (Route site)) => Read (BearerToken site)
deriving instance (Show (AuthId site), Show (Route site)) => Show (BearerToken site)

tokenToJSON :: forall m.
               ( MonadHandler m
               , HasCryptoUUID (AuthId (HandlerSite m)) m
               , RenderRoute (HandlerSite m)
               ) => BearerToken (HandlerSite m) -> m Value
tokenToJSON BearerToken{..} = do
  cID <- encrypt tokenAuthority :: m (CryptoUUID (AuthId (HandlerSite m)))
  let stdPayload = Jose.JwtClaims
        { jwtIss = Just $ toPathPiece tokenIssuedBy
        , jwtSub = Nothing
        , jwtAud = Nothing
        , jwtExp = IntDate . utcTimeToPOSIXSeconds <$> tokenExpiresAt
        , jwtNbf = IntDate . utcTimeToPOSIXSeconds <$> tokenStartsAt
        , jwtIat = Just . IntDate $ utcTimeToPOSIXSeconds tokenIssuedAt
        , jwtJti = Just $ toPathPiece tokenIdentifier
        }
  return . JSON.object $
    catMaybes [ Just $ "authority" .= cID
              , ("routes" .=) <$> tokenRoutes
              , ("add-auth" .=) <$> tokenAddAuth
              ]
    ++ let JSON.Object hm = toJSON stdPayload in HashMap.toList hm

tokenParseJSON :: forall site.
                  ( HasCryptoUUID (AuthId site) (ReaderT CryptoIDKey Parser)
                  , ParseRoute site
                  , Hashable (Route site)
                  )
               => Value
               -> ReaderT CryptoIDKey Parser (BearerToken site)
tokenParseJSON v@(Object o) = do
  tokenAuthority' <- lift (o .: "authority") :: ReaderT CryptoIDKey Parser (CryptoUUID (AuthId site))
  tokenAuthority <- decrypt tokenAuthority'

  tokenRoutes    <- lift $ o .:? "routes"
  tokenAddAuth   <- lift $ o .:? "add-auth"
  Jose.JwtClaims{..} <- lift $ parseJSON v

  let unIntDate (IntDate posix) = posixSecondsToUTCTime posix

  Just tokenIssuedBy <- return $ jwtIss >>= fromPathPiece
  Just tokenIdentifier <- return $ jwtJti >>= fromPathPiece
  Just tokenIssuedAt <- return $ unIntDate <$> jwtIat
  let tokenExpiresAt = unIntDate <$> jwtExp
      tokenStartsAt = unIntDate <$> jwtNbf

  return BearerToken{..}
tokenParseJSON v = lift $ JSON.typeMismatch "BearerToken" v

tokenParseJSON' :: forall 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)))
tokenParseJSON' = do
  cidKey <- cryptoIDKey return
  return $ flip runReaderT cidKey . tokenParseJSON


bearerToken :: forall m.
               ( MonadHandler m
               , HasInstanceID (HandlerSite m) InstanceId
               , HasCryptoUUID (AuthId (HandlerSite m)) m
               , HasAppSettings (HandlerSite m)
               )
            => 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))
bearerToken tokenAuthority tokenRoutes tokenAddAuth mTokenExpiresAt tokenStartsAt = do
  tokenIdentifier <- liftIO getRandom
  tokenIssuedAt <- liftIO getCurrentTime
  tokenIssuedBy <- getsYesod $ view instanceID

  defaultExpiration <- getsYesod $ view _appJWTExpiration
  
  let tokenExpiresAt
        | Just t <- mTokenExpiresAt
        = t
        | Just tDiff <- defaultExpiration
        = Just $ tDiff `addUTCTime` fromMaybe tokenIssuedAt tokenStartsAt
        | otherwise
        = Nothing

  return BearerToken{..}


jwtEncoding :: Jose.JwtEncoding
jwtEncoding = Jose.JwsEncoding Jose.HS256


encodeToken :: forall m.
               ( MonadHandler m
               , HasJSONWebKeySet (HandlerSite m) JwkSet
               , HasInstanceID (HandlerSite m) InstanceId
               , HasCryptoUUID (AuthId (HandlerSite m)) m
               , RenderRoute (HandlerSite m)
               )
            => BearerToken (HandlerSite m) -> m Jwt
encodeToken token = do
  payload <- Jose.Claims . toStrict . JSON.encode <$> tokenToJSON token
  JwkSet jwks <- getsYesod $ view jsonWebKeySet
  either throwM return =<< liftIO (Jose.encode jwks jwtEncoding payload)


data BearerTokenException
  = BearerTokenJwtError Jose.JwtError
  | BearerTokenUnsecured
  | BearerTokenInvalidFormat String
  | BearerTokenExpired | BearerTokenNotStarted
  deriving (Eq, Show, Generic, Typeable)

instance Exception BearerTokenException

decodeToken :: forall m.
               ( MonadHandler m
               , HasJSONWebKeySet (HandlerSite m) JwkSet
               , HasCryptoUUID (AuthId (HandlerSite m)) (ReaderT CryptoIDKey Parser)
               , MonadCryptoKey m ~ CryptoIDKey
               , MonadCrypto m
               , MonadThrow m
               , ParseRoute (HandlerSite m)
               , Hashable (Route (HandlerSite m))
               )
            => Jwt -> m (BearerToken (HandlerSite m))
decodeToken (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 <- tokenParseJSON'
  token@BearerToken{..} <- either (throwM . BearerTokenInvalidFormat . uncurry JSON.formatError) return $ JSON.eitherDecodeStrictWith JSON.jsonEOF' (JSON.iparse parser) content'
  now <- liftIO getCurrentTime
  unless (NTop tokenExpiresAt > NTop (Just now)) $
    throwM BearerTokenExpired
  unless (tokenStartsAt <= Just now) $
    throwM BearerTokenNotStarted
  return token


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