{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-}
{-# LANGUAGE UndecidableInstances #-}

module Model.Tokens.Bearer
  ( BearerToken(..), BearerTokenRouteMode(..)
  , _bearerIdentifier, _bearerAuthority, _bearerRoutes, _bearerAddAuth, _bearerRestrictions, _bearerRestrictionIx, _bearerRestrictionAt, _bearerIssuedAt, _bearerIssuedBy, _bearerExpiresAt, _bearerStartsAt
  , bearerRestrict
  , bearerToJSON, bearerParseJSON
  ) where

import ClassyPrelude.Yesod
import Yesod.Core.Instances ()
import Yesod.Servant (MonadSite(..))
  
import Model
import Model.Tokens.Lens
import Utils (assertM', foldMapM)
import Utils.Lens hiding ((.=))
import Data.Aeson.Lens (AsJSON(..))
import Utils.PathPiece

import Data.Universe

import Yesod.Auth (AuthId)

import Jose.Jwt (IntDate(..))
import qualified Jose.Jwt as Jose

import Jose.Jwt.Instances ()

import qualified Data.HashSet as HashSet
import qualified Data.HashMap.Strict as HashMap
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)

import qualified Data.CryptoID.Class.ImplicitNamespace as I



data BearerTokenRouteMode
  = BearerTokenRouteEval -- ^ Token is not to be evaluated for routes outside of the given restriction
  | BearerTokenRouteAccess -- ^ Token may be evaluated for routes outside of the given restriction, but not if the initial request was outside the restriction
  deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Typeable)
  deriving anyclass (Universe, Finite, Hashable, Binary)
nullaryPathPiece ''BearerTokenRouteMode $ camelToPathPiece' 3
pathPieceJSON ''BearerTokenRouteMode
pathPieceJSONKey ''BearerTokenRouteMode
instance Default BearerTokenRouteMode where
  def = BearerTokenRouteEval
  
-- | Presenting a `BearerToken` transfers some authorisation from `tokenAuthority` to /whoever/ presents the token
data BearerToken site = BearerToken
  { bearerIdentifier   :: TokenId
    -- ^ Unique identifier for each token; maybe useful for tracing usage of tokens
  , bearerAuthority    :: HashSet (Either Value (AuthId site))
    -- ^ Tokens only grant rights the `bearerAuthority` has (i.e. `AuthTag`s are evaluated with the user set to `bearerAuthority`)
  , bearerImpersonate  :: Maybe (AuthId site)
    -- ^ Token doubles as session token; i.e. if presented `maybeAuthId` etc. should evaluate to the given value
  , bearerRoutes       :: HashMap BearerTokenRouteMode (HashSet (Route site))
    -- ^ Tokens can optionally be restricted to only be usable on a subset of routes
  , bearerAddAuth      :: 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.
  , bearerRestrictions :: 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
  , bearerIssuedAt     :: UTCTime
  , bearerIssuedBy     :: InstanceId
  , bearerIssuedFor    :: ClusterId
  , bearerExpiresAt
  , bearerStartsAt     :: Maybe UTCTime
  } deriving (Generic, Typeable)

deriving stock instance (Eq (AuthId site), Eq (Route site)) => Eq (BearerToken site)
deriving stock instance (Ord (AuthId site), Ord (Route site)) => Ord (BearerToken site)
deriving stock instance (Read (AuthId site), Eq (Route site), Hashable (Route site), Read (Route site), Hashable (AuthId site), Eq (AuthId site)) => Read (BearerToken site)
deriving stock instance (Show (AuthId site), Show (Route site), Hashable (AuthId site)) => Show (BearerToken site)
deriving anyclass instance (Hashable (AuthId site), Hashable (Route site)) => Hashable (BearerToken site)
deriving anyclass instance (Binary (AuthId site), Binary (Route site), Hashable (Route site), Eq (Route site), Hashable (AuthId site), Eq (AuthId site)) => Binary (BearerToken site)

makeLenses_ ''BearerToken
instance HasTokenIdentifier (BearerToken site) TokenId where
  _tokenIdentifier = _bearerIdentifier
instance HasTokenIssuedBy (BearerToken site) InstanceId where
  _tokenIssuedBy = _bearerIssuedBy
instance HasTokenIssuedFor (BearerToken site) ClusterId where
  _tokenIssuedFor = _bearerIssuedFor
instance HasTokenIssuedAt (BearerToken site) UTCTime where
  _tokenIssuedAt = _bearerIssuedAt
instance HasTokenExpiresAt (BearerToken site) (Maybe UTCTime) where
  _tokenExpiresAt = _bearerExpiresAt
