module Webserver.Data.Sql.Cache.Session where

----------------------------------------------------------------------------------------------------

import qualified Data.Aeson                as Aeson
import qualified Data.ByteString           as ByteString
import qualified Data.Char                 as Char
import qualified Data.Text                 as Text
import qualified Data.Text.Encoding        as TSE
import qualified Data.Time.Clock           as Clock
import qualified Data.Time.Clock.POSIX     as Posix
import qualified Data.Time.Format          as TimeF
import qualified Data.UUID                 as UUID
import qualified GHC.Real                  as Real
import qualified Hasql.Decoders            as Decoders
import qualified Hasql.Encoders            as Encoders
import           Hasql.Generate            (HasInsert (..), HasUpdate (..))
import qualified Hasql.Pool                as Pool
import qualified Hasql.Session             as Session
import qualified Hasql.Statement           as Statement
import qualified Jose.Jwa                  as JWA
import qualified Jose.Jwk                  as JWK
import qualified Jose.Jwt                  as JWT
import           Network.HTTP.Types.Header (Header)
import qualified System.Environment        as System
import           Webserver.Control.Nil     (Nil (..))
import qualified Webserver.Control.Sql     as Sql

----------------------------------------------------------------------------------------------------

$(Sql.generateTable "cache" "sessions")

----------------------------------------------------------------------------------------------------

instance Nil Sessions where
  nil = Sessions {id = nil, token = nil, expiresOn = nil, userId = nil}

----------------------------------------------------------------------------------------------------

{-  Look up a session by token via raw SQL.
-}
sessionByTokenStatement :: Statement.Statement Text (Maybe Sessions)
sessionByTokenStatement =
  Statement.Statement
    "SELECT id, token, expires_on, user_id FROM cache.sessions WHERE token = $1"
    (Encoders.param (Encoders.nonNullable Encoders.text))
    (Decoders.rowMaybe sessionsDecoder)
    True

----------------------------------------------------------------------------------------------------

jwk :: IO JWK.Jwk
jwk = do
  mbSecret <- System.lookupEnv "JWT_SECRET"
  case mbSecret of
    Just secret ->
      return
        <| JWK.SymmetricJwk
          (ByteString.pack <| (Real.fromIntegral .> Char.ord) <$> secret)
          (Just <| JWT.KeyId "WebserverKeyId")
          (Just JWK.Sig)
          (Just <| JWA.Signed JWA.HS256)
    Nothing -> do
      JWK.generateSymmetricKey
        32
        (JWT.KeyId "WebserverKeyId")
        JWK.Sig
        (Just <| JWA.Signed JWA.HS256)
        >>= \key@(JWK.SymmetricJwk secret _ _ _) ->
          System.setEnv "JWT_SECRET" (show secret)
            >> return key

newSession :: JWK.Jwk -> Maybe UUID -> IO (Maybe Sessions)
newSession key mbUserId = do
  timeNow <- Clock.getCurrentTime
  let exp =
        Clock.addUTCTime
          (Clock.secondsToNominalDiffTime <| 60 * 60 * 24 * 30)
          timeNow
  let aud =
        case mbUserId of
          Nothing -> Nothing
          Just uuid_ ->
            Just [UUID.toText uuid_]
  let claims =
        JWT.JwtClaims
          (Just "yourwebsite.com")
          Nothing
          aud
          (Just <| JWT.IntDate <| Posix.utcTimeToPOSIXSeconds exp)
          Nothing
          (Just <| JWT.IntDate <| Posix.utcTimeToPOSIXSeconds timeNow)
          Nothing
  eitherJwt <-
    JWT.encode
      [key]
      (JWT.JwsEncoding JWA.HS256)
      (JWT.Claims <| ByteString.toStrict <| Aeson.encode claims)
  case eitherJwt of
    Left _ -> return Nothing
    Right jwt ->
      return
        <| Just
        <| Sessions
          nil
          (JWT.unJwt jwt |> TSE.decodeUtf8)
          exp
          mbUserId

