module Builtins where

import Code (Builtin(..), Code, Datum(..), Instr(..), InterpFn)
import CodeLens
import qualified Compiler as Co
import Control.Exception (IOException, catch)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Except (runExceptT)
import Data.Functor.Identity (runIdentity)
import Data.List (intercalate)
import qualified Data.Map as M
import Data.Maybe (fromJust)
import Env (PrlgEnv(..), findAtom, findStruct, prlgError)
import Heap (Dereferenced(..), derefHeap, heapStruct, newHeapVars)
import qualified IR
import Interpreter (backtrack)
import Lens.Family2.State
import Load (processInput)
import qualified Operators as O
import System.Console.Haskeline (getInputChar, outputStr, outputStrLn)

bi = Builtin

continue = pure Nothing

showTerm itos heap = runIdentity . heapStruct atom struct hrec heap
  where
    atom (Atom a) = pure $ "'" ++ itos M.! a ++ "'"
    atom (Number n) = pure (show n)
    atom VoidRef = pure "_"
    struct (Struct (IR.Id h _)) args =
      pure $ "'" ++ itos M.! h ++ "'(" ++ intercalate "," args ++ ")"
    hrec (HeapRef hr) ref =
      pure $
      (if hr == ref
         then "_X"
         else "_Rec") ++
      show hr

printLocals :: InterpFn
printLocals = do
  scope <- use (cur . gvar)
  heap <- use (cur . heap)
  IR.StrTable _ _ itos <- use strtable
  flip traverse (M.assocs scope) $ \(local, ref) ->
    lift . outputStrLn $ "_Local" ++ show local ++ " = " ++
    showTerm itos heap ref
  continue

promptRetry :: InterpFn
promptRetry = do
  last <- cho `uses` null
  if last
    then continue
    else promptRetry'

promptRetry' :: InterpFn
promptRetry' = do
  x <- lift $ getInputChar "? "
  case x of
    Just ';' -> backtrack
    _ -> continue

withArgs :: [Int] -> ([Int] -> InterpFn) -> InterpFn
withArgs as f = do
  scope <- use (cur . hvar)
  if all (`M.member` scope) as
    then f $ map (scope M.!) as
    else prlgError "arguments not bound"

write' :: InterpFn -> InterpFn
write' c =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    IR.StrTable _ _ itos <- use strtable
    lift . outputStr $ showTerm itos heap arg
    c --this now allows error fallthrough but we might like EitherT

write = write' continue

nl :: InterpFn
nl = do
  lift $ outputStrLn ""
  continue

writeln :: InterpFn
writeln = write' nl

assertFact :: (Code -> IR.Id -> PrlgEnv ()) -> InterpFn
assertFact addClause =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    case Co.compileGoal . Co.squashVars <$>
         Co.heapStructPrlgInt Nothing heap arg of
      Just (U (Struct s):head) -> do
        addClause (head ++ [Done]) s
        continue
      _ -> prlgError "assert fact failure"

assertRule :: (Code -> IR.Id -> PrlgEnv ()) -> InterpFn
assertRule addClause =
  withArgs [0, 1] $ \args -> do
    scope <- use (cur . hvar)
    heap <- use (cur . heap)
    comma <- findAtom ","
    semi <- findAtom ";"
    cut <- findAtom "!"
    case Co.squashVars . IR.CallI 0 <$>
         traverse (Co.heapStructPrlgInt Nothing heap) args of
      Just (IR.CallI 0 [hs, gs]) ->
        let (U (Struct s):cs) =
              Co.compileGoal hs ++ Co.compileGoals comma semi cut gs
         in do addClause cs s
               continue
      _ -> prlgError "assert clause failure"

retractall :: InterpFn
retractall =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    case derefHeap heap arg of
      BoundRef _ (Atom a) ->
        dropProcedure (IR.Id {IR.arity = 0, IR.str = a}) >> continue
      BoundRef _ (Struct id) -> dropProcedure id >> continue
      _ -> prlgError "retractall needs a struct"

exec' :: (Code -> Code) -> InterpFn
exec' fgol =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    case Co.squashVars <$> Co.heapStructPrlgInt Nothing heap arg of
      Just gs -> do
        comma <- findAtom ","
        semi <- findAtom ";"
        cut <- findAtom "!"
        zoom cur $ do
          hvar .= M.empty
          hed .= Co.compileGoals comma semi cut gs
          gol %= fgol
        continue
      _ -> prlgError "bad goal"

call :: InterpFn
call = exec' id

exec :: InterpFn
exec = exec' (const [Done])

{- struct assembly/disassembly -}
struct :: InterpFn
struct = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  case derefHeap heap <$> scope M.!? 0 of
    Just (BoundRef addr (Struct IR.Id {IR.arity = arity, IR.str = str})) ->
      structUnify arity str
    _ -> structAssemble

