{-# LANGUAGE UndecidableInstances #-}

module Model.Tokens
  ( BearerToken(..)
  , _tokenIdentifier, _tokenAuthority, _tokenRoutes, _tokenAddAuth, _tokenRestrictions, _tokenRestrictionIx, _tokenRestrictionAt, _tokenIssuedAt, _tokenIssuedBy, _tokenExpiresAt, _tokenStartsAt
  , tokenRestrict
  , tokenToJSON, tokenParseJSON
  ) where

import ClassyPrelude.Yesod
import Yesod.Core.Instances ()
  
import Model
import Utils (assertM')
import Utils.Lens hiding ((.=))
import Data.Aeson.Lens (AsJSON(..))

import Yesod.Auth (AuthId)

import Jose.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.HashMap.Strict.Instances ()
import Data.HashSet.Instances ()
import Data.Time.Clock.Instances ()

import Data.Aeson.Types (Parser, (.:?), (.:), (.!=), (.=))
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Types as JSON

import CryptoID

import Data.Time.Clock.POSIX

import Data.Binary (Binary)



-- | Presenting a `BearerToken` transfers some authorisation from `tokenAuthority` to /whoever/ presents the token
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.
  , tokenRestrictions :: HashMap (Route site) Value
    -- ^ Tokens can be restricted to certain actions within the context of a route (i.e. creating an user only with a certain username, resetting a password only if the old matches a hash, ...)
    --
    -- In general this is not encrypted; some care is required to not expose sensitive information to the bearer of the token
  , 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)

instance (Binary (AuthId site), Binary (Route site), Hashable (Route site), Eq (Route site)) => Binary (BearerToken site)

makeLenses_ ''BearerToken

_tokenRestrictionIx :: (FromJSON a, ToJSON a, Hashable (Route site), Eq (Route site)) => Route site -> Traversal' (BearerToken site) a
-- ^ Focus a singular restriction (by route) if it exists
--
-- This /cannot/ be used to add restrictions, use `_tokenRestrictionAt` or `tokenRestrict` instead
_tokenRestrictionIx route = _tokenRestrictions . ix route . _JSON

_tokenRestrictionAt :: (FromJSON a, ToJSON a, Hashable (Route site), Eq (Route site)) => Route site -> Traversal' (BearerToken site) (Maybe a)
-- ^ Focus a singular restriction (by route) whether it exists, or not
_tokenRestrictionAt route = _tokenRestrictions . at route . maybePrism _JSON

tokenRestrict :: (ToJSON a, Hashable (Route site), Eq (Route site)) => Route site -> a -> BearerToken site -> BearerToken site
-- ^ Add a restriction to a `BearerToken`
--
-- If a restriction already exists for the targeted route, it's silently overwritten
tokenRestrict route (toJSON -> resVal) = over _tokenRestrictions $ HashMap.insert route resVal



tokenToJSON :: forall m.
               ( MonadHandler m
               , HasCryptoUUID (AuthId (HandlerSite m)) m
               , RenderRoute (HandlerSite m)
               ) => BearerToken (HandlerSite m) -> m Value
-- ^ Encode a `BearerToken` analogously to `toJSON`
--
-- Monadic context is needed because `AuthId`s are encrypted during encoding
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
              , ("restrictions" .=) <$> assertM' (not . HashMap.null) tokenRestrictions
              ]
    ++ 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)
-- ^ Decode a `Value` to a `BearerToken` analogously to `parseJSON`
--
-- Monadic context is needed because `AuthId`s are encrypted during encoding
--
-- It's usually easier to use `Utils.Tokens.tokenParseJSON'`
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"
  tokenRestrictions  <- lift $ o .:? "restrictions" .!= HashMap.empty
  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