{-# LANGUAGE InstanceSigs #-}
module Main where

import Control.Monad 
import qualified Data.Map.Lazy as M
import Data.IORef

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)

data Vars
  = Vars { varsBindings :: M.Map Name Term }

emptyVars :: Vars 
emptyVars = Vars { varsBindings = 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 PrologStream a 
  = PYes (Maybe 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 Monad PrologStream where
  return :: a -> PrologStream a
  return = pure

  (>>=) :: PrologStream a -> (a -> PrologStream b) -> PrologStream b
  x >>= f = join_ (fmap f x)

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

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


join_ :: PrologStream (PrologStream a) -> PrologStream a
join_ (PNo) = PNo
join_ (PYes Nothing rest) = PYes Nothing (fmap join_ rest)
join_ (PYes (Just PNo) rest) = PYes Nothing (fmap join_ rest)
join_ (PYes (Just (PYes a xs)) rest) = 
  let joinedRest = liftM2 append xs (fmap join_ rest) in
  PYes a joinedRest

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


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


main :: IO ()
main = do
  putStrLn "Hello, world!"
{-
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!"

-}