module Builtins where

import Code
  ( Builtin(..)
  , Cho(..)
  , Datum(..)
  , Dereferenced(..)
  , Heap(..)
  , Instr(..)
  , Interp(..)
  , InterpFn
  , InterpFn
  , derefHeap
  , heapStruct
  , newHeapVars
  )
import qualified Compiler as Co
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Class (lift)
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 qualified Operators as O
import System.Console.Haskeline (getInputChar, outputStr, outputStrLn)

bi = Builtin

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
  return Nothing

promptRetry :: InterpFn
promptRetry = do
  last <- gets (null . cho)
  if last
    then return Nothing
    else promptRetry'

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

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' $ return Nothing

nl :: InterpFn
nl = do
  lift $ outputStrLn ""
  return Nothing

writeln :: InterpFn
writeln = write' nl

assertFact :: InterpFn
assertFact =
  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 s $ head ++ [NoGoal]
        return Nothing
      _ -> prlgError "assert fact failure"

assertClause :: InterpFn
assertClause =
  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 s cs
               return Nothing
      _ -> 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})
      BoundRef _ (Struct id) -> dropProcedure id
      _ -> 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
                    }
              }
          return Nothing
    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"

{- struct building -}
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}}
  return Nothing

{- adding the builtins -}
addOp op = modify $ \s -> s {ops = op : ops s}

addClause struct code =
  modify $ \s ->
    s {defs = M.alter (Just . maybe [code] (\hs -> code : hs)) struct $ defs s}

addProcedure struct heads =
  modify $ \s -> s {defs = M.insert struct heads $ defs s}

dropProcedure struct = do
  d <- gets defs
  if struct `M.member` d
    then do
      modify $ \s -> s {defs = M.delete struct d}
      return Nothing
    else prlgError "no such definition" -- this should backtrack?

addProc n a c = do
  sym <- findStruct n a
  addProcedure sym c

addBi n i b =
  addProc n i [[U (LocalRef $ r - 1) | r <- [1 .. i]] ++ [Invoke $ bi b]]

addPrelude :: PrlgEnv ()
addPrelude = do
  pure undefined
  {- primitives -}
  addBi "true" 0 (pure Nothing)
  addBi "fail" 0 backtrack
  addOp $ O.xfx "=" 700
  addProc "=" 2 [[U (LocalRef 0), U (LocalRef 0), NoGoal]]
  {- clauses -}
  addOp $ O.xfy "," 1000
  addOp $ O.xfx ":-" 1200
  horn2 <- findStruct ":-" 2
  --addOp $ O.fx ":-" 1200
  addProc
    "assert"
    1
    [ [ U (Struct horn2)
      , U (LocalRef 0)
      , U (LocalRef 1)
      , Cut
      , Invoke (bi assertClause)
      ]
    , [U (LocalRef 0), Invoke (bi assertFact)]
    ]
  addBi "retractall" 1 retractall
  addBi "call" 1 call
  addBi "struct" 3 struct
  {- query tools -}
  addBi "print_locals" 0 printLocals
  addBi "prompt_retry" 0 promptRetry'
  addBi "query" 0 (printLocals >> promptRetry)
  {- IO -}
  addBi "write" 1 write
  addBi "writeln" 1 writeln
  addBi "nl" 0 nl
  {- debug -}
  addBi "interpreter_trace" 0 (get >>= liftIO . print >> pure Nothing)