module Builtins where

import Paths_prlg

import Code (Builtin(..), Code, Datum(..), Id(..), Instr(..), InterpFn)
import CodeLens
import qualified Compiler as Co
import Constant
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.Bits
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, writeHeap)
import IR (PrlgInt(..), StrTable(..))
import Interpreter (backtrack)
import Lens.Micro.Mtl
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 (C (Atom a)) = pure $ "'" ++ itos M.! a ++ "'"
    atom (C (Number n)) = pure (show n)
    atom (C (Str str)) = pure (show str)
    atom VoidRef = pure "_"
    struct (Struct (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)
  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 <- null <$> use 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 <- 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)
    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 -> Id -> PrlgEnv ()) -> InterpFn
assertFact addClause =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    case Co.heapStructPrlgInt Nothing heap arg of
      Just x -> do
        case Co.compileGoal $ Co.squashVars x of
          Right (U (Struct s):head) -> do
            addClause (head ++ [Done]) s
            continue
          Left err -> prlgError err
      _ -> prlgError "assert fact failure"

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

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

call :: InterpFn
call =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    case derefHeap heap arg of
      BoundRef _ (C (Atom a)) -> do
        cur . hed .= [U (Struct (Id {str = a, arity = 0})), Done]
        continue
      BoundRef addr s@(Struct Id {arity = arity}) -> do
        cur . hed .= [U s] ++ [U (HeapRef $ addr + i) | i <- [1 .. arity]] ++
          [Done]
        continue
      _ -> prlgError "bad call"

stop :: InterpFn
stop =
  withArgs [0] $ \[arg] -> do
    StrTable _ _ itos <- use strtable
    heap <- use (cur . heap)
    prlgError $ "stop: " ++ showTerm itos heap arg

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