instance HasTokenStartsAt (BearerToken site) (Maybe UTCTime) where
  _tokenStartsAt = _bearerStartsAt

_bearerRestrictionIx :: (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 `_bearerRestrictionAt` or `bearerRestrict` instead
_bearerRestrictionIx route = _bearerRestrictions . ix route . _JSON

_bearerRestrictionAt :: (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
_bearerRestrictionAt route = _bearerRestrictions . at route . maybePrism _JSON

bearerRestrict :: (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
bearerRestrict route (toJSON -> resVal) = over _bearerRestrictions $ HashMap.insert route resVal


bearerToJSON :: forall site m.
               ( MonadSite site m
               , HasCryptoUUID (AuthId site) m
               , RenderRoute site 
               ) => BearerToken site -> m Value
-- ^ Encode a `BearerToken` analogously to `toJSON`
--
-- Monadic context is needed because `AuthId`s are encrypted during encoding
bearerToJSON BearerToken{..} = do
  cID <- foldMapM (fmap HashSet.singleton . either (return . Left) (fmap Right . I.encrypt)) bearerAuthority :: m (HashSet (Either Value (CryptoUUID (AuthId site))))
  let stdPayload = Jose.JwtClaims
        { jwtIss = Just $ toPathPiece bearerIssuedBy
        , jwtSub = Nothing
        , jwtAud = Just . pure $ toPathPiece bearerIssuedFor
        , jwtExp = IntDate . utcTimeToPOSIXSeconds <$> bearerExpiresAt
        , jwtNbf = IntDate . utcTimeToPOSIXSeconds <$> bearerStartsAt
        , jwtIat = Just . IntDate $ utcTimeToPOSIXSeconds bearerIssuedAt
        , jwtJti = Just $ toPathPiece bearerIdentifier
        }

      authorityToJSON auths | [auth] <- otoList auths = either toJSON toJSON auth
                            | otherwise = toJSON $ HashSet.map (either toJSON toJSON) auths
  iCID <- traverse I.encrypt bearerImpersonate :: m (Maybe (CryptoUUID (AuthId site)))
  
  return . JSON.object $
    catMaybes [ Just $ "authority" .= authorityToJSON cID
              , ("impersonate" .=) <$> iCID
              , ("routes" .=) <$> assertM' (not . HashMap.null) bearerRoutes
              , ("add-auth" .=) <$> bearerAddAuth
              , ("restrictions" .=) <$> assertM' (not . HashMap.null) bearerRestrictions
              ]
    ++ let JSON.Object hm = toJSON stdPayload in HashMap.toList hm

bearerParseJSON :: forall site.
                  ( Hashable (AuthId site), Eq (AuthId 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.bearerParseJSON'`
bearerParseJSON v@(Object o) = do
  bearerAuthority' <- lift $ asum
    [ HashSet.singleton . Right <$> o .: "authority"
    , (o .: "authority" :: Parser (HashSet Value)) >>= foldMapM (\v' -> fmap HashSet.singleton $ (Right <$> parseJSON v') <|> return (Left v'))
    , HashSet.singleton . Left  <$> o .: "authority"
    ] :: ReaderT CryptoIDKey Parser (HashSet (Either Value (CryptoUUID (AuthId site))))
  bearerAuthority <- foldMapM (fmap HashSet.singleton . either (return . Left) (fmap Right . I.decrypt)) bearerAuthority'

  bearerImpersonate   <- traverse I.decrypt <=< lift $ (o .:? "impersonate" :: Parser (Maybe (CryptoUUID (AuthId site))))
  bearerRoutes        <- lift $ (o .:? "routes" .!= HashMap.empty)
                            <|> (maybe HashMap.empty (HashMap.singleton def) <$> o .:? "routes")
  bearerAddAuth       <- lift $ o .:? "add-auth"
  bearerRestrictions  <- lift $ o .:? "restrictions" .!= HashMap.empty
  Jose.JwtClaims{..} <- lift $ parseJSON v

  let unIntDate (IntDate posix) = posixSecondsToUTCTime posix

  Just bearerIssuedBy <- return $ jwtIss >>= fromPathPiece
  Just bearerIssuedFor <- return $ do
    [aud] <- jwtAud
    fromPathPiece aud
  Just bearerIdentifier <- return $ jwtJti >>= fromPathPiece
  Just bearerIssuedAt <- return $ unIntDate <$> jwtIat
  let bearerExpiresAt = unIntDate <$> jwtExp
      bearerStartsAt = unIntDate <$> jwtNbf

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