module Builtins where

import Code
  ( Builtin(..)
  , Code
  , Datum(..)
  , Dereferenced(..)
  , Instr(..)
  , InterpFn
  , InterpFn
  , derefHeap
  , heapStruct
  , newHeapVars
  )
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 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 ","
    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.seqGoals (Co.compileGoals comma 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 ","
        cut <- findAtom "!"
        zoom cur $ do
          hvar .= M.empty
          hed .= Co.seqGoals (Co.compileGoals comma 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 "[]"
  let hcode = map U $ maybe VoidRef HeapRef . (scope M.!?) <$> [0 .. 2]
      (h', pvars) = newHeapVars arity h
      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
    heap .= h'
    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.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