heapListLength listAtom heap ref = go 0 ref ref (id, fromJust . step)
  where
    nil r
      | BoundRef _ str <- derefHeap heap r = str == C (Atom listAtom)
      | otherwise = False
    step r
      | BoundRef addr (Struct Id {arity = 2, 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 (C (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 Id {arity = arity, str = str} : map HeapRef pvars
      paramsData =
        concatMap
          (\pv -> [Struct Id {arity = 2, str = listAtom}, HeapRef pv])
          pvars ++
        [C $ Atom listAtom]
      gcode = map U $ structData ++ [C $ Atom str] ++ paramsData
  zoom cur $ do
    gol %= (gcode ++)
    hed %= (hcode ++)
    unis += 3
  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

atom :: InterpFn
atom = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  case derefHeap heap <$> scope M.!? 0 of
    Just (BoundRef _ (C (Atom _))) -> continue
    _ -> backtrack

number :: InterpFn
number = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  case derefHeap heap <$> scope M.!? 0 of
    Just (BoundRef _ (C (Number _))) -> continue
    _ -> backtrack

string :: InterpFn
string = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  --TODO unify with number/var/...
  case derefHeap heap <$> scope M.!? 0 of
    Just (BoundRef _ (C (Str _))) -> 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)
    StrTable _ _ itos <- use strtable
    case map (derefHeap heap) args of
      [BoundRef _ (C (Number prio)), BoundRef _ (C (Atom fixityAtom)), BoundRef _ (C (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)
    StrTable _ _ itos <- use strtable
    case derefHeap heap arg of
      BoundRef _ (C (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

{- integer arithmetics -}
intBinary :: (Int -> Int -> Int) -> InterpFn
intBinary op =
  withArgs [0, 1] $ \[arg1, arg2] -> do
    heap <- use (cur . heap)
    case derefHeap heap <$> [arg1, arg2] of
      [BoundRef _ (C (Number n1)), BoundRef _ (C (Number n2))] ->
        putInt (op n1 n2) 2
      _ -> prlgError "int binary needs numbers"

intBinPred :: (Int -> Int -> Bool) -> InterpFn
intBinPred op =
  withArgs [0, 1] $ \args -> do
    heap <- use (cur . heap)
    case derefHeap heap <$> args of
      [BoundRef _ (C (Number n1)), BoundRef _ (C (Number n2))] ->
        if op n1 n2
          then continue
          else backtrack
      _ -> prlgError "int binary pred needs numbers"

intUnary :: (Int -> Int) -> InterpFn
intUnary op =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    case derefHeap heap arg of
      BoundRef _ (C (Number n)) -> putInt (op n) 1
      _ -> prlgError "int unary needs number"

intUnPred :: (Int -> Bool) -> InterpFn
intUnPred op =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    case derefHeap heap arg of
      BoundRef _ (C (Number n)) ->
        if op n
          then continue
          else backtrack
      _ -> prlgError "int binary pred needs numbers"

putInt val sc = do
  heap <- use (cur . heap)
  scope <- use (cur . hvar)
  case scope M.!? sc of
    Nothing -> continue
    Just a ->
      case derefHeap heap a of
        BoundRef _ (C (Number val'))
          | val == val' -> continue
        FreeRef a' -> writeHeap a' (C (Number val)) >> continue
        _ -> backtrack

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

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

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

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

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

dropProcedure :: 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 -}
doLoad :: Bool -> String -> InterpFn
doLoad queryMode fn = do
  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)

load :: Bool -> InterpFn
load queryMode =
  withArgs [0] $ \[arg] -> do
    heap <- use (cur . heap)
    StrTable _ _ itos <- use strtable --TODO the argument here should preferably be a string, right?
    case derefHeap heap arg of
      BoundRef _ (C (Atom a)) -> do
        let fn = itos M.! a
        doLoad queryMode (itos M.! a)
      _ -> prlgError "load needs an atom"

{- actual prlgude -}
addPrelude :: PrlgEnv ()
addPrelude = do
  pure undefined
  {- absolute primitives -}
  addProc [[Done]] "true" 0
  addBi backtrack "fail" 0
  addBi stop "stop" 1
  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
  do [horn1, horn2] <- traverse (findStruct ":-") [1, 2]
     doCall <- U . Struct <$> findStruct "call" 1
     let assertCode ac =
           [ [ U (Struct horn2)
             , U (LocalRef 0)
             , U (LocalRef 1)
             , Cut
             , Invoke . bi $ assertRule ac
             ]
           , [ U (Struct horn1)
             , U (LocalRef 0)
             , Cut
             , doCall
             , U (LocalRef 0)
             , Done
             ]
           , [U (LocalRef 0), Invoke . bi $ assertFact ac]
           ]
     addProc (assertCode addClauseA) "asserta" 1
     addProc (assertCode addClauseZ) "assertz" 1
     addProc (assertCode addClauseZ) "assert" 1
  addBi retractall "retractall" 1
  do [comma, semi] <- traverse (flip findStruct 2) [",", ";"]
     doCall <- U . Struct <$> findStruct "call" 1
     addProc
       [ [ U (Struct comma)
         , U (LocalRef 0)
         , U (LocalRef 1)
         , Cut
         , doCall
         , U (LocalRef 0)
         , doCall
         , U (LocalRef 1)
         , Done
         ]
       , [ U (Struct semi)
         , U (LocalRef 0)
         , U (LocalRef 1)
         , Cut
         , Choices [[doCall, U (LocalRef 0)], [doCall, U (LocalRef 1)]]
         , Done
         ]
       , [U (LocalRef 0), Invoke $ bi call]
       ]
       "call"
       1
  {- terms -}
  addBi struct "struct" 3
  addBi var "var" 1
  addBi atom "atom" 1
  addBi number "number" 1
  addBi string "string" 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
            , Done
            ]
          , [U (LocalRef 0), U (LocalRef 0), Done]
          ]
          ("expand_" ++ q)
          2
   in do expandCode "load"
         expandCode "query"
  {- int primops -}
  let add2IntOp nm op = addBi (intBinary op) ("int2_" ++ nm) 3
      add1IntOp nm op = addBi (intUnary op) ("int1_" ++ nm) 2
      add2IntPred nm op = addBi (intBinPred op) ("int2p_" ++ nm) 2
      add1IntPred nm op = addBi (intUnPred op) ("int1p_" ++ nm) 1
   in do add2IntOp "add" (+)
         add2IntOp "sub" (-)
         add1IntOp "neg" negate
         add1IntOp "abs" abs
         add2IntOp "mul" (*)
         add2IntOp "div" div
         add2IntOp "mod" mod
         add2IntOp "bitand" (.&.)
         add2IntOp "bitor" (.|.)
         add2IntOp "bitxor" xor
         add2IntOp "shl" shiftL
         add2IntOp "shr" shiftR
         add1IntPred "zero" (== 0)
         add2IntPred "eq" (==)
         add2IntPred "lt" (<)
         add2IntPred "leq" (<=)
         add2IntPred "neq" (/=)
  {- 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
  {- load the prelude file -}
  preludeFile <- liftIO $ getDataFileName "prelude.pl"
  doLoad False preludeFile
  pure ()