module Frontend where

import Builtins
import Code (Interp(..))
import qualified Compiler as C
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Except (except, runExceptT)
import Control.Monad.Trans.State.Lazy (evalStateT, gets)
import qualified Data.Map as M
import Data.Traversable (for)
import Env (PrlgEnv, findAtom, findStruct, withStrTable)
import qualified IR
import qualified Interpreter as I
import qualified Parser as P
import System.Console.Haskeline
import qualified Text.Megaparsec as MP
import qualified Text.Pretty.Simple as Ppr

ppr :: Show a => a -> PrlgEnv ()
ppr x =
  liftIO $
  Ppr.pPrintOpt
    Ppr.CheckColorTty
    Ppr.defaultOutputOptionsDarkBg
      { Ppr.outputOptionsCompactParens = True
      , Ppr.outputOptionsIndentAmount = 2
      , Ppr.outputOptionsPageWidth = 80
      }
    x

left f = either (Left . f) Right

tokenize = left MP.errorBundlePretty . MP.parse P.lexPrlg "<stdin>"

parse = left MP.errorBundlePretty . MP.parse P.parsePrlg "<stdin>"

shunt ops =
  left (\err -> "operator resolution: " ++ err ++ "\n") . P.shuntPrlg ops

intern prlgs = do
  prlgi <- withStrTable $ \st -> IR.internPrlg st prlgs
  underscore <- findAtom "_"
  list <- findAtom "[]"
  withStrTable $ \st ->
    ( st
    , C.squashVars $ C.variablizePrlg underscore st $ C.desugarPrlg list prlgi)

compile prlgv = do
  comma <- findAtom ","
  cut <- findAtom "!"
  return $ C.seqGoals (C.compileGoals comma cut prlgv)

-- the signature of this is too ugly to include here
handleError m = do
  res <- runExceptT m
  case res of
    Left err -> lift $ outputStr err
    _ -> pure ()

processInput precompileHook good bad input =
  (True <$) . handleError $ do
    asts <- except $ tokenize input >>= parse
    ops <- lift $ gets ops
    for asts $ \ast' -> do
      ast <- except $ shunt ops ast'
      code <- lift $ intern ast >>= precompileHook >>= compile
      res <- lift (I.prove code) >>= except . left (++ "\n")
      lift . lift . outputStrLn $
        if res
          then good
          else bad

expansion noexpand expander output x = do
  es <- findStruct expander 2
  o <- findAtom output
  comma <- findAtom ","
  expand <- gets (M.member es . defs)
  pure $
    if expand
      then IR.CallI
             comma
             [ IR.CallI (IR.str es) [x, IR.VarI (-1) 0]
             , IR.CallI o [IR.VarI (-1) 0]
             ]
      else noexpand o x

query =
  processInput
    (expansion (\_ -> id) "query_expansion" "call")
    "yes."
    "no proof."

assert =
  processInput
    (expansion (\o x -> IR.CallI o [x]) "load_expansion" "assert")
    "ok."
    "rejected."

interpreterStart :: PrlgEnv ()
interpreterStart = do
  addPrelude
  interpreterLoop True

interpreterLoop :: Bool -> PrlgEnv ()
interpreterLoop queryMode = do
  minput <-
    lift $
    getInputLine
      (if queryMode
         then "prlg? "
         else "prlg |- ")
  case minput of
    Nothing -> return ()
    Just "." -> interpreterLoop (not queryMode)
    Just input -> do
      continue <-
        (if queryMode
           then query
           else assert)
          input
      if continue
        then interpreterLoop queryMode
        else return ()

interpreter :: InputT IO ()
interpreter =
  evalStateT
    interpreterStart
    (Interp
       { defs = M.empty
       , ops = []
       , opstash = []
       , strtable = IR.emptystrtable
       , cur = error "no cur"
       , cho = []
       })