--
--  Control.Concurrent.Session :: Session Types for Haskell
--  Copyright (C) 2007  Matthew Sackman (matthew@wellquite.org)
--
--  This library is free software; you can redistribute it and/or
--  modify it under the terms of the GNU Lesser General Public
--  License as published by the Free Software Foundation; either
--  version 2.1 of the License, or (at your option) any later version.
--
--  This library is distributed in the hope that it will be useful,
--  but WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
--  Lesser General Public License for more details.
--
--  You should have received a copy of the GNU Lesser General Public
--  License along with this library; if not, write to the Free Software
--  Foundation, Inc., 59 Temple Place, Suite 330, Boston,
--  MA  02111-1307  USA
--
{-# OPTIONS_GHC -fglasgow-exts #-}

module Tests where

import ChoiceTests
import Control.Concurrent.Session
import Control.Concurrent
import System.IO

sleep :: (JustSendsRecvs (SessionSpec s) (SessionSpec o) (SessionSpec i)) =>
         Int -> (SessionState (SessionSpec s) o i) ->
         IO ((), (SessionState (SessionSpec s) o i))
sleep = sliftIO' threadDelay

sendRecv a = SendS a $ RecvS a $ EndS

sendRecvTwice a = SendS a $ RecvS a $ SendS a $ RecvS a $ EndS

sendTwice a = SendS a $ SendS a $ EndS

sendTwice' a b = SendS a $ SendS b $ EndS

testSendRecvSimple = runSessionWithProcs (sendRecv BoolS) p1 p2
    where
      p1 = Proc $ (sendLog True)
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recvLog
           >~> (sendLog False)

testSendRecvTwicePartial :: IO (SessionSpec (SendT Bool (RecvT Bool EndT)))
testSendRecvTwicePartial = runSessionWithProcs (sendRecvTwice BoolS) p1 p2
    where
      p1 = Proc $ (sendLog True)
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recvLog
           >~> (sendLog False)

testSendTwiceRecvDelay :: IO (SessionSpec EndT)
testSendTwiceRecvDelay = runSessionWithProcs (sendTwice BoolS) p1 p2
    where
      p1 = Proc $ (sendLog True)
           >~> (sendLog False)
      p2 = Proc $ (sleep 1000000)
           >~> recvLog
           >~> sliftIO (putStrLn "boo")
           >~> recvLog
           >~> returnS ()

testSendTwiceMidDelay :: IO (SessionSpec EndT)
testSendTwiceMidDelay = runSessionWithProcs (sendTwice BoolS) p1 p2
    where
      p1 = Proc $ (sendLog True)
           >~> (sliftIO $ threadDelay 1000000)
           >~> (sendLog False)
      p2 = Proc $ recvLog
           >~> recvLog
           >~> returnS ()

testSendTwiceRecvDelay' :: IO (SessionSpec EndT)
testSendTwiceRecvDelay' = runSessionWithProcs (sendTwice' BoolS IntS) p1 p2
    where
      p1 = Proc $ (sendLog True)
           >~> (sendLog 5)
      p2 = Proc $ (sleep 1000000)
           >~> recvLog
           >~> recvLog
           >~> returnS ()