updateSession :: Sessions -> JWK.Jwk -> Maybe UUID -> IO (Maybe Sessions)
updateSession sn key mbUserId = do
  oldToken <-
    JWT.decode
      [key]
      (Just <| JWT.JwsEncoding JWA.HS256)
      (TSE.encodeUtf8 sn.token)
  case oldToken of
    Right (JWT.Jws (_, claimString)) ->
      case (Aeson.decodeStrict claimString :: Maybe JWT.JwtClaims) of
        Nothing -> return Nothing
        Just claims ->
          let newAud =
                case mbUserId of
                  Nothing -> Nothing
                  Just uuid_ ->
                    Just [UUID.toText uuid_]

              newClaims =
                claims {JWT.jwtAud = newAud}
           in do
                JWT.encode
                  [key]
                  (JWT.JwsEncoding JWA.HS256)
                  (JWT.Claims <| ByteString.toStrict <| Aeson.encode newClaims)
                  >>= \case
                    Left _ -> return Nothing
                    Right jwt ->
                      return
                        <| Just
                        <| sn
                          { token = JWT.unJwt jwt |> TSE.decodeUtf8
                          , userId = mbUserId
                          }
    _ ->
      return Nothing

----------------------------------------------------------------------------------------------------

createSession :: Pool.Pool -> Sessions -> IO Sessions
createSession pool sn =
  Sql.runSession pool (insert sn)

getSessionByToken :: Pool.Pool -> Text -> IO (Maybe Sessions)
getSessionByToken pool token_ =
  Sql.runSession pool (Session.statement token_ sessionByTokenStatement)

updateSessionRecord :: Pool.Pool -> Sessions -> IO ()
updateSessionRecord pool sn =
  Sql.runSession pool (update sn >> return ())

validateSession :: Pool.Pool -> JWK.Jwk -> ByteString.ByteString -> IO (Maybe Sessions)
validateSession pool key token_ =
  let decodeToken :: ByteString.ByteString -> IO (Maybe JWT.JwtClaims)
      decodeToken tokenBytes =
        JWT.decode [key] (Just <| JWT.JwsEncoding JWA.HS256) tokenBytes >>= \case
          Right (JWT.Jws (_, claimString)) ->
            return (Aeson.decodeStrict claimString :: Maybe JWT.JwtClaims)
          _ -> return Nothing

      isValidIssuer :: JWT.JwtClaims -> Bool
      isValidIssuer claims = claims.jwtIss == Just "yourwebsite.com"

      isNotExpired :: Clock.UTCTime -> JWT.JwtClaims -> Sessions -> Bool
      isNotExpired timeNow claims session =
        case claims.jwtExp of
          Nothing -> False
          Just (JWT.IntDate posix) ->
            let expiryTime = Posix.posixSecondsToUTCTime posix
                sameDayExpiry = Clock.utctDay expiryTime == Clock.utctDay session.expiresOn
                notExpiredYet = expiryTime > timeNow
             in sameDayExpiry && notExpiredYet

      isValidAudience :: JWT.JwtClaims -> Sessions -> Bool
      isValidAudience claims session =
        case (claims.jwtAud, session.userId) of
          (Just [mbUserId], Just userId_) -> mbUserId == UUID.toText userId_
          (Nothing, Nothing)              -> True
          _                               -> False

      validateClaims :: Clock.UTCTime -> JWT.JwtClaims -> Sessions -> Maybe Sessions
      validateClaims timeNow claims session =
        let validIssuer = isValidIssuer claims
            notExpired = isNotExpired timeNow claims session
            validAudience = isValidAudience claims session
         in if validIssuer && notExpired && validAudience
              then Just session
              else Nothing
   in do
        timeNow <- Clock.getCurrentTime
        mbSession <- getSessionByToken pool (TSE.decodeUtf8 token_)
        case mbSession of
          Nothing -> return Nothing
          Just session -> do
            mbClaims <- decodeToken token_
            case mbClaims of
              Nothing     -> return Nothing
              Just claims -> return (validateClaims timeNow claims session)

packSessionCookie :: Sessions -> Header
packSessionCookie sn =
  let expires = Text.pack <| TimeF.formatTime TimeF.defaultTimeLocale TimeF.rfc822DateFormat sn.expiresOn
      cookieValue =
        "WebserverToken="
          <> sn.token
          <> "; Expires="
          <> expires
          <> "; Path=/; Secure; HttpOnly; SameSite=Lax"
   in ("Set-Cookie", TSE.encodeUtf8 cookieValue)

packClearSessionCookie :: Header
packClearSessionCookie =
  ("Set-Cookie", "WebserverToken=; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Path=/; Secure; HttpOnly; SameSite=Lax")