{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}

module Thentos.CookieSession
    ( serveFAction
    , enterFAction
    , noopExtendClearanceOnSessionToken

    -- * types
    , SSession
    , FSession
    , FSessionMap
    , FSessionStore
    , FServantSession
    , FSessionKey
    , ExtendClearanceOnSessionToken
    )
where

import Control.Lens (use)
import Control.Monad (when)
import Control.Monad.Except.Missing (finally)
import Control.Monad.State.Class (MonadState(..))
import Control.Monad.Trans.Except (ExceptT)
import Crypto.Random (MonadRandom(..))
import Data.Char (ord)
import Data.Maybe (isJust)
import Data.Proxy (Proxy(Proxy))
import Data.String.Conversions
import Network.Wai (Middleware, Application, vault)
import Network.Wai.Session (SessionStore, Session, withSession)
import Servant (ServantErr, (:>), serve, HasServer, ServerT, Server, (:~>)(Nat), unNat)
import Servant.Server.Internal (route, passToServer)
import Servant.Utils.Enter (Enter, enter)
import Web.Cookie (SetCookie, setCookieName)

import qualified Data.ByteString as SBS
import qualified Data.Vault.Lazy as Vault
import qualified Network.Wai.Session.Map as SessionMap

import Servant.Missing (MonadError500, throwError500)
import Thentos.CookieSession.CSRF
import Thentos.CookieSession.Types (ThentosSessionToken, MonadUseThentosSessionToken, getThentosSessionToken)

-- * servant integration

-- | @SSession m k v@ represents a session storage with keys of type @k@,
-- values of type @v@, and operating under the monad @m@.
-- The underlying implementation uses the 'wai-session' package, and any
-- backend compatible with that package should work here too.
data SSession (m :: * -> *) (k :: *) (v :: *)

-- | 'HasServer' instance for 'SSession'.
instance (HasServer sublayout context) => HasServer (SSession n k v :> sublayout) context where
  type ServerT (SSession n k v :> sublayout) m
    = (Vault.Key (Session n k v) -> Maybe (Session n k v)) -> ServerT sublayout m
  route Proxy context subserver =
    route (Proxy :: Proxy sublayout) context (passToServer subserver go)
    where
      go request key = Vault.lookup key $ vault request


-- * middleware

type FSession        fsd = Session IO () fsd
type FSessionMap     fsd = Vault.Key (FSession fsd) -> Maybe (FSession fsd)
type FSessionStore   fsd = SessionStore IO () fsd
type FServantSession fsd = SSession IO () fsd
type FSessionKey     fsd = Vault.Key (FSession fsd)

cookieName :: SetCookie -> SBS
cookieName setCookie =
    if cookieNameValid n
        then n
        else error $ "Thentos.CookieSession: bad cookie name: " ++ show n
  where
    n = setCookieName setCookie

cookieNameValid :: SBS -> Bool
cookieNameValid = SBS.all (`elem` (fromIntegral . ord <$> '_':['a'..'z']))

sessionMiddleware :: Proxy fsd -> SetCookie -> IO (Middleware, FSessionKey fsd)
sessionMiddleware Proxy setCookie = do
    smap :: FSessionStore fsd <- SessionMap.mapStore_
    key  :: Vault.Key (FSession fsd) <- Vault.newKey
    return (withSession smap (cookieName setCookie) setCookie key, key)


-- * frontend action monad

type ExtendClearanceOnSessionToken m = ThentosSessionToken -> m ()

noopExtendClearanceOnSessionToken :: Monad m => ExtendClearanceOnSessionToken m
noopExtendClearanceOnSessionToken _ = pure ()

-- | The 'ExtendClearanceOnSessionToken' argument can be used in combination with the lio package to
-- write the access policy into the server monad state.  If your way of access restriction has no
-- need for that, simply pass @noopExtendClearanceOnSessionToken@.
serveFAction :: forall api m s e v.
        ( HasServer api '[]
        , Enter (ServerT api m) (m :~> ExceptT ServantErr IO) (Server api)
        , MonadRandom m, MonadError500 e m, MonadHasSessionCsrfToken s m
        , MonadViewCsrfSecret v m, MonadUseThentosSessionToken s m
        )
     => Proxy api
     -> Proxy s
     -> SetCookie
     -> ExtendClearanceOnSessionToken m
     -> IO :~> m
     -> m :~> ExceptT ServantErr IO
     -> ServerT api m -> IO Application
serveFAction _ sProxy setCookie extendClearanceOnSessionToken ioNat nat fServer =
    app <$> sessionMiddleware sProxy setCookie
  where
    app :: (Middleware, FSessionKey s) -> Application
    app (mw, key) = mw $ serve (Proxy :: Proxy (FServantSession s :> api)) (server' key)

    server' :: FSessionKey s -> FSessionMap s -> Server api
    server' key smap = enter nt fServer
      where
        nt :: m :~> ExceptT ServantErr IO
        nt = enterFAction key smap extendClearanceOnSessionToken ioNat nat

enterFAction
    :: ( MonadRandom m, MonadError500 e m, MonadHasSessionCsrfToken s m
       , MonadViewCsrfSecret v m, MonadUseThentosSessionToken s m)
    => FSessionKey s
    -> FSessionMap s
    -> ExtendClearanceOnSessionToken m
    -> IO :~> m
    -> m :~> ExceptT ServantErr IO
    -> m :~> ExceptT ServantErr IO
enterFAction key smap extendClearanceOnSessionToken ioNat nat = Nat $ \fServer -> unNat nat $ do
    case smap key of
        Nothing ->
            -- FIXME: this case should not be code 500, as it can (probably) be provoked by
            -- the client.
            throwError500 "Could not read cookie."
        Just (lkup, ins) -> do
            cookieToFSession ioNat (lkup ())
            maybeSessionToken <- use getThentosSessionToken

            -- update privileges
            mapM_ extendClearanceOnSessionToken maybeSessionToken

            -- refresh the CSRF token if there is a session token
            when (isJust maybeSessionToken) refreshCsrfToken

            fServer `finally` (do
                clearCsrfToken  -- could be replaced by 'refreshCsrfToken'
                cookieFromFSession ioNat (ins ()))

-- | Write 'FrontendSessionData' from the 'SSession' state to 'MonadFAction' state.  If there
-- is no state, do nothing.
cookieToFSession :: MonadState s m => IO :~> m -> IO (Maybe s) -> m ()
cookieToFSession ioNat r = unNat ioNat r >>= mapM_ put

-- | Read 'FrontendSessionData' from 'MonadFAction' and write back into 'SSession' state.
cookieFromFSession :: MonadState s m => IO :~> m -> (s -> IO ()) -> m ()
cookieFromFSession ioNat w = get >>= unNat ioNat . w