heapListLength listAtom heap ref = go 0 ref ref (id, fromJust . step)
  where
    nil r
      | BoundRef _ str <- derefHeap heap r = str == Atom listAtom
      | otherwise = False
    step r
      | BoundRef addr (Struct IR.Id {IR.arity = 2, IR.str = listAtom'}) <-
         derefHeap heap r
      , listAtom == listAtom' = Just (addr + 2)
      | otherwise = Nothing
    go n fast slow (f1, f2)
      | nil fast = Just n
      | Just fast' <- step fast =
        if slow == fast'
          then Nothing
          else go (n + 1) fast' (f1 slow) (f2, f1)
      | otherwise = Nothing

structAssemble :: InterpFn
structAssemble = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  case derefHeap heap <$> scope M.!? 1 of
    Just (BoundRef addr (Atom str)) -> do
      listAtom <- findAtom "[]"
      case scope M.!? 2 >>= heapListLength listAtom heap of
        Just arity -> structUnify arity str
        _ -> prlgError "struct arity unknown"
    _ -> prlgError "struct id unknown"

structUnify arity str = do
  h <- use (cur . heap)
  scope <- use (cur . hvar)
  listAtom <- findAtom "[]"
  pvars <- newHeapVars arity
  let hcode = map U $ maybe VoidRef HeapRef . (scope M.!?) <$> [0 .. 2]
      structData =
        Struct IR.Id {IR.arity = arity, IR.str = str} : map HeapRef pvars
      paramsData =
        concatMap
          (\pv -> [Struct IR.Id {IR.arity = 2, IR.str = listAtom}, HeapRef pv])
          pvars ++
        [Atom listAtom]
      gcode = map U $ structData ++ [Atom str] ++ paramsData
  zoom cur $ do
    gol %= (gcode ++)
    hed %= (hcode ++)
  continue

{- terms -}
var :: InterpFn
var = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  case derefHeap heap <$> scope M.!? 0 of
    Nothing -> continue
    Just (FreeRef _) -> continue
    _ -> backtrack

sameTerm :: InterpFn
sameTerm = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  case map (fmap (derefHeap heap) . (scope M.!?)) [0, 1] of
    [Just a, Just b]
      | a == b -> continue
    _ -> backtrack

currentPredicate :: InterpFn
currentPredicate =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    ds <- use defs
    case derefHeap heap arg of
      BoundRef _ (Struct s) ->
        if s `M.member` ds
          then continue
          else backtrack
      _ -> prlgError "not a predicate"

{- operator management -}
op :: InterpFn
op =
  withArgs [0, 1, 2] $ \args -> do
    heap <- use (cur . heap)
    IR.StrTable _ _ itos <- use strtable
    case map (derefHeap heap) args of
      [BoundRef _ (Number prio), BoundRef _ (Atom fixityAtom), BoundRef _ (Atom opatom)]
        | Just op <-
           (,) <$> itos M.!? opatom <*>
           (O.Op prio <$> ((itos M.!? fixityAtom) >>= O.fixity)) -> do
          ops %= (op :)
          continue
      _ -> prlgError "bad op spec"

deop :: InterpFn
deop =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    IR.StrTable _ _ itos <- use strtable
    case derefHeap heap arg of
      BoundRef _ (Atom opatom)
        | Just op <- itos M.!? opatom -> do
          ops %= filter ((/= op) . fst)
          continue
      _ -> prlgError "bad op spec"

stashOps :: InterpFn
stashOps = do
  currentOps <- use ops
  opstash %= (currentOps :)
  continue

popOps :: InterpFn
popOps = do
  currentOps <- use opstash
  case currentOps of
    [] -> prlgError "no ops stashed"
    (ops':opss) -> do
      ops .= ops'
      opstash .= opss
      continue

{- expansion environment -}
stashExpansions :: InterpFn
stashExpansions = do
  ds <- use defs
  les <- findStruct "load_expansion" 2
  qes <- findStruct "query_expansion" 2
  let [le, qe] = map (ds M.!?) [les, qes]
  macrostash %= ((le, qe) :)
  continue

popExpansions :: InterpFn
popExpansions = do
  currentMacros <- use macrostash
  les <- findStruct "load_expansion" 2
  qes <- findStruct "query_expansion" 2
  case currentMacros of
    [] -> prlgError "no expansions stashed"
    ((le, qe):stash') -> do
      defs %= M.alter (const le) les . M.alter (const qe) qes
      macrostash .= stash'
      continue

{- adding the builtins -}
addOp :: (String, O.Op) -> PrlgEnv ()
addOp op = ops %= (op :)

modDef :: ([Code] -> Maybe [Code]) -> IR.Id -> PrlgEnv ()
modDef fn struct = defs %= M.alter (maybe (fn []) fn) struct

addClauseA :: Code -> IR.Id -> PrlgEnv ()
addClauseA code = modDef $ Just . (code :)

addClauseZ :: Code -> IR.Id -> PrlgEnv ()
addClauseZ code = modDef $ Just . (++ [code])

addProcedure :: [Code] -> IR.Id -> PrlgEnv ()
addProcedure heads = modDef $ Just . const heads

dropProcedure :: IR.Id -> PrlgEnv ()
dropProcedure = modDef $ const Nothing

addProc :: [Code] -> String -> Int -> PrlgEnv ()
addProc c n a = findStruct n a >>= addProcedure c

addBi :: InterpFn -> String -> Int -> PrlgEnv ()
addBi b n a =
  addProc [[U (LocalRef $ r - 1) | r <- [1 .. a]] ++ [Invoke $ bi b]] n a

{- loading code -}
load :: Bool -> InterpFn
load queryMode =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    IR.StrTable _ _ itos <- use strtable --TODO the argument here should preferably be a string, right?
    case derefHeap heap arg of
      BoundRef _ (Atom a) -> do
        let fn = itos M.! a
        src' <- liftIO $ catch (Right <$> readFile fn) (pure . Left)
        case src' of
          Right src -> do
            res <- runExceptT $ processInput fn queryMode src
            case res of
              Right _ -> continue
              Left e -> prlgError $ "loading from '" ++ fn ++ "': " ++ e
          Left e -> prlgError $ show (e :: IOException)
      _ -> prlgError "load needs an atom"

{- actual prlgude -}
addPrelude :: PrlgEnv ()
addPrelude = do
  pure undefined
  {- primitives -}
  addBi (pure Nothing) "true" 0
  addBi backtrack "fail" 0
  addOp $ O.xfx "=" 700
  addProc [[U (LocalRef 0), U (LocalRef 0), Done]] "=" 2
  {- clauses -}
  addOp $ O.xfy "," 1000
  addOp $ O.xfy ";" 1100
  addOp $ O.xfx ":-" 1200
  addOp $ O.fx ":-" 1200
  horn1 <- findStruct ":-" 1
  horn2 <- findStruct ":-" 2
  let assertCode ac =
        [ [ U (Struct horn2)
          , U (LocalRef 0)
          , U (LocalRef 1)
          , Cut
          , Invoke . bi $ assertRule ac
          ]
        , [U (Struct horn1), U (LocalRef 0), Cut, Invoke $ bi exec]
        , [U (LocalRef 0), Invoke . bi $ assertFact ac]
        ]
   in do addProc (assertCode addClauseA) "asserta" 1
         addProc (assertCode addClauseZ) "assertz" 1
         addProc (assertCode addClauseZ) "assert" 1
  addBi retractall "retractall" 1
  addBi call "call" 1
  {- terms -}
  addBi struct "struct" 3
  addBi var "var" 1
  addBi sameTerm "same_term" 2
  addBi currentPredicate "current_predicate" 1
  {- code loading -}
  addBi (load False) "load" 1
  addBi (load True) "source" 1
  {- operators -}
  addBi op "op" 3
  addBi deop "deop" 1
  addBi stashOps "stash_operators" 0
  addBi popOps "pop_operators" 0
  {- macroenvironment -}
  addBi stashExpansions "stash_expansions" 0
  addBi popExpansions "pop_expansions" 0
  let expandCode q = do
        s <- findStruct (q ++ "_expansion") 2
        cp <- findStruct "current_predicate" 1
        addProc
          [ [ U (LocalRef 0)
            , U (LocalRef 1)
            , U (Struct cp) -- current_predicate(expand_something(_,_)),
            , U (Struct s)
            , U VoidRef
            , U VoidRef
            , U (Struct s) -- expand_something(Arg1, Arg2).
            , U (LocalRef 0)
            , U (LocalRef 1)
            , Cut -- TODO  check that the cut works here; this was the whole reason why we migrated off vienna
            , Done
            ]
          , [U (LocalRef 0), U (LocalRef 0), Done]
          ]
          ("expand_" ++ q)
          2
  expandCode "load"
  expandCode "query"
  {- query tools -}
  addBi printLocals "print_locals" 0
  addBi promptRetry' "prompt_retry" 0
  addBi (printLocals >> promptRetry) "query" 0
  {- IO -}
  addBi write "write" 1
  addBi writeln "writeln" 1
  addBi nl "nl" 0
  {- debug -}
  addBi (use id >>= liftIO . print >> pure Nothing) "interpreter_trace" 0