{-# LANGUAGE InstanceSigs #-}
module Main where

import Control.Monad 
import qualified Data.Map.Lazy as M
import Data.IORef
import Control.Monad.RWS (MonadTrans(lift))
import Control.Applicative

data Atomic
  = AString String
  deriving (Eq, Show)

data Name
  = NString String
  deriving (Ord, Eq, Show)

data Term 
  = TCompound Atomic [Term] -- functor, args
  | TVariable Name  -- name
  deriving (Eq, Show)

type Vars = M.Map Name Term

emptyVars :: Vars 
emptyVars = M.empty


data Globals = Globals {
  psVars :: IORef Vars
}

newtype PrologIO a = PrologIO (Globals -> IO a)

instance Functor PrologIO where
  fmap :: (a -> b) -> PrologIO a -> PrologIO b
  fmap f (PrologIO a) = PrologIO (\globals -> fmap f (a globals))

instance Applicative PrologIO where
  pure :: a -> PrologIO a
  pure x = PrologIO (\_ -> pure x)

  (PrologIO f) <*> (PrologIO x) = PrologIO (
    \globals -> do
      f' <- f globals
      x' <- x globals
      pure (f' x')
    )

instance Monad PrologIO where
  return :: a -> PrologIO a
  return = pure

  (PrologIO a) >>= f = PrologIO (
    \globals -> do
      a' <- a globals
      let PrologIO b' = f a'
      b' globals
    )

runPrologIO :: Globals -> PrologIO a -> IO a
runPrologIO globals (PrologIO a) = a globals

data Suspend a = Yield a | Suspend

instance Functor Suspend where
  fmap f (Yield a) = Yield (f a)
  fmap _ Suspend = Suspend

data PrologStream a 
  = PYes (Suspend a) (PrologIO (PrologStream a))
  | PNo

instance Functor PrologStream where
  fmap :: (a -> b) -> PrologStream a -> PrologStream b
  fmap f (PYes a next) = PYes (fmap f a) (fmap (fmap f) next)
  fmap _ PNo = PNo

instance Applicative PrologStream where
  pure :: a -> PrologStream a
  pure x = PYes (Yield x) (pure PNo)

  (<*>) :: PrologStream (a -> b) -> PrologStream a -> PrologStream b
  f <*> x = f >>= (\f' -> x >>= (\x' -> return (f' x')))

instance Monad PrologStream where
  return :: a -> PrologStream a
  return = pure

  (>>=) :: PrologStream a -> (a -> PrologStream b) -> PrologStream b
  PNo >>= _ = PNo
  PYes (Yield x) xs >>= f = f x <|> PYes Suspend (fmap (>>= f) xs)
  PYes Suspend xs >>= f = PYes Suspend (fmap (>>= f) xs)

instance Alternative PrologStream where 
  empty :: PrologStream a
  empty = PNo

  (<|>) :: PrologStream a -> PrologStream a -> PrologStream a
  (<|>) = append


instance MonadPlus PrologStream where
  mzero :: PrologStream a
  mzero = empty

  mplus :: PrologStream a -> PrologStream a -> PrologStream a
  mplus = (<|>)


append :: PrologStream a -> PrologStream a -> PrologStream a
append (PYes x xs) next = PYes x (liftM2 append xs (return next))
append PNo next = next


data Thread a = Thread 
  { threadGlobals :: Globals
  , threadStream :: IORef (PrologStream a)
  }

startThread :: PrologStream a -> IO (Thread a)
startThread stream = do
  vars <- newIORef emptyVars
  streamRef <- newIORef stream 
  let globals = Globals { psVars = vars }
  return (Thread { threadGlobals = globals, threadStream = streamRef })

advanceThread :: Thread a -> IO (Maybe a)
advanceThread (Thread { threadGlobals = globals, threadStream = streamRef }) = do
  stream <- readIORef streamRef
  (value, newStream) <- runPrologIO globals (squeeze stream)
  writeIORef streamRef newStream
  return value

completeThread :: Thread a -> IO [a]
completeThread thread = do
  value <- advanceThread thread
  case value of 
    Nothing -> return []
    Just x -> fmap (x:) (completeThread thread)

squeeze :: PrologStream a -> PrologIO (Maybe a, PrologStream a)
squeeze PNo = return (Nothing, PNo)
squeeze (PYes Suspend rest) = rest >>= squeeze
squeeze (PYes (Yield x) rest) = fmap ((,) (Just x)) rest 

exitUndo :: PrologIO a -> PrologIO () -> PrologStream a
exitUndo exit undo = 
  PYes Suspend $ exit >>= \v -> return $ 
    PYes (Yield v) $ undo >>= \_ -> return $
      PNo

liftPrologIO :: PrologIO a -> PrologStream a
liftPrologIO exit = exitUndo exit (return ())

liftIO :: IO a -> PrologStream a
liftIO x = liftPrologIO $ PrologIO (\_ -> x)

setVar :: Name -> Term -> PrologStream ()
setVar name value = exitUndo exit undo
  where 
  exit = PrologIO (
    \globals -> do 
      let varsRef = psVars globals
      modifyIORef varsRef (M.insert name value)
    )
  undo = PrologIO (
    \globals -> do
      let varsRef = psVars globals
      modifyIORef varsRef (M.delete name)
    )

getVar :: Name -> PrologStream (Maybe Term)
getVar name = exitUndo exit undo
  where
  exit = PrologIO (
    \globals -> do
      let varsRef = psVars globals
      vars <- readIORef varsRef
      return (M.lookup name vars)
    )
  undo = PrologIO (\_ -> return ())

  
demoProgram :: PrologStream ()
demoProgram = do
  setVar (NString "bat") (TCompound (AString "Nyeogmi") []) <|> setVar (NString "bat") (TCompound (AString "Pyrex") []) 
  setVar (NString "activity") (TCompound (AString "drinks blood") []) <|> setVar (NString "activity") (TCompound (AString "spits venom") []) 
  who <- getVar (NString "bat")
  what <- getVar (NString "activity")
  liftIO (print ("Value: ", who, "What: ", what))


main :: IO ()
main = do
  thread <- startThread demoProgram
  results <- completeThread thread
  putStrLn ("Results: " ++ show results)
{-
no :: Prolog a
no = Prolog (\_ -> [])

readVars :: Prolog Vars
readVars = Prolog (\vars -> [(vars, vars)])

writeVars :: Vars -> Prolog ()
writeVars vars = Prolog (\_ -> [((), vars)])

unify :: Term -> Term -> Prolog ()
unify = unify_ []
  where
  unify_ :: [Name] -> Term -> Term -> Prolog ()
  unify_ names (TCompound h0 args0) (TCompound h1 args1) = do
    when (h0 /= h1) no
    when (length args0 /= length args1) no
    forM_ (zip args0 args1) $ \(arg0, arg1) -> unify_ names arg0 arg1
  unify_ _ (TVariable name0) (TVariable name1) | name0 == name1 = return ()
  unify_ names (TVariable name) y = do
    -- occurs check
    when (name `elem` names) no
    -- unify as usual
    (Vars vs0) <- readVars
    case M.lookup name vs0 of
      Just existing -> unify_ (name:names) existing y
      Nothing -> do
        let vs1 = M.insert name y vs0
        writeVars (Vars vs1)
  unify_ names x@(TCompound _ _) y@(TVariable _) = unify_ names y x





main :: IO ()
main = putStrLn "Hello, Haskell!"

-}