tricu/src/REPL.hs

module REPL where

import Eval
import FileEval
import Lexer
import Parser
import Research

import Control.Exception         (IOException, SomeException, catch
                                 , displayException)
import Control.Monad             (forM_)
import Control.Monad.Catch       (handle, MonadCatch)
import Control.Monad.IO.Class    (liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe (MaybeT(..), runMaybeT)
import Data.Char                 (isSpace, isUpper)
import Data.List                 (dropWhile, dropWhileEnd, isPrefixOf)
import Data.Version (showVersion)
import Paths_tricu (version)
import System.Console.Haskeline

import qualified Data.Map as Map
import qualified Data.Text as T
import qualified Data.Text.IO as T

repl :: Env -> IO ()
repl env = runInputT settings (withInterrupt (loop env Decode))
  where
    settings :: Settings IO
    settings = Settings
      { complete = completeWord Nothing " \t" completeCommands
      , historyFile = Just "~/.local/state/tricu/history"
      , autoAddHistory = True
      }

    completeCommands :: String -> IO [Completion]
    completeCommands str = return $ map simpleCompletion $
      filter (str `isPrefixOf`) commands
      where
        commands = [ "!exit"
                   , "!output"
                   , "!definitions"
                   , "!import"
                   , "!clear"
                   , "!save"
                   , "!reset"
                   , "!version"
                   ]

    loop :: Env -> EvaluatedForm -> InputT IO ()
    loop env form = handle (interruptHandler env form) $ do
      minput <- getInputLine "tricu < "
      case minput of
        Nothing -> outputStrLn "Exiting tricu"
        Just s
          | strip s == "" -> loop env form
          | strip s == "!exit" -> outputStrLn "Exiting tricu"
          | strip s == "!clear" -> do
              liftIO $ putStr "\ESC[2J\ESC[H"
              loop env form
          | strip s == "!reset" -> do
              outputStrLn "Environment reset to initial state"
              loop Map.empty form
          | strip s == "!version" -> do
              outputStrLn $ "tricu version " ++ showVersion version
              loop env form
          | "!save" `isPrefixOf` strip s -> handleSave env form
          | strip s == "!output" -> handleOutput env form
          | strip s == "!definitions" -> do
              let defs = Map.keys $ Map.delete "!result" env
              if null defs
                then outputStrLn "No definitions discovered."
                else do
                  outputStrLn "Available definitions:"
                  mapM_ outputStrLn defs
              loop env form
          | "!import" `isPrefixOf` strip s -> handleImport env form
          | take 2 s == "--" -> loop env form
          | otherwise -> do
              newEnv <- liftIO $ processInput env s form `catch` errorHandler env
              loop newEnv form

    handleOutput :: Env -> EvaluatedForm -> InputT IO ()
    handleOutput env currentForm = do
      let formats = [Decode, TreeCalculus, FSL, AST, Ternary, Ascii]
      outputStrLn "Available output formats:"
      mapM_ (\(i, f) -> outputStrLn $ show i ++ ". " ++ show f)
        (zip [1..] formats)

      result <- runMaybeT $ do
        input <- MaybeT $ getInputLine "Select output format (1-6) < "
        case reads input of
          [(n, "")] | n >= 1 && n <= 6 ->
            return $ formats !! (n-1)
          _ -> MaybeT $ return Nothing

      case result of
        Nothing -> do
          outputStrLn "Invalid selection. Keeping current output format."
          loop env currentForm
        Just newForm -> do
          outputStrLn $ "Output format changed to: " ++ show newForm
          loop env newForm

    handleImport :: Env -> EvaluatedForm -> InputT IO ()
    handleImport env form = do
      res <- runMaybeT $ do
        let fset = setComplete completeFilename defaultSettings
        path <- MaybeT $ runInputT fset $
          getInputLineWithInitial "File path to load < " ("", "")

        text <- MaybeT $ liftIO $ handle (\e -> do
            putStrLn $ "Error reading file: " ++ displayException (e :: IOException)
            return Nothing
          ) $ Just <$> readFile (strip path)

        case parseProgram (lexTricu text) of
          Left err -> do
            lift $ outputStrLn $ "Parse error: " ++ handleParseError err
            MaybeT $ return Nothing
          Right ast -> do
            ns <- MaybeT $ runInputT defaultSettings $
              getInputLineWithInitial "Namespace (or !Local for no namespace) < " ("", "")

            let name = strip ns
            if (name /= "!Local" && (null name || not (isUpper (head name)))) then do
                lift $ outputStrLn "Namespace must start with an uppercase letter"
                MaybeT $ return Nothing
            else do
                prog <- liftIO $ preprocessFile (strip path)
                let code = case name of
                      "!Local" -> prog
                      _ -> nsDefinitions name prog
                    env' = evalTricu env code
                return env'
      case res of
        Nothing -> do
          outputStrLn "Import cancelled"
          loop env form
        Just env' ->
          loop (Map.delete "!result" env') form

    interruptHandler :: Env -> EvaluatedForm -> Interrupt -> InputT IO ()
    interruptHandler env form _ = do
      outputStrLn "Interrupted with CTRL+C\n\
                  \You can use the !exit command or CTRL+D to exit"
      loop env form

    processInput :: Env -> String -> EvaluatedForm -> IO Env
    processInput env input form = do
      let asts   = parseTricu input
          newEnv = evalTricu env asts
      case Map.lookup "!result" newEnv of
        Just r -> do
          putStrLn $ "tricu > " ++ formatT form r
        Nothing -> pure ()
      return newEnv

    errorHandler :: Env -> SomeException -> IO (Env)
    errorHandler env e = do
      putStrLn $ "Error: " ++ show e
      return env

    strip :: String -> String
    strip = dropWhileEnd isSpace . dropWhile isSpace

    handleSave :: Env -> EvaluatedForm -> InputT IO ()
    handleSave env form = do
      let fset = setComplete completeFilename defaultSettings
      path <- runInputT fset $
        getInputLineWithInitial "File to save < " ("", "")

      case path of
        Nothing -> do
          outputStrLn "Save cancelled"
          loop env form
        Just p -> do
          let definitions = Map.toList $ Map.delete "!result" env
              filepath = strip p

          outputStrLn "Starting save..."
          liftIO $ writeFile filepath ""
          outputStrLn "File created..."
          forM_ definitions $ \(name, value) -> do
            let content = name ++ " = " ++ formatT TreeCalculus value ++ "\n"
            outputStrLn $ "Writing definition: " ++ name ++ " with length " ++ show (length content)
            liftIO $ appendFile filepath content
          outputStrLn $ "Saved " ++ show (length definitions) ++ " definitions to " ++ p

          loop env form