{-# 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.Base16 as Hex
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 Data.Time (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, (:<|>) (..), err404)
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
    :<|> unregisterUser
    :<|> getLeaderboard
    :<|> getUser
    :<|> 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 Nothing
          updatedMap = Map.insert token score usersMap
      writeTVar tvar updatedMap
  pure token

unregisterUser :: (MonadReader Env m, MonadIO m) => Maybe Token -> m Score
unregisterUser Nothing = liftIO $ throwM $ err400 {errBody = "Missing token"}
unregisterUser (Just token) = do
  checkUserToken token
  score <-
    asks users
      >>= liftIO . atomically . \tvar -> do
        usersMap <- readTVar tvar
        let score = Map.lookup token usersMap
            updatedMap = Map.delete token usersMap
        writeTVar tvar updatedMap
        pure score
  case score of
    Nothing -> liftIO $ throwM $ err400 {errBody = "User does not exist"}
    Just s -> pure s

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 $ Hex.encode $ LBS.toStrict tokenBytes

getLeaderboard :: (MonadIO m, MonadReader Env m) => m [Score]
getLeaderboard = do
  usersMap <- asks users >>= liftIO . readTVarIO
  pure $ Map.elems usersMap

getUser :: (MonadIO m, MonadReader Env m) => Maybe Token -> m Score
getUser mToken = do
  usersMap <- asks users >>= liftIO . readTVarIO
  case mToken >>= (`Map.lookup` usersMap) of
    Nothing -> liftIO $ throwM $ err404 {errBody = "User does not exist"}
    Just score -> pure score

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 -> TestResult -> m ()
updateUserLastTest token result = 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, lastResult = Just (isSuccess result)}
            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
  Challenge {expr} <- getCurrentChallenge token
  let PurerEnv {outputs} = withInputs initialREPL [In expr]
      result = evaluateOutputs outputs expected
  updateUserLastTest token result
  case result of
    TestSucceeded ->
      upgradeUserChallenge token
    TestFailed msg ->
      liftIO $ throwM $ err400 {errBody = "Test failed: " <> LBS.fromStrict (encodeUtf8 msg)}
    TestErrored msg ->
      liftIO $ throwM $ err500 {errBody = "Error evaluating answer: " <> LBS.fromStrict (encodeUtf8 msg)}

data TestResult = TestSucceeded | TestFailed Text | TestErrored Text

isSuccess :: TestResult -> Bool
isSuccess TestSucceeded = True
isSuccess _ = False

evaluateOutputs :: [Out] -> Text -> TestResult
evaluateOutputs outputs expected =
  case outputs of
    [Bye, Evaluated ast _ _] ->
      if pretty ast /= expected
        then TestFailed "Incorrect result"
        else TestSucceeded
    [Bye, Msg msg] -> TestFailed msg
    _ -> TestErrored $ 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")