{-# LANGUAGE TupleSections #-}

module Code where

import Control.Monad.Trans.State.Lazy (StateT)
import qualified Data.Map as M
import IR (Id(..), StrTable)
import Operators (Ops)
import Parser (PAST)
import System.Console.Haskeline (InputT)

data Datum
  = Atom Int -- unifies a symbolic constant
  | Number Int -- unifies a numeric constant
  | Struct Id -- unifies a structure with arity
  | VoidRef -- unifies with anything
  | LocalRef Int -- code-local variable idx (should never occur on heap)
  | HeapRef Int -- something further on the heap
  deriving (Show, Eq, Ord)

data Instr
  = U Datum -- something unifiable
  | NoGoal -- trivial goal (directly after head)
  | Invoke Builtin -- also directly after head
  | Goal -- a new goal (set head)
  | Call -- all seems okay, call the head's hoal
  | LastCall -- tail call the head's goal
  | Cut -- remove all alternative clauses of the current goal
  deriving (Show)

type Code = [Instr]

type Defs = M.Map Id [Code]

data Heap =
  Heap Int (M.Map Int Datum)
  deriving (Show)

emptyHeap = Heap 1 M.empty

type Scope = M.Map Int Int

emptyScope :: Scope
emptyScope = M.empty

data Cho =
  Cho
    { hed :: Code -- head pointer
    , hvar :: Scope -- variables unified in head (so far)
    , gol :: Code -- goal pointer
    , gvar :: Scope -- variables unified in the goal
    , heap :: Heap -- a snapshot of the heap (there's no trail; we rely on CoW copies in other choicepoints)
    , stk :: [(Code, Scope, [Cho])] -- remaining goals together with their vars and cuts
    , cut :: [Cho] -- snapshot of choicepoints before entering
    }
  deriving (Show)

data Interp =
  Interp
    { defs :: Defs -- global definitions for lookup
    , cur :: Cho -- the choice that is being evaluated right now
    , cho :: [Cho] -- remaining choice points
    , ops :: Ops -- currently defined operators
    , opstash :: [Ops] -- saved operators
    , strtable :: StrTable -- string table
    , cmdq :: [(Bool, PAST)] -- isQuery, lexemes
    }
  deriving (Show)

type PrlgEnv a = StateT Interp (InputT IO) a

type InterpFn = PrlgEnv (Maybe (Either String Bool))

data Builtin =
  Builtin InterpFn

instance Show Builtin where
  show _ = "Builtin _"

data Dereferenced
  = FreeRef Int
  | BoundRef Int Datum
  | NoRef
  deriving (Show, Eq)

-- TRICKY: HeapRefs alone must not form a cycle other than the FreeRef case.
derefHeap :: Heap -> Int -> Dereferenced
derefHeap h@(Heap _ hmap) x =
  case hmap M.!? x of
    Just (HeapRef x') ->
      if x == x'
        then FreeRef x'
        else derefHeap h x'
    Just x' -> BoundRef x x'
    _ -> NoRef

writeHeap :: Int -> Datum -> Heap -> Heap
writeHeap addr x (Heap nxt m) = Heap nxt (M.adjust (const x) addr m)

newHeapVar :: Heap -> (Heap, Int)
newHeapVar heap = head <$> newHeapVars 1 heap

newHeapVars :: Int -> Heap -> (Heap, [Int])
newHeapVars n (Heap nxt m) =
  let addrs = [nxt + i - 1 | i <- [1 .. n]]
   in ( Heap (nxt + n) $
        foldr (uncurry M.insert) m [(a, HeapRef a) | a <- addrs]
      , addrs)

withNewHeapStruct :: Int -> Datum -> Heap -> ([Datum] -> Heap -> a) -> a
withNewHeapStruct addr s@(Struct Id {arity = arity}) heap cont =
  let (Heap nxt' m', addrs) = newHeapVars (arity + 1) heap
      m'' = M.insert addr (HeapRef $ head addrs) . M.insert (head addrs) s $ m'
   in cont (map HeapRef $ tail addrs) (Heap nxt' m'')

-- TODO are we actually going to use this?
codeStruct ::
     Monad m
  => (Datum -> m a)
  -> (Datum -> [a] -> m a)
  -> (Datum -> m (Either Int a))
  -> (Datum -> Int -> m a)
  -> m a
  -> Heap
  -> Code
  -> m (Code, a)
codeStruct atom struct local rec end heap = go
  where
    go [] = ([], ) <$> end
    go (U lr@(LocalRef _):cs) = do
      x <- local lr
      case x of
        Left ref -> (cs, ) <$> heapStruct atom struct rec heap ref
        Right a -> pure (cs, a)
    go (U s@(Struct (IR.Id _ arity)):cs) = eat arity cs >>= traverse (struct s)
    go (U x:cs) = (cs, ) <$> atom x
    go cs = (cs, ) <$> end
    eat n cs
      | n <= 0 = pure (cs, [])
      | otherwise = do
        (rest, a) <- go cs
        fmap (a :) <$> eat (n - 1) rest

heapStruct ::
     Monad m
  => (Datum -> m a)
  -> (Datum -> [a] -> m a)
  -> (Datum -> Int -> m a)
  -> Heap
  -> Int
  -> m a
heapStruct atom struct rec (Heap _ heap) hr = go [hr] hr
  where
    go visited ref
      | rr@(HeapRef r) <- heap M.! ref =
        if r == ref || r `elem` visited
          then rec rr ref
          else go (r : visited) r
      | s@(Struct (IR.Id _ arity)) <- heap M.! ref =
        sequence [go (ref + i : visited) (ref + i) | i <- [1 .. arity]] >>=
        struct s
      | x <- heap M.! ref = atom x