{-# 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.

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

-- 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. 

type Concurrent a = Cont Action a

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

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"



fork :: Concurrent () -> Concurrent ()
fork = undefined
  
action2 :: Concurrent ()
action2 = undefined