import Control.Monad (ap, liftM, void, (<=<), (>=>))
import Data.IORef (IORef, newIORef, readIORef, writeIORef)

-- What do I owe you?
-- Nothing, it's for free!
-- Monads? For freee????!

-- I guess I'm forced to write haskell the old fashioned way now -- Writing with a bat? A butterfly

newtype FState s a = FState {runFState :: s -> (a, s)}

instance Functor (FState s) where
  fmap = liftM

instance Applicative (FState s) where
  -- pure x = FState $ \s -> (x, s)
  -- pure x = FState $ (x,)
  pure = FState . (,)
  (<*>) = ap

instance Monad (FState s) where
  c >>= f = FState $ \s ->
    let (a, s') = runFState c s
     in runFState (f a) s'

get :: FState s s
get = FState (\s -> (s, s))

put :: s -> FState s ()
-- put s = FState (\_ -> ((), s))
-- put s = FState (const ((), s))
put = FState . const . ((),) -- ~~Pointless!~~ Pointfree! -.-

-- The get, put and runState functions define the State monad
-- Alongside the following rules:
-- put s >> get = put s >> return s
-- put s >> put s' = put s'
-- get >>= put = return ()

-- type Monad :: (* -> *) -> Constraint
-- class Applicative m => Monad m where
--   ...
--   (>>) :: m a -> m b -> m b
--   ...
--         -- Defined in ‘GHC.Internal.Base’
-- infixl 1 >>

-- yee 🦕 Dino nuggies, omg i want dino nuggies

class (Monad m) => StateMonad m s where
  getm :: m s
  putm :: s -> m ()

modify :: (StateMonad m s) => (s -> s) -> m ()
-- modify f = getm >>= (putm . f)
modify f = f <$> getm >>= putm

tick :: (StateMonad m Int) => m Int
-- tick = do{n <- getm; putm (n+1); return n}
tick = do n <- getm; n <$ putm (n + 1)

push :: (StateMonad m [a]) => a -> m ()
push x = modify (x :)

pop :: (StateMonad m [a]) => m (Maybe a)
pop =
  do
    xs <- getm
    case xs of
      [] -> return Nothing
      (x : xs') -> do putm xs'; return $ Just x

newtype IState s a = IState {runIState :: IORef s -> IO a}

evalIState :: s -> IState s a -> IO a
evalIState def m = newIORef def >>= runIState m

instance Functor (IState s) where
  fmap = liftM

instance Applicative (IState s) where
  pure = IState . const . pure
  (<*>) = ap

instance Monad (IState s) where
  c >>= f = IState $ \ref -> do
    val <- runIState c ref
    runIState (f val) ref

instance StateMonad (IState s) s where
  getm = IState $ readIORef
  putm = IState . flip writeIORef

-- A `continuation` is passed into the function
-- Continuation passing style!
putk :: (StateMonad m s) => s -> (() -> m a) -> m a
putk s = (putm s >>=)

data FSM s a where
  Return :: a -> FSM s a
  Getk :: (s -> FSM s a) -> FSM s a
  Putk :: s -> FSM s a -> FSM s a

instance Functor (FSM s) where
  fmap = liftM

instance Applicative (FSM s) where
  pure = Return
  (<*>) = ap

instance Monad (FSM s) where
  Return x >>= f = f x
  Getk n >>= f = Getk $ \s -> n s >>= f -- GetK (n >=> f)
  Putk v n >>= f = Putk v $ n >>= f -- PutK ()

instance StateMonad (FSM s) s where
  getm = Getk return
  putm state' = Putk state' (return ())

-- This function allows FSM to be evaluated under any other StateMonad
evalFSM :: (StateMonad m s) => FSM s a -> m a
evalFSM (Return x) = return x
evalFSM (Getk next) = getm >>= \state -> evalFSM (next state)
evalFSM (Putk state' next) = putm state' *> evalFSM next

runFSM :: s -> FSM s a -> (a, s)
runFSM s (Return a) = (a, s)
runFSM s (Getk n) = runFSM s (n s)
runFSM s (Putk s' n) = runFSM s' n

-- I'm not sure I get why these constructors are needed. Which one?
-- FSM itself, there are just the functions as constructors
-- Exactly! FSM is a Free State Monad that encodes what a state typeclass would have, but inside a datatype
-- So you can have a type that models state, but in the end you can evaluate it however you actually want

data StateOp s a where
  StateGet :: (s -> a) -> StateOp s a
  StatePut :: s -> a -> StateOp s a

data Free e a where
  Pure :: a -> Free e a
  Free :: e (Free e a) -> Free e a

instance (Functor e) => Functor (Free e) where
  fmap = liftM

instance (Functor e) => Applicative (Free e) where
  pure = Pure
  (<*>) = ap

instance (Functor e) => Monad (Free e) where
  (Pure a) >>= f = f a
  (Free c) >>= f = Free $ fmap (>>= f) c

type FreeState s a = Free (StateOp s) a

-- Space