{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}

import Control.Monad ( ap, liftM ) 
import qualified System.IO as IO
import qualified Data.IORef as IO
-- import qualified Network.Socket as Socket -- from the `network` library

-- * Currency as a monad

-- Just like we represent
-- arithmetic expression as a tree, 
-- we represent all the concurrent action
-- as a tree.


-- Just like we can define an evaluator
-- for expressions, we define a scheduler
-- for performing actions. Moreover, we
-- assume to work with an action queue. 

scheduler :: [Action] -> IO ()
scheduler [] = return ()
scheduler (x:xs) =
  case x of
    Stop -> scheduler xs
    Fork a1 a2 ->
      scheduler (xs ++ [a2, a1])
    Atomic m -> do
      a <- m
      scheduler (xs ++ [a])


-- We can define the following write action
-- which atomically write one character at a time.  
writeAction :: String -> Action
writeAction [] = Stop
writeAction (x:xs) = Atomic $ do
  putChar x
  return $ writeAction xs

action0 = Fork (writeAction "hello") (writeAction "world") 

test0 = scheduler [action0]


-- Q: How to sequence two actions? 
seq_action :: Action -> Action -> Action
seq_action = error "can't be defined!"

-- Solution: Continuation passing style (CPS)!

-- * A quick CPS style tutorial

-- direct style 
add :: Int -> Int -> Int
add x y = x + y

square :: Int -> Int
square x = x * x

pythagoras :: Int -> Int -> Int
pythagoras x y = add (square x) (square y)

-- CPS style

add' :: Int -> Int -> (Int -> r) -> r
add' x y k = k (x + y)  

square' :: Int -> (Int -> r) -> r
square' x k = k (x * x)  

-- The evaluation order will be explicit in
-- CPS style. 
pythagoras' :: Int -> Int -> (Int -> r) -> r 
pythagoras' x y k =
  square' x $ \ i ->
  square' y $ \ j ->
  add' i j (\ s -> k s)

-- CPS in recursion              
fact :: Int -> Int
fact 0 = 1
fact n = n * fact (n-1)


fact' :: Int -> (Int -> r) -> r
fact' 0 k = k 1
fact' n k = fact' (n-1) (\ x -> k (n * x))


instance Applicative (Cont r) where  
instance Functor (Cont r) where 


join :: Monad m => m (m a) -> m a
join x = do
  x' <- x
  x'
  
instance Monad (Cont r) where
  return :: a -> Cont r a
  return x = Cont $ \ k -> k x

  (>>=) :: Cont r a -> (a -> Cont r b) -> Cont r b
  (Cont g) >>= f =
    Cont $ \ k -> g $ \ x -> runCont (f x) k

    -- runCont (f x) :: (b -> r) -> r
    -- x :: a
    -- k :: b -> r
    -- g :: (a -> r) -> r
    -- f :: a -> Cont r b
    
    

-- Do notation for CPS.
add'' :: Int -> Int -> Cont r Int
add'' x y = return $ x + y

square'' :: Int -> Cont r Int
square'' x = return $ x * x 

pythagoras'' :: Int -> Int -> Cont r Int
pythagoras'' x y = do
  a <- square'' x
  b <- square'' y
  z <- add'' a b
  return z
  




-- In our setting, we will use Action for the type r.
-- So rather than returning an action, we take a continuation (Action -> Action),
-- and return Action. 


atomic :: IO a -> Concurrent a
atomic io = Cont $ \ k -> Atomic (do
                                     x <- io
                                     return $ k x
                                     )

run :: Concurrent a -> IO ()
run m = scheduler [runCont m $ \ x -> Stop]

  -- m :: Concurrent a = Cont Action a
  -- runCont m :: (a -> Action) -> Action 
  
  -- k :: a -> Action
  -- Atomic 



writeAction' :: String -> Concurrent () 
writeAction' [] = return ()
writeAction' (x:xs) = do
  atomic $ putChar x 
  writeAction' xs

action1 :: Concurrent ()
action1 = do
  writeAction' "hello"
  writeAction' "world"

type Concurrent a = Cont Action a

-- In fact, CPS forms a monad. 
data Cont r a = Cont ((a -> r) -> r)

runCont :: Cont r a -> (a -> r) -> r
runCont (Cont f) = f

data Action =
  Atomic (IO Action)
  | Fork Action Action
  | Stop

fork :: Concurrent () -> Concurrent ()
fork m =
  Cont $ \ k ->
           Fork (k ()) (runCont m (\ () -> Stop)) 
  -- m :: Concurrent () = Cont Action ()
  -- runCont m :: (() -> Action) -> Action
  -- k :: () -> Action 
  -- :: Cont Action ()
  
action2 :: Concurrent ()
action2 = do
  writeAction' "hello"
  fork $ writeAction' "my"
  fork $ writeAction' "name is"
  writeAction' "world"

action3 :: Concurrent ()
action3 = do
  fork $ writeAction' "hello"
  writeAction' "world"

action4 :: Concurrent ()
action4 = do
  writeAction' "hello"
  fork $ writeAction' "world"

action5 :: Concurrent ()
action5 = do
  writeAction' "------\n"
  fork $ writeAction' "hello\n"
  fork $ writeAction' "world\n"
  fork $ writeAction' "my name\n"
  fork $ writeAction' "is Haskell\n"

  


-- * Message passing through share state


-- IO.newIORef
-- IO.writeIORef
-- IO.readIORef


-- Let us define the following input/output helper functions. 
write :: String -> Concurrent ()
write s = atomic (putStr s)  


input :: Concurrent (Maybe String)
input = atomic $ do x <- IO.hReady IO.stdin
                    if x then do
                      s <- getLine
                      return $ Just s
                      else return Nothing

-- A data type for message. 
data Msg =
  Incr
  | Reset
  | Print
  | Quit
  deriving (Show, Eq)

-- Mutable state in the IO monad.
-- imported from the module Data.IORef. 

type Mailbox = IO.IORef (Maybe Msg)

newMailbox :: Concurrent Mailbox
newMailbox = 
  atomic $ IO.newIORef Nothing 


sendMsg :: Mailbox -> Msg -> Concurrent ()
sendMsg k m =
  atomic $ IO.writeIORef k (Just m) 

checkMsg :: Mailbox -> Concurrent (Maybe Msg)
checkMsg k =
  atomic $ do
    x <- IO.readIORef k
    case x of
      Just m -> IO.writeIORef k Nothing >> return (Just m)
      Nothing -> return Nothing

server :: Mailbox -> Integer -> Concurrent ()
server k i = do
  m <- checkMsg k
  case m of
    Nothing -> server k i
    Just Incr -> server k (i+1)
    Just Print -> do
      write $ "the current value is " ++ show i ++ "\n"
      server k i
    Just Reset -> do
      server k 0
    Just Quit -> do
      write $ "shutting down\n"
      return ()
      
  

client :: Mailbox -> Concurrent ()
client k = do
  x <- input
  case x of
    Nothing -> client k
    Just "a" -> do
      sendMsg k Incr
      client k
    Just "q" -> do
      sendMsg k Quit
      return ()
    Just "p" -> do
      sendMsg k Print
      client k
    Just "r" -> do
      sendMsg k Reset 
      client k
    Just _ -> do
      write "unkown command \n"
      client k


counter :: Concurrent ()
counter = do
  k <- newMailbox
  fork $ server k 0
  fork $ client k