testProperEcho :: Int -> (Int -> Int) -> IO (SessionSpec EndT)
testProperEcho n f = runSessionWithProcs (sendRecv IntS) p1 p2
    where
      p1 = Proc $ sendLog n
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recvLog
           >~>= (sliftIO' (\t -> threadDelay t >> return t)) >~>= sendLog . f

sendEvalFunc :: (SType a) -> (SType b) -> SessionSpec (SendT c (SendT a (RecvT b EndT)))
sendEvalFunc a b = SendS AnyS . SendS a . RecvS b $ EndS

testSendEvalFunc :: Int -> (Int -> Int) -> IO (SessionSpec EndT)
testSendEvalFunc n f = runSessionWithProcs (sendEvalFunc IntS IntS) p1 p2
    where
      p1 = Proc $ send f
           >~> sendLog n
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recv
           >~>= \func -> recvLog
           >~>= \val -> sendLog (func val)

testSendEvalFunc' :: Int -> (Int -> Double) -> IO (SessionSpec EndT)
testSendEvalFunc' n f = runSessionWithProcs (sendEvalFunc IntS DoubleS) p1 p2
    where
      p1 = Proc $ send f
           >~> sendLog n
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recv
           >~>= \func -> recvLog
           >~>= \val -> sendLog (func val)

reduceSession :: SType a -> SType b ->
                 SessionSpec (SendT f (SendT a (SendT a (RecvT b EndT))))
reduceSession a b = SendS AnyS . SendS a . SendS a . RecvS b $ EndS

testEq :: (Show a, Eq a) => (SType a) -> (a -> a -> Bool) -> a -> a -> IO (SessionSpec EndT)
testEq t f v1 v2 = runSessionWithProcs (reduceSession t BoolS) p1 p2
    where
      p1 = Proc $ send f
           >~> sendLog v1
           >~> sendLog v2
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recv
           >~>= \func -> recvLog
           >~>= \valA -> recvLog
           >~>= \valB -> sendLog (func valA valB)

reducePolySession :: SType a -> SType b ->
                     SessionSpec (SendT f
                                  (SendT a
                                   (SendT a
                                    (RecvT a
                                     (SendT b
                                      (SendT b
                                       (RecvT b EndT)))))))
reducePolySession a b = SendS AnyS . SendS a . SendS a . RecvS a . SendS b . SendS b . RecvS b $ EndS

reducePolyToListSession :: SType a -> SType b ->
                           SessionSpec (SendT f
                                        (SendT a
                                         (SendT a
                                          (RecvT [a]
                                           (SendT b
                                            (SendT b
                                             (RecvT [b] EndT)))))))
reducePolyToListSession a b = SendS AnyS . SendS a . SendS a . RecvS (ListS a) . SendS b . SendS b . RecvS (ListS b) $ EndS

data FnWrapper :: * where
                     NumW :: (forall a . (Num a) => a -> a -> a) -> FnWrapper
                     EnumW  :: (forall a . (Enum a) => a -> a -> [a]) -> FnWrapper

testPolyEvalFunc :: (Num a, Num b) => FnWrapper -> a -> a -> b -> b -> IO (SessionSpec EndT)
testPolyEvalFunc f al ar bl br = runSessionWithProcs (reducePolySession AnyS AnyS) p1 p2
    where
      p1 = Proc $ send f
           >~> sendLog al
           >~> sendLog ar
           >~> recvLog
           >~> sendLog bl
           >~> sendLog br
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recv
           >~>= \f -> recv
           >~>= \l -> recv
           >~>= \r -> send (let (NumW x) = f in l `x` r)
           >~> recv
           >~>= \l -> recv
           >~>= \r -> send (let (NumW x) = f in l `x` r)
           >~> returnS ()

testPolyEnumEvalFunc :: (Show a, Show b, Enum a, Enum b) => FnWrapper -> a -> a -> b -> b -> IO (SessionSpec EndT)
testPolyEnumEvalFunc f al ar bl br = runSessionWithProcs (reducePolyToListSession AnyS AnyS) p1 p2
    where
      p1 = Proc $ send f
           >~> sendLog al
           >~> sendLog ar
           >~> recvLog
           >~> sendLog bl
           >~> sendLog br
           >~> recvLog
           >~> returnS ()
      p2 = Proc $ recv
           >~>= \f -> recv
           >~>= \l -> recv
           >~>= \r -> send (let (EnumW x) = f in x l r)
           >~> recv
           >~>= \l -> recv
           >~>= \r -> send (let (EnumW x) = f in x l r)
           >~> returnS ()

simpleLoopSpec :: SessionSpec (LoopT (SessionSpec (SendT Int LoopEndT)))
simpleLoopSpec = mkLoopS (SendS IntS)

testSimpleLoop :: IO (SessionSpec EndT)
testSimpleLoop = runSessionWithProcs simpleLoopSpec p1 p2
    where
      p1 = Proc $ mkLoop (\n -> (if (n `mod` 85000) == 0 then (sendLog n) else (send n))
                          >~> returnS (n+1)
                         ) 0
      p2 = Proc $ mkLoop (\n -> (if (n `mod` 85000) == 0 then recvLog else recv)
                          >~>= \m -> if n == m then returnS (n+1) else error "Wrong order!")
                         0

sendRecvLoop = mkLoopS (SendS IntS . RecvS IntS)
testSendRecvLoop = runSessionWithProcs sendRecvLoop p1 p2
    where
      p1 = Proc $ mkLoop (\n ->
                          (if (n `mod` 10000) == 0 then sendLog n else send n) >~>
                          recv >~>=
                          returnS . (+1)
                         ) 0
      p2 = Proc $ mkLoop (\n ->
                          (if (n `mod` 10000) == 0 then recvLog else recv) >~>= \m ->
                          send m >~>
                          returnS (m+1)
                         ) 0

simpleLoopSpec' = LoopS (SendS IntS EndS)
testSimpleLoop' = runSessionWithProcs simpleLoopSpec' p1 p2
    where
      p1 = Proc $ mkLoop (\() -> send 6 >~> end ()) ()
      p2 = Proc $ mkLoop (\() -> recvLog >~> end ()) ()

{-
calcOffer = LoopS (OfferS calcOfferList)
calcOfferList =  ((CalcOpQuit, EndS)
                  ~||~
                  (CalcOpAdd, (RecvS IntS . RecvS IntS . SendS IntS $ LoopEndS))
                  ~||~
                  (CalcOpSubtract, (RecvS IntS . RecvS IntS . SendS IntS $ LoopEndS))
                  ~||~
                  (CalcOpNegate, (RecvS IntS . SendS IntS $ LoopEndS))
                  ~||~ specList
                 )

calcOfferChoices = mkLoop (\() ->
                           offer ((CalcOpQuit, Proc $ returnS ())
                                  ~||~
                                  (CalcOpAdd,
                                   Proc (recv >~>= \l -> recv >~>= \r ->
                                         send (l + r) >~> returnS ()))
                                  ~||~
                                  (CalcOpSubtract,
                                   Proc (recv >~>= \l -> recv >~>= \r ->
                                         send (l - r) >~> returnS ()))
                                  ~||~
                                  (CalcOpNegate,
                                   Proc (recv >~>= send . negate >~> returnS ()))
                                  ~||~ procList)) ()

testCalc = runSessionWithProcs calcOffer p1 p2
    where
      p1 = Proc $ calcOfferChoices
      p2 = Proc $ mkLoop (\() -> (sliftIO talkToUser) >~>= \f -> f) ()
      add = select CalcOpAdd (Proc (sendLog 5
                                    >~> sendLog 8
                                    >~> recvLog
                                    >~> returnS ()))
      subtract = select CalcOpSubtract (Proc (sendLog 3
                                              >~> sendLog 6
                                              >~> recvLog
                                              >~> returnS ()))
      neg = select CalcOpNegate (Proc (sendLog 10 >~> recvLog >~> returnS ()))
--      quit = select CalcOpQuit (Proc (returnS ()))
      talkToUser
          = do { hFlush stdout
               ; hPutStrLn stdout "Enter 'Add' or 'Subtract' or 'Negate' or anything else to Quit:"
               ; hFlush stdout
               ; line <- hGetLine stdin
               ; return $ case line of
                            a | (Just _) <- prefixMatch "Add" line -> add
                            n | (Just _) <- prefixMatch "Negate" line -> neg
                            s | (Just _) <- prefixMatch "Subtract" line -> subtract
--                            _ -> quit
               }

prefixMatch :: String -> String -> Maybe String
prefixMatch prefix str
    = case prefix == strP of
        True -> Just strT
        False -> Nothing
    where
      len = length prefix
      (strP, strT) = splitAt len str
-}

calcOffer =  OfferS ((CalcOpAdd, (RecvS IntS . RecvS IntS . SendS IntS $ EndS))
                     ~||~
                     (CalcOpSubtract, (RecvS IntS . RecvS IntS . SendS IntS $ EndS))
                     ~||~
                     (CalcOpNegate, (RecvS IntS . SendS IntS $ EndS))
                     ~||~
                     (CalcOpQuit, EndS)
                     ~||~ specList
                    )

calcOfferChoices = offer ((CalcOpAdd,
                           Proc (recv >~>= \l -> recv >~>= \r ->
                                 send (l + r) >~> returnS ()))
                          ~||~
                          (CalcOpSubtract,
                           Proc (recv >~>= \l -> recv >~>= \r ->
                                 send (l - r) >~> returnS ()))
                          ~||~
                          (CalcOpNegate,
                           Proc (recv >~>= send . negate >~> returnS ()))
                          ~||~
                          (CalcOpQuit,
                           Proc (returnS ()))
                          ~||~ procList)

testCalc = runSessionWithProcs calcOffer p1 p2
    where
      p1 = Proc $ calcOfferChoices
      p2 = Proc $ (sliftIO talkToUser) >~>= id
      add = select CalcOpAdd
            (Proc $
             (sliftIO grabTwoNumbers) >~>= \[x, y] ->
             sendLog x >~> sendLog y >~> recvLog >~> returnS ())
      subtract = select CalcOpSubtract
                 (Proc $
                  (sliftIO grabTwoNumbers) >~>= \[x, y] ->
                  sendLog x >~> sendLog y >~> recvLog >~> returnS ())
      neg = select CalcOpNegate
            (Proc $ (sliftIO grabNumber) >~>= \x ->
             sendLog x >~> recvLog >~> returnS ())
      quit = select CalcOpQuit (Proc (returnS ()))
      talkToUser
          = do { hPutStrLn stdout "Enter 'Add' or 'Subtract' or 'Negate' or anything else to Quit:"
               ; line <- hGetLine stdin
               ; return $ case line of
                            "Add" -> add
                            "Negate" -> neg
                            "Subtract" -> subtract
                            _ -> quit
               }
      grabTwoNumbers = sequence (replicate 2 grabNumber)
      grabNumber = putStrLn "Enter a number:" >> hGetLine stdin >>= return . read