module Interpreter where

import Data.Function
import qualified Data.Map as M

{- VAM 2P, done the lazy way -}
data StrTable =
  StrTable Int (M.Map String Int) (M.Map Int String)
  deriving (Show)

emptystrtable = StrTable 0 M.empty M.empty

strtablize t@(StrTable nxt fwd rev) str =
  case fwd M.!? str of
    Just i -> (t, i)
    _ -> (StrTable (nxt + 1) (M.insert str nxt fwd) (M.insert nxt str rev), nxt)

data Id =
  Id
    { str :: Int
    , arity :: Int
    }
  deriving (Show, Eq, Ord)

data Datum
  = Atom Int -- unifies a constant
  | Struct Id -- unifies a structure with arity
  | VoidVar -- unifies with anything
  -- | LocalVar Int -- unifies with a local variable (possibly making a new one when it's not in use yet)
  -- | Ref Int -- unifies with the referenced value on the heap (not to be used in code)
  deriving (Show, Eq, Ord)

data Instr
  = U Datum -- something unifiable
  | NoGoal -- trivial goal
  | Goal -- we start a new goal, set up backtracking etc
  | Call -- all seems okay, call the goal
  | LastCall -- tail call the goal
  | Cut -- remove all alternative clauses of the current goal
  deriving (Show)

type Code = [Instr]

type Defs = M.Map Id [Code]

data Cho =
  Cho
    { hed :: Code -- head pointer
    , gol :: Code -- goal pointer
    , stk :: [(Code, [Cho])] -- remaining "and" goals
    , 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
    }
  deriving (Show)

prove :: Code -> Defs -> (Interp, Either String Bool)
prove g ds =
  let i0 =
        Interp
          { defs = ds
          , cur = Cho {hed = g, gol = [LastCall], stk = [], cut = []}
          , cho = []
          }
      run (Left x) = x
      run (Right x) = run $ proveStep Right (\i e -> Left (i, e)) x
   in run (Right i0)

{- this gonna need Either String Bool for errors later -}
proveStep :: (Interp -> a) -> (Interp -> Either String Bool -> a) -> Interp -> a
proveStep c f i = go i
  where
    ifail msg = f i $ Left msg
    tailcut [LastCall] chos _ = Just chos
    tailcut [LastCall, Cut] _ cut = Just cut
    tailcut _ _ _ = Nothing
    withDef fn
      | Just d <- defs i M.!? fn = ($ d)
      | otherwise = const $ ifail $ "no definition: " ++ show fn
    {- Backtracking -}
    backtrack i@Interp {cho = chos}
      {- if available, restore the easiest choicepoint -}
      | (cho:chos) <- chos = c i {cur = cho, cho = chos}
      {- if there's no other choice, answer no -}
      | otherwise = f i $ Right False
    {- Unification -}
    go i@Interp {cur = cur@Cho {hed = U a:hs, gol = U b:gs}} -- unify constants
     = unify a b
      where
        uok = c i {cur = cur {hed = hs, gol = gs}}
        unify VoidVar VoidVar = uok
        unify (Atom a) (Atom b) | a == b = uok
        unify VoidVar (Atom _) = uok
        unify (Atom _) VoidVar = uok
        unify (Struct a) (Struct b) | a == b = uok
        unify VoidVar (Struct Id{arity=a}) = c i {cur = cur {hed = replicate a (U VoidVar) ++ hs, gol = gs}}
        unify (Struct Id{arity=a}) VoidVar = c i {cur = cur {hed = hs, gol = replicate a (U VoidVar) ++ gs}}
        unify _ _ = backtrack i
    {- Resulution -}
    go i@Interp { cur = cur@Cho {hed = hed, gol = gol, stk = stk, cut = cut}
                , cho = chos
                }
      {- top-level success -}
      | [NoGoal] <- hed
      , Just nchos <- tailcut gol chos cut
      , [] <- stk =
        f i {cur = cur {hed = [], gol = []}, cho = nchos} $ Right True
      {- cut before the first goal (this solves all cuts in head) -}
      | Cut:hs <- hed = c i {cur = cur {hed = hs}, cho = cut}
      {- succeed and return to caller -}
      | [NoGoal] <- hed
      , Just nchos <- tailcut gol chos cut
      , (Goal:U (Struct fn):gs, _):ss <- stk =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, gol = gs, stk = ss}
              , cho = [Cho oh gs ss nchos | oh <- ohs] ++ nchos
              }
      {- succeed and return to caller, and the caller wants a cut -}
      | [NoGoal] <- hed
      , Just _ <- tailcut gol chos cut
      , (Cut:Goal:U (Struct fn):gs, rchos):ss <- stk =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, gol = gs, stk = ss}
              , cho = [Cho oh gs ss rchos | oh <- ohs] ++ rchos
              }
      {- start matching next goal -}
      | [NoGoal] <- hed
      , (Call:Goal:U (Struct fn):gs) <- gol =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, gol = gs}
              , cho = [Cho oh gs stk chos | oh <- ohs] ++ chos
              }
      {- start matching next goal after a cut -}
      | [NoGoal] <- hed
      , (Call:Cut:Goal:U (Struct fn):gs) <- gol =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, gol = gs}
              , cho = [Cho oh gs stk cut | oh <- ohs] ++ cut
              }
      {- goal head matching succeeded, make a normal call -}
      | (Goal:U (Struct fn):ngs) <- hed
      , (Call:gs) <- gol =
        withDef fn $ \(hs:ohs) ->
          let nstk = (gs, chos) : stk
           in c i
                  { cur = cur {hed = hs, gol = ngs, stk = nstk}
                  , cho = [Cho oh ngs nstk chos | oh <- ohs] ++ chos
                  }
      {- successful match continued by tail call -}
      | (Goal:U (Struct fn):ngs) <- hed
      , Just nchos <- tailcut gol chos cut =
        withDef fn $ \(hs:ohs) ->
          c
            i
              { cur = cur {hed = hs, gol = ngs}
              , cho = [Cho oh ngs stk nchos | oh <- ohs] ++ nchos
              }
    {- The End -}
    go _ = ifail "impossible instruction combo"