module Load where

import Code (Code, Id(..), PrlgEnv)
import CodeLens
import qualified Compiler as C
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Except (ExceptT, except)
import qualified Data.Map as M
import Env (findAtom, findStruct, withStrTable)
import qualified IR
import Lens.Micro.Mtl
import qualified Parser as P
import qualified Text.Megaparsec as MP

left :: (a -> b) -> Either a c -> Either b c
left f = either (Left . f) Right

tokenize :: String -> String -> Either String [P.Lexeme]
tokenize fn = left MP.errorBundlePretty . MP.parse P.lexPrlg fn

parse :: String -> [P.Lexeme] -> Either String [P.PAST]
parse fn = left MP.errorBundlePretty . MP.parse P.parsePrlg fn

shunt :: P.PAST -> ExceptT String PrlgEnv P.PrlgStr
shunt past = do
  ops <- lift $ use ops
  except . left (\err -> "operator resolution: " ++ err ++ "\n") $
    P.shuntPrlg ops past

intern :: P.PrlgStr -> PrlgEnv IR.PrlgInt
intern prlgs = do
  prlgi <- withStrTable $ \st -> IR.internPrlg st prlgs
  underscore <- findAtom "_"
  withStrTable $ \st ->
    (st, C.squashVars $ C.variablizePrlg underscore st prlgi)

compile :: IR.PrlgInt -> ExceptT String PrlgEnv Code
compile prlgv = do
  [comma, semi, cut] <- lift $ traverse findAtom [",", ";", "!"]
  except $ C.compileGoals comma semi cut prlgv

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

queryExpansion = expansion (\_ -> id) "expand_query" "call"

loadExpansion = expansion (\o x -> IR.CallI o [x]) "expand_load" "assert"

processInput :: String -> Bool -> String -> ExceptT String PrlgEnv ()
processInput fn queryMode input = do
  asts <- except $ tokenize fn input >>= parse fn
  lift $ cmdq .= [(queryMode, ast) | ast <- asts]