content: {-# LANGUAGE BangPatterns, FlexibleContexts, ScopedTypeVariables #-}
  module Ros.Node.RosTcp (subStream, runServer, runServers, callServiceWithMaster) where
  import Control.Applicative ((<$>))
  import Control.Arrow (first)
  import Control.Concurrent (forkIO, killThread, newEmptyMVar, takeMVar, putMVar)
  import Control.Concurrent.STM (atomically)
  import Control.Concurrent.STM.TVar
  import qualified Control.Exception as E
  import Control.Monad.Reader
  import Data.Binary.Put (runPut, putWord32le)
  import Data.Binary.Get (runGet, getWord32le)
  import Data.ByteString.Lazy (ByteString)
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy as BL
  
  import Network.BSD (getHostByName, hostAddress)
  import Network.Socket hiding (send, sendTo, recv, recvFrom, Stream, ServiceName)
  import qualified Network.Socket as Sock
  import Network.Socket.ByteString
  import Prelude hiding (getContents)
  
  import System.IO (IOMode(ReadMode), hClose)
  import Text.URI (parseURI, uriRegName, uriPort)
  
  import Ros.Node.BinaryIter (streamIn, getServiceResult)
  import Ros.Internal.Msg.MsgInfo
  import Ros.Internal.Msg.SrvInfo
  import Ros.Internal.RosBinary
  import Ros.Internal.RosTypes
  import Ros.Internal.Util.RingChan
  import Ros.Internal.Util.AppConfig (Config, debug, forkConfig)
  import Ros.Topic (Topic(..))
  import Ros.Node.ConnectionHeader
  import Ros.Graph.Slave (requestTopicClient)
  import Ros.Graph.Master (lookupService)
  import Data.Maybe (fromMaybe)
  import Ros.Service.ServiceTypes
  import Control.Monad.Except
  import System.IO.Error (tryIOError)
  
  -- |Push each item from this client's buffer over the connected
  -- socket.
  serviceClient :: RingChan ByteString -> Socket -> IO ()
  serviceClient c s = forever $ do bs <- readChan c
                                   sendBS s bs
  
  sendBS :: Socket -> ByteString -> IO ()
  sendBS sock bs =
    let len = runPut $ 
              putWord32le . fromIntegral $ 
              BL.length bs
    in
     sendAll sock (BL.toStrict $ BL.append len bs)
  
  recvAll :: Socket -> Int -> IO B.ByteString
  recvAll s = flip go []
      where go len acc = do bs <- recv s len
                            if B.length bs < len
                              then go (len - B.length bs) (bs:acc)
                              else return $ B.concat (reverse (bs:acc))
  
  negotiatePub :: String -> String -> Socket -> IO ()
  negotiatePub ttype md5 sock = 
      do headerLength <- runGet (fromIntegral <$> getWord32le) <$>
                         BL.fromChunks . (:[]) <$> recvAll sock 4
         headerBytes <- recvAll sock headerLength
         let connHeader = parseHeader headerBytes
             wildCard = case lookup "type" connHeader of
                          Just t | t == "*" -> True
                                 | t == ttype -> False
                                 | otherwise -> error $ 
                                                "Disagreeing Topic types: " ++
                                                "publisher expected "++ttype++
                                                ", but client asked for "++t
                          Nothing -> error $ "Client did not include the "++
                                             "topic type in its "++
                                             "connection request."
         when (not wildCard) 
              (case lookup "md5sum" connHeader of
                 Just s | s == md5 -> return ()
                        | otherwise -> error "Disagreement on Topic type MD5"
                 Nothing -> error $ "Client did not include MD5 sum "++
                                    "in its request.")
         case lookup "tcp_nodelay" connHeader of
           Just "1" -> setSocketOption sock NoDelay 0
           _ -> return ()
         sendAll sock . genHeader $
           [("md5sum",md5), ("type",ttype), ("callerid","roshask")]
  
  -- |Accept new client connections. A new send buffer is allocated for
  -- each new client and added to the client list along with an action
  -- for cleaning up the client connection.
  -- FIXME: cleaning up a disconnected client should be reflected at a
  -- higher level, too.
  acceptClients :: Socket -> TVar [(Config (), RingChan ByteString)] -> 
                   (Socket -> IO ()) -> IO (RingChan ByteString) -> Config ()
  acceptClients sock clients negotiate mkBuffer = forever acceptClient
      where acceptClient = do (client,_) <- liftIO $ accept sock
                              debug "Accepted client socket"
                              liftIO $ negotiate client
                              chan <- liftIO mkBuffer
                              let cleanup1 = 
                                      do debug "Closing client socket"
                                         liftIO $ 
                                           shutdown client ShutdownBoth `E.catch`
                                             \(_::E.SomeException) -> return ()
                              r <- ask
                              t <- liftIO . forkIO $ 
                                   serviceClient chan client `E.catch` 
                                     \(_::E.SomeException) -> runReaderT cleanup1 r
                              let cleanup2 = cleanup1 >>
                                             (liftIO $ killThread t)
                              liftIO . atomically $ 
                                readTVar clients >>= 
                                writeTVar clients . ((cleanup2,chan) :)
  
  -- |Publish each item obtained from a 'Topic' to each connected client.
  pubStream :: RosBinary a
            => Topic IO a -> TVar [(b, RingChan ByteString)] -> Config ()
  pubStream t0 clients = liftIO $ go 0 t0
    where go !n t = do (x, t') <- runTopic t
                       let bytes = runPut $ putMsg n x
                       cs <- readTVarIO clients
                       mapM_ (flip writeChan bytes . snd) cs
                       go (n+1) t'
  
  -- |Produce a publishing action associated with a list of
  -- clients. This is used by runServers.
  pubStreamIO :: RosBinary a => IO (TVar [(b, RingChan ByteString)] -> Config (), 
                                    a -> IO ())
  pubStreamIO = do m <- newEmptyMVar
                   let feed clients = 
                         let go !n = do x <- takeMVar m
                                        let bytes = runPut $ putMsg n x
                                        cs <- readTVarIO clients
                                        mapM_ (flip writeChan bytes . snd) cs
                                        go (n+1)
                         in liftIO $ go 0
                   return (feed, putMVar m)
  
  -- Negotiate a TCPROS subscriber connection.
  -- Precondition: The socket is connected
  negotiateSub :: Socket -> String -> String -> String -> IO ()
  negotiateSub sock tname ttype md5 = 
      do sendAll sock $ genHeader [ ("callerid", "roshask"), ("topic", tname)
                                  , ("md5sum", md5), ("type", ttype) 
                                  , ("tcp_nodelay", "1") ]
         responseLength <- runGet (fromIntegral <$> getWord32le) <$>
                           BL.fromChunks . (:[]) <$> recvAll sock 4
         headerBytes <- recvAll sock responseLength
         let connHeader = parseHeader headerBytes
         case lookup "type" connHeader of
           Just t | t == ttype -> return ()
                  | otherwise -> error $ "Disagreeing Topic types: " ++
                                         "subscriber expected "++ttype++
                                         ", but server replied with "++t
           Nothing -> error $ "Server did not include the topic type "++
                              "in its response."
         case lookup "md5sum" connHeader of
           Just s | s == md5 -> return ()
                  | otherwise -> error "Disagreement on Topic type MD5"
           Nothing -> error "Server did not include MD5 sum in its response."
         setSocketOption sock KeepAlive 1
  
  -- |Connect to a publisher and return the stream of data it is
  -- publishing.
  subStream :: forall a. (RosBinary a, MsgInfo a) => 
               URI -> String -> (Int -> IO ()) -> Config (Topic IO a)
  subStream target tname _updateStats = 
      do debug $ "Opening stream to " ++target++" for "++tname
         h <- liftIO $ 
              do response <- requestTopicClient target "/roshask" tname 
                                                [["TCPROS"]]
                 let port = case response of
                              (1,_,("TCPROS",_,port')) -> fromIntegral port'
                              _ -> error $ "Couldn't get publisher's port for "++
                                           tname++" from node "++target
                 sock <- socket AF_INET Sock.Stream defaultProtocol
                 ip <- hostAddress <$> getHostByName host
                 connect sock $ SockAddrInet port ip
                 let md5 = sourceMD5 (undefined::a)
                     ttype = msgTypeName (undefined::a)
                 negotiateSub sock tname ttype md5
                 socketToHandle sock ReadMode
         --hSetBuffering h NoBuffering
         debug $ "Streaming "++tname++" from "++target
         return $ streamIn h
      where host = parseHost target
  
  parseHost :: URI -> String
  parseHost target = case parseURI target of
    Just u -> fromMaybe
              (error $ "Couldn't parse hostname "++ "from "++target)
              (uriRegName u)
    Nothing -> error $ "Couldn't parse URI "++target
  
  parseHostAndPort :: URI -> Either ServiceResponseExcept (String, PortNumber)
  parseHostAndPort target = do
    uri <- maybeToEither (ConnectExcept $ "Could not parse URI "++target) $ parseURI target
    host <- maybeToEither (ConnectExcept $ "Could not parse hostname from "++target) $ uriRegName uri
    port <- maybeToEither (ConnectExcept $ "Could not parse port from "++target) $ uriPort uri
    return (host, fromIntegral port)
  
  maybeToEither :: a -> Maybe b -> Either a b
  maybeToEither left m = case m of
    Just x -> Right x
    Nothing -> Left left
  
  --TODO: use the correct callerID
  callServiceWithMaster :: forall a b. (RosBinary a, SrvInfo a, RosBinary b, SrvInfo b) =>
                           URI -> ServiceName -> a -> IO (Either ServiceResponseExcept b)
  callServiceWithMaster rosMaster serviceName message = runExceptT $ do
    checkServicesMatch message (undefined::b)
    --lookup the service with the master
    (code, statusMessage, serviceUrl) <- lookupService rosMaster callerID serviceName
    checkLookupServiceCode code statusMessage
    (host, port) <- ExceptT . return $ parseHostAndPort serviceUrl
    -- make a socket
    let makeSocket = socket AF_INET Sock.Stream defaultProtocol
        closeSocket sock = liftIO $ close sock
        withSocket sock = do
          ioErrorToExceptT ConnectExcept "Problem connecting to server. Got exception : " $ do
            --Connect to the socket
            ip <- hostAddress <$> getHostByName host
            connect sock $ SockAddrInet port ip
          let reqMd5 = srvMD5 message
              reqServiceType = srvTypeName message
          negotiateService sock serviceName reqServiceType reqMd5
          let bytes = runPut $ putMsg 0 message
          ioErrorToExceptT SendRequestExcept "Problem sending request. Got exception: " $
            sendBS sock bytes
          liftIO $ socketToHandle sock ReadMode
    handle <- bracketOnErrorME (liftIO makeSocket) closeSocket withSocket
    result <- getServiceResult handle
    liftIO $ hClose handle
    return result
      where
        callerID = "roshask"
        checkLookupServiceCode 1 _ = return ()
        checkLookupServiceCode code statusMessage =
          throwError $ MasterExcept
          ("lookupService failed, code: " ++ show code ++ ", statusMessage: " ++ statusMessage)
        checkServicesMatch x y =
          unless match $
            -- throw error here since the calling code needs to be changed
            error "Request and response type do not match"
          where
            match = srvMD5 x == srvMD5 y && srvTypeName x == srvTypeName y
  
  -- | bracketOnError equivalent for MonadError
  bracketOnErrorME :: MonadError e m => m a -> (a -> m b) -> (a -> m c) -> m c
  bracketOnErrorME before after thing = do
    a <- before
    let handler e = after a >> throwError e
    catchError (thing a) handler
  
  -- | Catch any IOErrors and convert them to a different type
  catchConvertIO :: (String -> a) -> IO b -> IO (Either a b)
  catchConvertIO excep action = do
    err <- tryIOError action
    return $ case err of
      Left e -> Left . excep $ show e
      Right r -> Right r
  
  -- | Catch all IOErrors that might occur and convert them to a custom error type
  -- with the IOError message postpended to the given string
  ioErrorToExceptT :: (String -> e) -> String -> IO a -> ExceptT e IO a
  ioErrorToExceptT except msg acc =
    ExceptT . catchConvertIO (\m -> except $ msg ++ m) $ acc
  
  -- Precondition: The socket is already connected to the server
  -- Exchange ROSTCP connection headers with the server
  negotiateService :: Socket -> String -> String -> String -> ExceptT ServiceResponseExcept IO ()
  negotiateService sock serviceName serviceType md5 = do
      headerBytes <- liftIO $
        do sendAll sock $ genHeader [ ("callerid", "roshask"), ("service", serviceName)
                                  , ("md5sum", md5), ("type", serviceType) ]
           responseLength <- runGet (fromIntegral <$> getWord32le) <$>
                             BL.fromChunks . (:[]) <$> recvAll sock 4
           recvAll sock responseLength
      let connHeader = parseHeader headerBytes
      case lookup "error" connHeader of
        Nothing -> return ()
        Just _ -> throwError . ConHeadExcept $
                  "Connection header from server has error, connection header is: " ++ show connHeader
                                      
  -- Helper to run the publisher's side of a topic negotiation with a
  -- new client.
  mkPubNegotiator :: MsgInfo a => a -> Socket -> IO ()
  mkPubNegotiator x = negotiatePub (msgTypeName x) (sourceMD5 x)
  
  -- Run a publication server given a function that returns a
  -- negotiation action given a client 'Socket', a function that returns
  -- a publication action given a client list, a statistics updater, and
  -- the size of the send buffer.
  runServerAux :: (Socket -> IO ()) -> 
                  (TVar [(Config (), RingChan ByteString)] -> Config ()) -> 
                  (URI -> Int -> IO ()) -> Int -> Config (Config (), Int)
  runServerAux negotiate pubAction _updateStats bufferSize = 
      do r <- ask
         liftIO . withSocketsDo $ runReaderT go r
    where go = do sock <- liftIO $ socket AF_INET Sock.Stream defaultProtocol
                  liftIO $ bindSocket sock (SockAddrInet aNY_PORT iNADDR_ANY)
                  port <- liftIO (fromInteger . toInteger <$> socketPort sock)
                  liftIO $ listen sock 5
                  clients <- liftIO $ newTVarIO []
                  let mkBuffer = newRingChan bufferSize
                  acceptThread <- forkConfig $
                                  acceptClients sock clients negotiate mkBuffer
                  pubThread <- forkConfig $ pubAction clients
                  let cleanup = liftIO (atomically (readTVar clients)) >>= 
                                sequence_ . map fst >> 
                                liftIO (shutdown sock ShutdownBoth >>
                                        killThread acceptThread >>
                                        killThread pubThread)
                  return (cleanup, port)
  
  -- |The server starts a thread that peels elements off the stream as
  -- they become available and sends them to all connected
  -- clients. Returns an action for cleaning up resources allocated by
  -- this publication server along with the port the server is listening
  -- on.
  runServer :: forall a. (RosBinary a, MsgInfo a) => 
               Topic IO a -> (URI -> Int -> IO ()) -> Int -> 
               Config (Config (), Int)
  runServer stream = runServerAux (mkPubNegotiator (undefined::a)) 
                                  (pubStream stream)
  
  -- |The 'MsgInfo' type class dictionary made explicit to strip off the
  -- actual message type.
  data MsgInfoRcd = MsgInfoRcd { _md5, _typeName :: String }
  
  -- |A 'Feeder' represents a 'Topic' fully prepared to accept
  -- subscribers.
  data Feeder = Feeder MsgInfoRcd -- Explicit MsgInfo dictionary
                       Int -- Transmit buffer size
                       (URI -> Int -> IO ()) -- Update topic stats
                       (TVar [(Config (), RingChan ByteString)] -> Config ())
                       -- 'pubStream' partial application
  
  -- |Prepare an action for publishing messages. Arguments are a monadic
  -- function for updating topic statistics, and a transmit buffer
  -- size. The returned 'Feeder' value may be supplied to 'runServers',
  -- while the returned 'IO' function may be used to push out new
  -- messages.
  feedTopic :: forall a. (MsgInfo a, RosBinary a) => 
               (URI -> Int -> IO ()) -> Int -> IO (Feeder, a -> IO ())
  feedTopic updateStats bufSize = 
    do (feed,pub) <- pubStreamIO
       let f = Feeder info bufSize updateStats feed
       return (f, pub)
    where info = mkInfo (undefined::a)
          mkInfo x = MsgInfoRcd (msgTypeName x) (sourceMD5 x)
  
  -- |Publish several 'Topic's. A single cleanup action for all 'Topic's
  -- is returned, along with each 'Topic's server port in the order of
  -- the input 'Feeder's.
  runServers :: [Feeder] -> Config (Config (), [Int])
  runServers = return . first sequence_ . unzip <=< mapM feed
    where feed (Feeder (MsgInfoRcd md5 typeName) bufSize stats push) = 
            let pub = negotiatePub typeName md5
            in runServerAux pub push stats bufSize
  