content: {-# LANGUAGE OverloadedStrings, TupleSections #-}
  module Analysis (MsgInfo, liftIO, getTypeInfo, withMsg, getMsg, addMsg,
                   runAnalysis, isFlat, SerialInfo(..)) where
  import Control.Applicative
  import Control.Arrow ((&&&))
  import Control.Monad.State
  import Data.ByteString.Char8 (pack, unpack, ByteString)
  import qualified Data.ByteString.Char8 as B
  import qualified Data.Map as M
  import Data.Maybe (isJust)
  import System.FilePath (takeFileName, dropExtension)
  import Ros.Internal.DepFinder (findMessagesInPkg)
  import Types hiding (msgName)
  import Parse
  import ResolutionTypes
  
  -- Synonym for a Msg paired with its SerialInfo. This tuple is cached
  -- for message types as they are visited.
  type SerialMsg = (SerialInfo, Msg)
  
  -- Front-end to run analyses.
  runAnalysis :: MsgInfo a -> IO a
  runAnalysis m = evalStateT (cachePackage "rosgraph_msgs" >> m) emptyMsgContext
  
  -- All the .msg files in a package are cached for quick lookup on
  -- subsequent type resolutions.
  cachePackage :: ByteString -> MsgInfo PkgCache
  cachePackage pkgName = do (dir, msgs) <- liftIO $ findMessagesInPkg (unpack pkgName)
                            let cache = M.fromList (map prepMsg msgs)
                                pkg = (dir, M.empty, cache)
                            alterPkgMap (M.insert pkgName pkg)
                            return pkg
      where prepMsg = pack . dropExtension . takeFileName &&& Left
  
  -- Get a package's 'MsgCache' from the package cache.
  getPackage :: ByteString -> MsgInfo PkgCache
  getPackage pkgName = 
      do pkgs <- msgDefs <$> get
         maybe (cachePackage pkgName) return $ M.lookup pkgName pkgs
  
  -- Try to get a 'Msg' from a given package. The first argument is the
  -- package name, the second argument is the unqualified message name.
  getMsgFromPkg :: ByteString -> ByteString -> MsgInfo (Maybe SerialMsg)
  getMsgFromPkg pkgName msgName = getPackage pkgName >>= lookupMsg . msgCache
      where lookupMsg :: MsgCache -> MsgInfo (Maybe SerialMsg)
            lookupMsg cache = maybe (return Nothing) 
                                    (return . Just <=< loadMsg)
                                    (M.lookup msgName cache)
            loadMsg :: (Either FilePath SerialMsg) -> MsgInfo SerialMsg
            loadMsg (Left fp) = either error addMsg =<< liftIO (parseMsg fp)
            loadMsg (Right m) = return m
            msgCache (_,_,c)  = c
  
  getMsg :: ByteString -> MsgInfo SerialMsg
  getMsg msgName = check <$>
                   if B.null msgType
                   then getMsgFromPkg "rosgraph_msgs" msgName <||>
                        getMsgFromPkg "std_msgs" msgName <||>
                        (flip getMsgFromPkg msgName . homePkg =<< get)
                   else getMsgFromPkg msgPkg (B.tail msgType)
      where (msgPkg, msgType) = B.span (/= '/') msgName
            check :: Maybe SerialMsg -> SerialMsg
            check Nothing = error $ "Couldn't resolve type " ++ unpack msgName
            check (Just m) = m
            -- checkLocal :: Maybe SerialMsg -> MsgInfo (Maybe SerialMsg)
            -- checkLocal Nothing = do home <- homePkg <$> get
            --                         getMsgFromPkg home msgName
            -- checkLocal info    = return info
            (<||>) :: (Applicative f, Alternative g) => f (g a) -> f (g a) -> f (g a)
            (<||>) = liftA2 (<|>)
  
  isFlat :: Msg -> MsgInfo Bool
  isFlat = fmap (all isStorable) . mapM (getTypeInfo . fieldType) . fields
  
  isStorable :: SerialInfo -> Bool
  isStorable = isJust . size
  
  -- Add bindings for every MsgType referenced by this Msg to a Haskell
  -- type and serialization information for that type to a 'MsgInfo'
  -- context.
  addMsg :: Msg -> MsgInfo SerialMsg
  addMsg msg = do oldHome <- homePkg <$> get
                  let pkgName = B.pack $ msgPackage msg
                      sName = pack $ shortName msg
                      tName = B.concat [sName, ".", sName]
                  setHomePkg pkgName
                  flat <- isFlat msg
                  let ser = (if flat then defaultFlat else defaultNonFlat) $ tName
                  addParsedMsg pkgName sName ser msg
                  when (not (B.null oldHome)) (setHomePkg oldHome)
                  return (ser,msg)
  
  withMsg :: Msg -> MsgInfo a -> MsgInfo a
  withMsg msg action = do _ <- addMsg msg
                          oldHome <- homePkg <$> get
                          setHomePkg . B.pack $ msgPackage msg
                          r <- action
                          setHomePkg oldHome
                          return r
  
  -- Get a 'SerialInfo' value for a given 'MsgType'. If the specified
  -- 'MsgType' has not yet been parsed, it will be resolved and parsed.
  getTypeInfo :: MsgType -> MsgInfo SerialInfo
  getTypeInfo mt = do pkg <- typeCache <$> (getPackage . homePkg =<< get)
                      aux . M.lookup mt $ pkg
      where aux Nothing = addField mt
            aux (Just info) = return info
            typeCache (_,tc,_) = tc
  
  -- Default SerialInfo value for inductive (non-flat) values
  defaultNonFlat :: ByteString -> SerialInfo
  defaultNonFlat t = SerialInfo t "put" "get" Nothing
  
  -- Default SerialInfo for a flat value.
  defaultFlat :: ByteString -> SerialInfo
  defaultFlat t = SerialInfo t "put" "get" (Just si)
      where si = B.concat ["sizeOf (P.undefined::", t, ")"]
  
  -- Build a RHS for a sizeOf definition that multiplies a Storable
  -- element's size by the number of elements a FixedArray of that type
  -- contains.
  mulSize :: Int -> SerialInfo -> ByteString
  mulSize n (SerialInfo _ _ _ (Just sz)) = 
      B.concat [pack $ show n, " * (", sz, ")"]
  mulSize _ (SerialInfo _ _ _ Nothing) = 
      error "Can't generate a Storable instance for a RFixedArray\
            \with a non-Storable element type."
  
  -- NOTE: ROS specifies that we serialize booleans as a single byte,
  -- while there is an instance of Haskell's 'Storable' class for 'Bool'
  -- that uses four bytes for each boolean value. We handle individual
  -- boolean values with the 'RosBinary' instance for the Haskell 'Bool'
  -- type (i.e. one byte for each boolean). We must handle arrays of
  -- booleans specially as the serialized form must still be one byte
  -- per value.
  
  -- | Deserialization source code string to read a vector of bytes,
  -- then convert that to a vector of 'Bool's.
  getBoolFromWord :: ByteString
  getBoolFromWord = "P.fmap (V.map (P.> 0) :: V.Vector Word.Word8 \
                    \-> V.Vector P.Bool) get"
  
  -- | Serialization source code string to convert a vector of 'Bool's
  -- to a vector of bytes before serializing.
  putWordFromBool :: ByteString
  putWordFromBool = "(put . (V.map (P.fromIntegral . P.fromEnum)\
                    \ :: V.Vector P.Bool -> V.Vector Word.Word8))"
  
  -- Add a SerialInfo value for a 'MsgType' to the 'MsgInfo' context.
  addField :: MsgType -> MsgInfo SerialInfo
  addField RString = ros2Hask RString >>= setTypeInfo RString . defaultNonFlat
  addField t@(RVarArray RBool) =
    do lst <- vecOf RBool
       let arr = SerialInfo lst putWordFromBool getBoolFromWord Nothing
       setTypeInfo t arr
  addField t@(RVarArray el) = 
      do elInfo <- getTypeInfo el
         arr <- if isStorable elInfo
                then do lst <- vecOf el
                        return $ SerialInfo lst "put" "get" Nothing
                else do lst <- listOf el
                        return $ SerialInfo lst "putList" "getList" Nothing
         setTypeInfo t arr
  addField t@(RFixedArray n RBool) =
    do lst <- vecOf RBool
       let arr = SerialInfo lst putWordFromBool getBoolFromWord
                            (Just . B.pack $ show n)
       setTypeInfo t arr
  addField t@(RFixedArray n el) = 
      do elInfo <- getTypeInfo el
         arr <- if isStorable elInfo
                then do lst <- vecOf el
                        return $ SerialInfo lst "put" "get" 
                                            (Just $ mulSize n elInfo)
                else do lst <- listOf el
                        return $ SerialInfo lst "putFixedList" "getFixedList"
                                            Nothing
         setTypeInfo t arr
  addField t@(RUserType n) = do userFlat <- isStorable . fst <$> getMsg n
                                t' <- ros2Hask t
                                setTypeInfo t $ (if userFlat
                                                 then defaultFlat
                                                 else defaultNonFlat) t'
  addField t = ros2Hask t >>= setTypeInfo t . defaultFlat
  
  -- Generate the name of the Haskell type that corresponds to a flat
  -- (i.e. non-array) ROS type.
  mkFlatType :: MsgType -> ByteString
  mkFlatType RBool         = "P.Bool"
  mkFlatType RInt8         = "Int.Int8"
  mkFlatType RUInt8        = "Word.Word8"
  mkFlatType RByte         = "Word.Word8"
  mkFlatType RChar         = "Int.Int8"
  mkFlatType RInt16        = "Int.Int16"
  mkFlatType RUInt16       = "Word.Word16"
  mkFlatType RInt32        = "P.Int"
  mkFlatType RUInt32       = "Word.Word32"
  mkFlatType RInt64        = "Int.Int64"
  mkFlatType RUInt64       = "Word.Word64"
  mkFlatType RFloat32      = "P.Float"
  mkFlatType RFloat64      = "P.Double"
  mkFlatType RTime         = "ROSTime"
  mkFlatType RDuration     = "ROSDuration"
  mkFlatType (RUserType t) = qualify . pack . takeFileName . unpack $ t
      where qualify b = B.concat [b, ".", b]
  mkFlatType t             = error $ show t ++ " is not a flat type"
  
  -- Given a home package name and a ROS 'MsgType', generate a Haskell
  -- type name.
  ros2Hask :: MsgType -> MsgInfo ByteString
  ros2Hask (RFixedArray _ t) = mkArrayType t
  ros2Hask (RVarArray t)     = mkArrayType t
  ros2Hask RString           = return "P.String"
  ros2Hask t                 = return $ mkFlatType t
  
  vecOf :: MsgType -> MsgInfo ByteString
  vecOf = getTypeInfo >=> return . B.append "V.Vector " . hType
  
  listOf :: MsgType -> MsgInfo ByteString
  listOf = getTypeInfo >=> return . buildString . hType
      where buildString t = B.concat ["[", t, "]"]
  
  -- Make an array type declaration. If the element type of the
  -- collection is flat (i.e. does not itself have a field of type
  -- 'RVarArray'), then it will have a 'Storable' instance and can be
  -- stored in a 'Vector'.
  mkArrayType :: MsgType -> MsgInfo ByteString
  mkArrayType t = toArr . isStorable =<< getTypeInfo t
      where toArr True = vecOf t
            toArr False = listOf t
  