{-# 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, pack)
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 Minilang.Lambda.Parser (pretty)
import Minilang.REPL (withInputs)
import Minilang.REPL.Purer (PurerEnv (..))
import Minilang.REPL.Types (In (..), Out (..), initialREPL)
import Network.Wai.Application.Static (defaultFileServerSettings, staticApp)
import Servant (ServerT, err400, err500, errBody, (:<|>) (..))
import Servant.Server (Application, Handler (..), Tagged (..))
import System.Random (StdGen)

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

makeEnv :: StdGen -> [Challenge] -> IO Env
makeEnv gen challenges =
  Env
    <$> newTVarIO gen
    <*> newTVarIO mempty
    <*> pure challenges

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
    :<|> postAnswer
    :<|> 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, MonadReader Env m) => m [Score]
getLeaderboard = do
  usersMap <- asks users >>= liftIO . readTVarIO
  pure $ Map.elems usersMap

getChallenge :: (MonadReader Env m, MonadIO m) => Maybe Token -> m Challenge
getChallenge Nothing = liftIO $ throwM $ err400 {errBody = "Missing token"}
getChallenge (Just token) = do
  checkUserToken token
  getCurrentChallenge token

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 -> m ()
updateUserLastTest token = do
  time <- liftIO getCurrentTime
  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

postAnswer :: (MonadReader Env m, MonadIO m) => Maybe Token -> Text -> m ()
postAnswer Nothing _ = liftIO $ throwM $ err400 {errBody = "Missing token"}
postAnswer (Just token) expected = do
  checkUserToken token
  updateUserLastTest token
  Challenge {expr} <- getCurrentChallenge token
  let PurerEnv {outputs} = withInputs initialREPL [In expr]
  case outputs of
    [Bye, Evaluated ast _ _] -> do
      when (pretty ast /= expected) $
        liftIO $
          throwM $
            err400 {errBody = "Incorrect answer"}
      upgradeUserChallenge token
    [Bye, Msg msg] -> liftIO $ throwM $ err400 {errBody = "Evaluation error: " <> LBS.fromStrict (encodeUtf8 msg)}
    _ -> liftIO $ throwM $ err500 {errBody = "Unexpected REPL output: " <> LBS.fromStrict (encodeUtf8 (pack $ show outputs))}

upgradeUserChallenge :: (MonadIO m, MonadReader Env m) => Token -> m ()
upgradeUserChallenge token = do
  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 {grade = succ (grade score)}
            updatedMap = Map.insert token updatedScore usersMap
        writeTVar user updatedMap


getCurrentChallenge :: (MonadIO m, MonadReader Env m) => Token -> m Challenge
getCurrentChallenge token = do
  Env{users, challenges} <- ask
  liftIO $ atomically $ do
    usersMap <- readTVar users
    case Map.lookup token usersMap of
      Nothing -> throwM $ err400 {errBody = "Cannot find challenge: Invalid token"}
      Just Score {grade} ->
        pure $ challenges !! grade

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