module Builtins where

import Code
  ( Builtin(..)
  , Cho(..)
  , Code
  , Datum(..)
  , Dereferenced(..)
  , Heap(..)
  , Instr(..)
  , Interp(..)
  , InterpFn
  , InterpFn
  , derefHeap
  , heapStruct
  , newHeapVars
  )
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 Control.Monad.Trans.State.Lazy (get, gets, modify)
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 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 <- gets (gvar . cur)
  heap <- gets (heap . cur)
  IR.StrTable _ _ itos <- gets strtable
  flip traverse (M.assocs scope) $ \(local, ref) ->
    lift . outputStrLn $
    "_Local" ++ show local ++ " = " ++ showTerm itos heap ref
  continue

promptRetry :: InterpFn
promptRetry = do
  last <- gets (null . cho)
  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 <- gets (hvar . cur)
  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 <- gets (heap . cur)
    IR.StrTable _ _ itos <- gets 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 <- gets (heap . cur)
    case Co.compileGoal . Co.squashVars <$>
         Co.heapStructPrlgInt Nothing heap arg of
      Just (U (Struct s):head) -> do
        addClause (head ++ [NoGoal]) s
        continue
      _ -> prlgError "assert fact failure"

assertRule :: (Code -> IR.Id -> PrlgEnv ()) -> InterpFn
assertRule addClause =
  withArgs [0, 1] $ \args -> do
    scope <- gets (hvar . cur)
    heap <- gets (heap . cur)
    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 <- gets (heap . cur)
    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"

call' :: InterpFn
call' =
  withArgs [0] $ \[arg] -> do
    heap@(Heap _ hmap) <- gets (heap . cur)
    let exec base struct arity = do
          cur <- gets cur
          modify $ \s ->
            s
              { cur =
                  cur
                    { gol =
                        [Call, Goal, U struct] ++
                        [U $ hmap M.! (base + i) | i <- [1 .. arity]] ++ gol cur
                    }
              }
          continue
    case derefHeap heap arg of
      BoundRef addr struct@(Struct IR.Id {IR.arity = arity}) ->
        exec addr struct arity
      BoundRef addr (Atom a) ->
        exec addr (Struct IR.Id {IR.arity = 0, IR.str = a}) 0
      _ -> prlgError "not callable"

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

call :: InterpFn
call = exec' id

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

{- struct assembly/disassembly -}
struct :: InterpFn
struct = do
  heap <- gets (heap . cur)
  scope <- gets (hvar . cur)
  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 <- gets (heap . cur)
  scope <- gets (hvar . cur)
  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
  cur <- gets cur
  let h = heap cur
      scope = hvar cur
  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
  modify $ \s ->
    s {cur = cur {heap = h', gol = gcode ++ gol cur, hed = hcode ++ hed cur}}
  continue

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

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

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

stashOps :: InterpFn
stashOps = do
  currentOps <- gets ops
  modify $ \s -> s {opstash = currentOps : opstash s}
  continue

popOps :: InterpFn
popOps = do
  currentOps <- gets opstash
  case currentOps of
    [] -> prlgError "no op stash to pop"
    (ops':opss) -> do
      modify $ \s -> s {ops = ops', opstash = opss}
      continue

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

modDef :: ([Code] -> Maybe [Code]) -> IR.Id -> PrlgEnv ()
modDef fn struct =
  modify $ \s -> s {defs = M.alter (maybe (fn []) fn) struct $ defs s}

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 <- gets (heap . cur)
    IR.StrTable _ _ itos <- gets 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), NoGoal]] "=" 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 same_term "same_term" 2
  {- code loading -}
  addBi (load False) "load" 1
  addBi (load True) "source" 1
  {- operators -}
  addBi op "op" 3
  addBi stashOps "stash_operators" 0
  addBi popOps "pop_operators" 0
  {- 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 (get >>= liftIO . print >> pure Nothing) "interpreter_trace" 0