module Utils.Tokens
  ( bearerToken
  , encodeToken, BearerTokenException(..), decodeToken
  , tokenParseJSON'
  , askJwt
  , formEmbedJwtPost, formEmbedJwtGet
  ) where

import ClassyPrelude.Yesod

import Yesod.Auth (AuthId)

import Utils (NTop(..), hoistMaybe, SessionKey(..))
import Utils.Parameters
import Utils.Lens

import Model
import Model.Tokens

import Jose.Jwk (JwkSet(..))
import Jose.Jwt (Jwt(..))
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 Data.Time.Clock

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

import Settings

import CryptoID

import Text.Blaze (Markup)


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)))
-- ^ Read `CryptoIDKey` from monadic context and construct a `Parser` for `BearerToken`s
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))
-- ^ Smart constructor for `BearerToken`, does not set route restrictions (due to polymorphism), use `tokenRestrict`
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
      tokenRestrictions = HashMap.empty

  return BearerToken{..}


encodeToken :: forall m.
               ( MonadHandler m
               , HasJSONWebKeySet (HandlerSite m) JwkSet
               , HasInstanceID (HandlerSite m) InstanceId
               , HasAppSettings (HandlerSite m)
               , HasCryptoUUID (AuthId (HandlerSite m)) m
               , RenderRoute (HandlerSite m)
               )
            => BearerToken (HandlerSite m) -> m Jwt
-- ^ Call `tokenToJSON` and encode the result as a `Jwt` according to `appJwtEncoding`
encodeToken token = do
  payload <- Jose.Claims . toStrict . JSON.encode <$> tokenToJSON token
  JwkSet jwks <- getsYesod $ view jsonWebKeySet
  jwtEncoding <- getsYesod $ view _appJwtEncoding
  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

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))
-- ^ Decode a `Jwt` and call `tokenParseJSON`
--
-- Throws `bearerTokenException`s
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
  , fmap Jwt . MaybeT $ lookupSessionBS (toPathPiece SessionBearer)
  ]

formEmbedJwtPost, formEmbedJwtGet :: MonadHandler m => (Markup -> m a) -> (Markup -> m a)
formEmbedJwtPost f fragment = do
  mJwt <- askJwt
  f [shamlet|
       $newline never
       $maybe jwt <- mJwt
         <input type=hidden name=#{toPathPiece PostBearer} value=#{toPathPiece jwt}>
       #{fragment}
    |]
formEmbedJwtGet f fragment = do
  mJwt <- askJwt
  f [shamlet|
       $newline never
       $maybe jwt <- mJwt
         <input type=hidden name=#{toPathPiece GetBearer} value=#{toPathPiece jwt}>
       #{fragment}
    |]