{-# LANGUAGE GeneralizedNewtypeDeriving #-}

-- | A version of REPL that interacts through a websocket.
module Minilang.REPL.Net where

import Control.Concurrent.STM (atomically, newTVar, readTVar)
import Control.Concurrent.STM.TVar
  ( TVar,
    newTVarIO,
    readTVarIO,
    writeTVar,
  )
import Control.Monad.Catch (MonadCatch (..), MonadThrow (..))
import Control.Monad.State
  ( MonadIO (..),
    MonadState,
    MonadTrans (..),
    StateT (StateT),
    evalStateT,
    gets,
  )
import Data.Aeson
  ( ToJSON,
    eitherDecode,
    encode,
    object,
    (.=),
  )
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as LBS
import Data.Either (fromRight)
import Data.Map as Map (Map, insert, lookup)
import Data.Serialize
  ( Serialize (..),
    getLazyByteString,
    getWord64be,
    putLazyByteString,
    putWord64be,
  )
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8With)
import Data.Text.Encoding.Error (lenientDecode)
import Minilang.Log (LoggerEnv (logInfo))
import Minilang.REPL.Pure (render)
import Minilang.REPL.Run (runREPL)
import Minilang.REPL.Types
  ( In (In),
    MonadREPL (..),
    Out,
    REPLEnv,
    initialREPL,
  )
import Network.Wai (Application)
import Network.Wai.Handler.WebSockets (websocketsOr)
import qualified Network.WebSockets as WS

-- | Configuration of the WS server
data ServerConfig = ServerConfig
  { -- | The host/IP to listen on. If 0.0.0.0, then the server listens
    configHost :: String,
    -- | The port to listen on. If 0, a random port will be assigned.
    configPort :: Int
  }
  deriving (Eq, Show)

data NetEnv = NetEnv
  { connection :: WS.Connection,
    logger :: LoggerEnv IO,
    repl :: TVar REPLEnv
    --    , storage :: FileStorage
  }

newtype Net m a = Net {runNet :: StateT NetEnv m a}
  deriving (Functor, Applicative, Monad, MonadTrans, MonadIO, MonadState NetEnv)

instance Serialize In where
  put out = do
    let bs = encode out
    putWord64be (fromIntegral $ LBS.length bs)
    putLazyByteString bs

  get = do
    len <- getWord64be
    bs <- getLazyByteString (fromIntegral len)
    case eitherDecode bs of
      Right out -> pure out
      Left err -> fail err

-- doStore :: In -> Net IO (Either Text ())
-- doStore out = do
--     fs <- gets storage
--     res <- liftIO $ store fs (pure $ Right out) handleStorageResult
--     case res of
--         WriteSucceed _ -> pure $ Right ()
--         err -> pure $ Left (pack $ show err)
--   where
--     handleStorageResult (Right (WriteSucceed _)) = pure $ Right ()
--     handleStorageResult (Right err) = pure $ Left $ pack $ show err
--     handleStorageResult (Left err) = pure $ Left $ pack err

netLog :: (ToJSON a) => a -> Net IO ()
netLog logEntry = gets logger >>= liftIO . flip logInfo logEntry

instance MonadREPL (Net IO) where
  input = do
    inp <- gets connection >>= lift . wsReceive
    -- res <- doStore inp
    -- case res of
    --   Right () -> netLog (object ["action" .= ("input" :: Text), "content" .= inp])
    --   Left txt -> netLog (object ["action" .= ("storageError" :: Text), "content" .= txt])
    netLog (object ["action" .= ("input" :: Text), "content" .= inp])
    pure inp

  output a = do
    gets connection >>= lift . wsSend a
    netLog
      ( object
          [ "action" .= ("output" :: Text),
            "content" .= a,
            "rendered" .= render a
          ]
      )

  prompt = pure ()
  getEnv = gets repl >>= lift . readTVarIO
  setEnv e' = gets repl >>= lift . atomically . flip writeTVar e'

  -- TODO what does this mean? perhaps it coiuld be used to load some
  -- virtual files edited by user in their env?
  load = const (pure (Right ""))

instance (MonadThrow m) => MonadThrow (Net m) where
  throwM = lift . throwM

instance (MonadCatch m) => MonadCatch (Net m) where
  Net m `catch` f = Net $ m `catch` \e -> runNet (f e)

-- instance TypeChecker (Net IO) where
--     emit (CheckD ev) = gets logger >>= \env -> liftIO $ logInfo env ev
--     emit _ = pure ()

wsReceive :: WS.Connection -> IO In
wsReceive cnx =
  fromRight (In "") . eitherDecode <$> WS.receiveData cnx

wsSend :: Out -> WS.Connection -> IO ()
wsSend out cnx =
  WS.sendBinaryData cnx (encode out)

safeDecodeUtf8 :: ByteString -> Text
safeDecodeUtf8 = decodeUtf8With lenientDecode

clientHandler :: LoggerEnv IO -> TVar (Map ByteString (TVar REPLEnv)) -> WS.PendingConnection -> IO ()
clientHandler loggerEnv envs cnx = do
  let path = WS.requestPath $ WS.pendingRequest cnx
  conn <- WS.acceptRequest cnx
  -- storageFile <- mkStorageFile path
  -- withStorage storageFile $ \fileStorage -> do
  replEnv <- findOrCreateREPL path
  WS.withPingThread conn 30 (pure ()) $
    evalStateT (runNet runREPL) (NetEnv conn loggerEnv replEnv)
  where
    findOrCreateREPL :: ByteString -> IO (TVar REPLEnv)
    findOrCreateREPL path = do
      (newEnv, reconnect) <- atomically $ do
        envMaps <- readTVar envs
        (newEnv, reconnect) <- case Map.lookup path envMaps of
          Nothing -> newTVar initialREPL >>= (\e -> pure (e, True))
          Just e -> pure (e, False)
        writeTVar envs (Map.insert path newEnv envMaps)
        pure (newEnv, reconnect)
      if reconnect
        then logInfo loggerEnv $ object ["action" .= ("ReconnectingREPL" :: Text), "path" .= safeDecodeUtf8 path]
        else logInfo loggerEnv $ object ["action" .= ("StartingREPL" :: Text), "path" .= safeDecodeUtf8 path]
      pure newEnv

-- mkStorageFile path = do
--   let f = takeFileName (unpack (safeDecodeUtf8 path))
--   dir <- getCurrentDirectory
--   pure $ defaultOptions {storageFilePath = dir </> f}

-- replay :: (H.Store IO store) => store -> TVar REPLEnv -> IO (TVar REPLEnv)
-- replay storage envVar = do
--   res <- H.load storage
--   case res of
--     LoadSucceed coms -> do
--       atomically $ do
--         env <- readTVar envVar
--         let newEnv = withInputs env coms
--         writeTVar envVar (purerREPL newEnv)
--       logInfo loggerEnv $
--         object
--           [ "action" .= ("Loaded" :: Text),
--             "numCommands" .= length coms
--           ]
--       pure envVar
--     err -> do
--       logInfo loggerEnv $ object ["action" .= ("ErrorLoading" :: Text), "error" .= pack (show err)]
--       pure envVar

runNetREPL :: LoggerEnv IO -> TVar (Map ByteString (TVar REPLEnv)) -> Application -> Application
runNetREPL loggerEnv envs = websocketsOr WS.defaultConnectionOptions (clientHandler loggerEnv envs)