{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module Minilang.App where

import Control.Concurrent.STM (TVar, atomically, newTVarIO, readTVar, readTVarIO, writeTVar)
import Control.Exception (try)
import Control.Monad (when)
import Control.Monad.Catch (MonadThrow (throwM))
import Control.Monad.Except (ExceptT (..))
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Reader (MonadReader, ReaderT, ask, asks, runReaderT)
import qualified Data.ByteString.Base64 as Base64
import qualified Data.ByteString.Lazy as LBS
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import qualified Data.Text.Lazy.Encoding as LT
import Data.Time (UTCTime, getCurrentTime)
import Minilang.API (Challenge (..), LambdaAPI, Score (..), Token (..), User (..), newToken)
import Network.Wai.Application.Static (defaultFileServerSettings, staticApp)
import Servant (ServerT, err400, errBody, (:<|>) (..))
import Servant.Server (Application, Handler (..), Tagged (..))
import System.Random (StdGen)

data Env = Env
  { gen :: TVar StdGen,
    users :: TVar (Map Token Score)
  }

makeEnv :: StdGen -> IO Env
makeEnv gen =
  Env
    <$> newTVarIO gen
    <*> newTVarIO mempty

newtype AppM a = AppM {runAppM :: ReaderT Env IO a}
  deriving newtype (Functor, Applicative, Monad, MonadIO, MonadReader Env)

runApp :: Env -> AppM a -> Handler a
runApp env = Handler . ExceptT . try . flip runReaderT env . runAppM

lambdaApp :: (MonadIO m, MonadReader Env m) => ServerT LambdaAPI m
lambdaApp =
  registerUser
    :<|> getLeaderboard
    :<|> getChallenge
    :<|> Tagged staticResources

registerUser :: (MonadReader Env m, MonadIO m) => User -> m Token
registerUser user@(User name) = do
  checkDuplicateUser name
  token <- generateToken
  asks users
    >>= liftIO . atomically . \tvar -> do
      usersMap <- readTVar tvar
      let score = Score user 0 Nothing
          updatedMap = Map.insert token score usersMap
      writeTVar tvar updatedMap
  pure token

checkDuplicateUser :: (MonadIO m, MonadReader Env m) => Text -> m ()
checkDuplicateUser name =
  asks users >>= \user -> liftIO $ atomically $ do
    usersMap <- readTVar user
    when (isUserRegistered name usersMap) $
      throwM $
        err400 {errBody = "User already registered"}

isUserRegistered :: Text -> Map Token Score -> Bool
isUserRegistered name = Map.foldr (\Score {user = User {username}} acc -> username == name || acc) False

generateToken :: (MonadReader Env m, MonadIO m) => m Token
generateToken =
  asks gen
    >>= liftIO . atomically . \tvar -> do
      genVal <- readTVar tvar
      let (tokenBytes, newGen) = newToken genVal
      writeTVar tvar newGen
      pure $ Token $ decodeUtf8 $ Base64.encode $ LBS.toStrict tokenBytes

getLeaderboard :: (MonadIO m) => m [Score]
getLeaderboard = do
  now <- liftIO getCurrentTime
  pure [Score (User "Alice") 95 (Just now), Score (User "Bob") 88 (Just now)]

getChallenge :: (MonadReader Env m, MonadIO m) => Maybe Token -> m Challenge
getChallenge Nothing = liftIO $ throwM $ err400 {errBody = "Missing token"}
getChallenge (Just token) = do
  checkUserToken token
  now <- liftIO getCurrentTime
  updateUserLastTest token now
  pure $ Challenge 1 "x"

checkUserToken :: (MonadReader Env m, MonadIO m) => Token -> m ()
checkUserToken token = do
  usersMap <- asks users >>= liftIO . readTVarIO
  case Map.lookup token usersMap of
    Nothing -> liftIO $ throwM $ err400 {errBody = "Token " <> LBS.fromStrict (encodeUtf8 (unToken token)) <> " does not exist"}
    Just _ -> pure ()

updateUserLastTest :: (MonadIO m, MonadReader Env m) => Token -> UTCTime -> m ()
updateUserLastTest token time =
  asks users >>= \user -> liftIO $ atomically $ do
    usersMap <- readTVar user
    case Map.lookup token usersMap of
      Nothing -> throwM $ err400 {errBody = "Cannot update last test: Invalid token"}
      Just score -> do
        let updatedScore = score {lastTest = Just time}
            updatedMap = Map.insert token updatedScore usersMap
        writeTVar user updatedMap

-- | Serve static resources under `public/` directory
staticResources :: Application
staticResources = staticApp (defaultFileServerSettings "public")