Transposed the matrix for pointCount to account for a change in OpenCV (somewhere between 2.4.3 and 2.4.5)

{-#LANGUAGE ForeignFunctionInterface, ScopedTypeVariables #-}
#include "cvWrapLEO.h"
-- | This module exposes opencv functions for camera calibration using a chessboard rig. This module follows opencv quite closely and the best documentation
--   is probably found there. As quick example however, the following program detects and draws chessboard corners from an image.
--
-- @
-- module Main where
-- import CV.Image
-- import CV.Calibration
--
-- main = do
--     Just i <- loadColorImage \"chess.png\"
--     let corners = findChessboardCorners (unsafeImageTo8Bit i) (4,5) (FastCheck:defaultFlags)
--     let y = drawChessboardCorners (unsafeImageTo8Bit i) (4,5) corners
--     mapM_ print (corners)
--     saveImage \"found_chessboard.png\" y
-- @
module CV.Calibration
    (
     -- * Finding chessboard calibration rig
     FindFlags(..)
    ,defaultFlags
    ,findChessboardCorners
    ,refineChessboardCorners
    -- * Visualization
    ,drawChessboardCorners
    -- * Camera calibration
    ,calibrateCamera2
    -- * Rectification
    ,stereoRectifyUncalibrated
    ,findFundamentalMat
    ,c'CV_FM_7POINT
    ,c'CV_FM_8POINT
    ,c'CV_FM_RANSAC
    ,c'CV_FM_LMEDS
    ) where
{-#OPTIONS-GHC -fwarn-unused-imports #-}

import Foreign.C.Types
import Foreign.C.String
import Foreign.ForeignPtr
import Foreign.Storable
import Foreign.Marshal.Array
import Foreign.Marshal.Utils
import Foreign.Ptr
import Data.Bits

import CV.Image

import System.IO.Unsafe
import Utils.Point
import Control.Applicative

import CV.Matrix
import CV.Bindings.Calibrate
import CV.Bindings.Types

{#import CV.Image#}

#c
enum FindFlags {
     AdaptiveThresh  = CV_CALIB_CB_ADAPTIVE_THRESH
    ,NormalizeImage  = CV_CALIB_CB_NORMALIZE_IMAGE
    ,FilterQuads     = CV_CALIB_CB_FILTER_QUADS
    ,FastCheck       = CV_CALIB_CB_FAST_CHECK
    };
#endc

-- | Flags for the chessboard corner detector. See opencv documentation for cvFindChessboardCorners.
{#enum FindFlags {}#}

flagsToNum fs = foldl (.|.) 0 $ map (fromIntegral . fromEnum) fs

-- |Default flags for finding corners
defaultFlags :: [FindFlags]
defaultFlags = [AdaptiveThresh]

-- | Find the inner corners of a chessboard in a given image.
findChessboardCorners :: CV.Image.Image RGB D8 -> (Int, Int) -> [FindFlags] -> [(Float,Float)]
findChessboardCorners image (w,h) flags =
   unsafePerformIO $
    with 1 $ \(c_corner_count::Ptr CInt) ->
     allocaArray len $ \(c_corners :: Ptr CvPoint )->
      withGenImage image $ \c_image -> do
        r <- {#call wrapFindChessBoardCorners#} c_image (fromIntegral w) (fromIntegral h)
                                           (castPtr c_corners) c_corner_count
                                           (flagsToNum flags)
        count <- peek c_corner_count
        arr <- peekArray (fromIntegral count) c_corners
        return (map cvPt2Pt arr)
  where len = w*h

-- |Given an estimate of chessboard corners, provide a subpixel estimation of actual corners.
refineChessboardCorners :: Image GrayScale D8 -> [(Float,Float)] -> (Int,Int) -> (Int,Int) -> [(Float,Float)]
refineChessboardCorners img pts (winW,winH) (zeroW,zeroH) = unsafePerformIO $ do
    with 1 $ \(c_corner_count::Ptr CInt) ->
      withImage img $ \c_img ->
      withArray (map mkPt pts) $ \(c_corners :: Ptr C'CvPoint2D32f ) -> do
        c'wrapFindCornerSubPix c_img c_corners (length pts) winW winH zeroW zeroH tType maxIter epsilon
        map fromPt `fmap` peekArray (length pts) c_corners
 where
    mkPt (x,y) = C'CvPoint2D32f (realToFrac x) (realToFrac y)
    fromPt (C'CvPoint2D32f x y) = (realToFrac x,realToFrac y)
    tType = c'CV_TERMCRIT_ITER
    maxIter = 100
    epsilon = 0.01

-- | Draw the found chessboard corners to an image
drawChessboardCorners :: CV.Image.Image RGB D8 -> (Int, Int) -> [(Float,Float)] -> CV.Image.Image RGB D8
drawChessboardCorners image (w,h) corners =
   unsafePerformIO $
    withCloneValue image $ \clone ->
     withArray (map pt2CvPt corners) $ \(c_corners :: Ptr CvPoint )->
      withGenImage clone$ \c_image -> do
        r <- {#call wrapDrawChessBoardCorners#} c_image (fromIntegral w) (fromIntegral h)
                                           (castPtr c_corners) (fromIntegral $ length corners)
                                           (found)
        return clone
  where
    len = w*h
    found | (w*h) == length corners = 1
          | otherwise = 0

newtype CvPoint = CvPt (CFloat,CFloat) deriving (Show)
cvPt2Pt (CvPt (a,b)) = (realToFrac a , realToFrac b)
pt2CvPt (a,b) = CvPt (realToFrac a , realToFrac b)

instance Storable CvPoint where
  sizeOf _ = {#sizeof CvPoint #}
  alignment _ = {#alignof CvPoint2D32f #}
  peek p = CvPt <$> ((,)
    <$> {#get CvPoint2D32f->x #} p
    <*> {#get CvPoint2D32f->y #} p)
  poke p (CvPt (hx,hy)) = do
    {#set CvPoint2D32f.x #} p (hx)
    {#set CvPoint2D32f.y #} p (hy)

-- | See opencv function cvCalibrateCamera2. This function takes a list of world-screen coordinate pairs acquired with find-chessboard corners
--   and attempts to find the camera parameters for the system. It returns the fitting error, the camera matrix, list of distortion co-efficients
--   and rotation and translation vectors for each coordinate pair.
calibrateCamera2 ::
     [[((Float, Float, Float), (Float, Float))]]
     -> (Int, Int)
     -> IO (Double, Matrix Float, [[Float]], [[Float]], [[Float]])
calibrateCamera2 views (w,h) = do
    let
        pointCounts :: Matrix Int
        pointCounts  = fromList (length views,1) (map (length) views)
        m = length views
        totalPts = length (concat views)
        objectPoints :: Matrix Float
        objectPoints = fromList (3,totalPts) $ concat [[x,y,z] | ((x,y,z),_) <- concat views]
        imagePoints :: Matrix Float
        imagePoints  = fromList (2,totalPts) $ concat [[x,y]   | (_,(x,y))   <- concat views]
        flags = c'CV_CALIB_FIX_K1
                .|.  c'CV_CALIB_FIX_K1
                .|.  c'CV_CALIB_FIX_K2
                .|.  c'CV_CALIB_FIX_K3
                .|.  c'CV_CALIB_FIX_K4
                .|.  c'CV_CALIB_FIX_K5
                .|.  c'CV_CALIB_FIX_K6
                .|.  c'CV_CALIB_ZERO_TANGENT_DIST

        size = C'CvSize (fromIntegral w) (fromIntegral h)
        cameraMatrix,distCoeffs,rvecs,tvecs :: Matrix Float
        cameraMatrix = emptyMatrix (3,3)
        distCoeffs   = emptyMatrix (1,8)
        rvecs        = emptyMatrix (m,3)
        tvecs        = emptyMatrix (m,3)

    err <- with size $ \c_size ->
     withMatPtr objectPoints $ \c_objectPoints ->
     withMatPtr imagePoints $ \c_imagePoints ->
     withMatPtr pointCounts $ \c_pointCounts ->
     withMatPtr cameraMatrix $ \c_cameraMatrix ->
     withMatPtr distCoeffs $ \c_distCoeffs ->
     withMatPtr rvecs $ \c_rvecs ->
     withMatPtr tvecs $ \c_tvecs ->
      c'wrapCalibrateCamera2 c_objectPoints c_imagePoints c_pointCounts c_size
                             c_cameraMatrix c_distCoeffs c_rvecs c_tvecs flags

    -- print ( objectPoints, imagePoints, pointCounts,cameraMatrix, distCoeffs, rvecs, tvecs )
    return (err, transpose cameraMatrix, toCols distCoeffs, toCols rvecs, toCols tvecs)

stereoRectifyUncalibrated :: Matrix (Float,Float) -> Matrix (Float,Float) -> Matrix Float -> (Int,Int) -> Double
                             -> (Matrix Float, Matrix Float)
stereoRectifyUncalibrated pts1 pts2 fund (w,h) threshold = unsafePerformIO $
    let h1 = emptyMatrix (3,3)
        h2 = emptyMatrix (3,3)
    in withMatPtr pts1     $ \c_pts1 ->
       withMatPtr pts2     $ \c_pts2 ->
       withMatPtr fund     $ \c_fund ->
       withMatPtr h1       $ \c_h1 ->
       withMatPtr h2       $ \c_h2 ->
       with (C'CvSize (fromIntegral w) (fromIntegral h)) $ \c_size -> do
        r <- c'wrapStereoRectifyUncalibrated c_pts1 c_pts2 c_fund c_size c_h1 c_h2 (realToFrac threshold)
        return (h1, h2)

findFundamentalMat :: Matrix Float -> Matrix Float -> CInt -> Double -> Double
                             -> Matrix Float --, Matrix Float)
findFundamentalMat pts1 pts2  method p1 p2 = unsafePerformIO $
    let fund = emptyMatrix (3,3)
    in withMatPtr pts1     $ \c_pts1 ->
       withMatPtr pts2     $ \c_pts2 ->
       withMatPtr fund     $ \c_fund ->
        do
         r <- c'cvFindFundamentalMat c_pts1 c_pts2 c_fund method (realToFrac p1) (realToFrac p2) nullPtr
         return (fund)

stereoCalibrate ::  Matrix (Float,Float,Float) -- ^ Object points
                 -> Matrix (Float,Float)       -- ^ ImagePoints 1
                 -> Matrix (Float,Float)       -- ^ ImagePoints 2
                 -> Matrix (Int)               -- ^ point counts
                 -> Matrix Float               -- ^ camera matrix 1
                 -> Matrix Float               -- ^ camera matrix 2
                 -> Matrix Float               -- ^ dist coeffs 1
                 -> Matrix Float               -- ^ dist coeffs 2
                 -> (Int,Int)                  -- ^ image size
                 -> C'CvTermCriteria           -- ^ stopping condition
                 -> CInt                       -- ^ Flags
                 -> (Double, Matrix Float, Matrix Float ,Matrix Float, Matrix Float)
stereoCalibrate obj im1 im2 counts cam1 cam2 dist1 dist2 (w,h) tc flags =  unsafePerformIO $ do
    r <- CV.Matrix.create (3,3)
    e <- CV.Matrix.create (3,3)
    t <- CV.Matrix.create (3,1)
    f <- CV.Matrix.create (3,3)

    withMatPtr obj     $ \c_obj ->
       withMatPtr im1     $ \c_im1 ->
       withMatPtr im2     $ \c_im2 ->
       withMatPtr counts     $ \c_counts ->
       withMatPtr cam1     $ \c_cam1 ->
       withMatPtr cam2     $ \c_cam2 ->
       withMatPtr dist1     $ \c_dist1 ->
       withMatPtr dist2     $ \c_dist2 ->
       withMatPtr r     $ \c_r ->
       withMatPtr t     $ \c_t ->
       withMatPtr e     $ \c_e ->
       withMatPtr f     $ \c_f ->
       with tc          $ \c_tc ->
       with (CV.Bindings.Types.C'CvSize (fromIntegral w) (fromIntegral h)) $ \c_size -> do
        pr <- c'wrapStereoCalibrate
            c_obj
            c_im1
            c_im2
            c_counts
            c_cam1
            c_dist1
            c_cam2
            c_dist2
            c_size
            c_r
            c_t
            c_e
            c_f
            c_tc
            flags
        return (realToFrac pr,r,t,e,f)

{-#LANGUAGE ForeignFunctionInterface, ScopedTypeVariables #-}
#include "cvWrapLEO.h"
-- | This module exposes opencv functions for camera calibration using a chessboard rig. This module follows opencv quite closely and the best documentation
--   is probably found there. As quick example however, the following program detects and draws chessboard corners from an image.
--
-- @
-- module Main where
-- import CV.Image
-- import CV.Calibration
--
-- main = do
--     Just i <- loadColorImage \"chess.png\"
--     let corners = findChessboardCorners (unsafeImageTo8Bit i) (4,5) (FastCheck:defaultFlags)
--     let y = drawChessboardCorners (unsafeImageTo8Bit i) (4,5) corners
--     mapM_ print (corners)
--     saveImage \"found_chessboard.png\" y
-- @
module CV.Calibration
    (
     -- * Finding chessboard calibration rig
     FindFlags(..)
    ,defaultFlags
    ,findChessboardCorners
    ,refineChessboardCorners
    -- * Visualization
    ,drawChessboardCorners
    -- * Camera calibration
    ,calibrateCamera2
    -- * Rectification
    ,stereoRectifyUncalibrated
    ,findFundamentalMat
    ,c'CV_FM_7POINT
    ,c'CV_FM_8POINT
    ,c'CV_FM_RANSAC
    ,c'CV_FM_LMEDS
    ) where
{-#OPTIONS-GHC -fwarn-unused-imports #-}

import Foreign.C.Types
import Foreign.C.String
import Foreign.ForeignPtr
import Foreign.Storable
import Foreign.Marshal.Array
import Foreign.Marshal.Utils
import Foreign.Ptr
import Data.Bits

import CV.Image

import System.IO.Unsafe
import Utils.Point
import Control.Applicative

import CV.Matrix
import CV.Bindings.Calibrate
import CV.Bindings.Types

{#import CV.Image#}

#c
enum FindFlags {
     AdaptiveThresh  = CV_CALIB_CB_ADAPTIVE_THRESH
    ,NormalizeImage  = CV_CALIB_CB_NORMALIZE_IMAGE
    ,FilterQuads     = CV_CALIB_CB_FILTER_QUADS
    ,FastCheck       = CV_CALIB_CB_FAST_CHECK
    };
#endc

-- | Flags for the chessboard corner detector. See opencv documentation for cvFindChessboardCorners.
{#enum FindFlags {}#}

flagsToNum fs = foldl (.|.) 0 $ map (fromIntegral . fromEnum) fs

-- |Default flags for finding corners
defaultFlags :: [FindFlags]
defaultFlags = [AdaptiveThresh]

-- | Find the inner corners of a chessboard in a given image.
findChessboardCorners :: CV.Image.Image RGB D8 -> (Int, Int) -> [FindFlags] -> [(Float,Float)]
findChessboardCorners image (w,h) flags =
   unsafePerformIO $
    with 1 $ \(c_corner_count::Ptr CInt) ->
     allocaArray len $ \(c_corners :: Ptr CvPoint )->
      withGenImage image $ \c_image -> do
        r <- {#call wrapFindChessBoardCorners#} c_image (fromIntegral w) (fromIntegral h)
                                           (castPtr c_corners) c_corner_count
                                           (flagsToNum flags)
        count <- peek c_corner_count
        arr <- peekArray (fromIntegral count) c_corners
        return (map cvPt2Pt arr)
  where len = w*h

-- |Given an estimate of chessboard corners, provide a subpixel estimation of actual corners.
refineChessboardCorners :: Image GrayScale D8 -> [(Float,Float)] -> (Int,Int) -> (Int,Int) -> [(Float,Float)]
refineChessboardCorners img pts (winW,winH) (zeroW,zeroH) = unsafePerformIO $ do
    with 1 $ \(c_corner_count::Ptr CInt) ->
      withImage img $ \c_img ->
      withArray (map mkPt pts) $ \(c_corners :: Ptr C'CvPoint2D32f ) -> do
        c'wrapFindCornerSubPix c_img c_corners (length pts) winW winH zeroW zeroH tType maxIter epsilon
        map fromPt `fmap` peekArray (length pts) c_corners
 where
    mkPt (x,y) = C'CvPoint2D32f (realToFrac x) (realToFrac y)
    fromPt (C'CvPoint2D32f x y) = (realToFrac x,realToFrac y)
    tType = c'CV_TERMCRIT_ITER
    maxIter = 100
    epsilon = 0.01

-- | Draw the found chessboard corners to an image
drawChessboardCorners :: CV.Image.Image RGB D8 -> (Int, Int) -> [(Float,Float)] -> CV.Image.Image RGB D8
drawChessboardCorners image (w,h) corners =
   unsafePerformIO $
    withCloneValue image $ \clone ->
     withArray (map pt2CvPt corners) $ \(c_corners :: Ptr CvPoint )->
      withGenImage clone$ \c_image -> do
        r <- {#call wrapDrawChessBoardCorners#} c_image (fromIntegral w) (fromIntegral h)
                                           (castPtr c_corners) (fromIntegral $ length corners)
                                           (found)
        return clone
  where
    len = w*h
    found | (w*h) == length corners = 1
          | otherwise = 0

newtype CvPoint = CvPt (CFloat,CFloat) deriving (Show)
cvPt2Pt (CvPt (a,b)) = (realToFrac a , realToFrac b)
pt2CvPt (a,b) = CvPt (realToFrac a , realToFrac b)

instance Storable CvPoint where
  sizeOf _ = {#sizeof CvPoint #}
  alignment _ = {#alignof CvPoint2D32f #}
  peek p = CvPt <$> ((,)
    <$> {#get CvPoint2D32f->x #} p
    <*> {#get CvPoint2D32f->y #} p)
  poke p (CvPt (hx,hy)) = do
    {#set CvPoint2D32f.x #} p (hx)
    {#set CvPoint2D32f.y #} p (hy)

-- | See opencv function cvCalibrateCamera2. This function takes a list of world-screen coordinate pairs acquired with find-chessboard corners
--   and attempts to find the camera parameters for the system. It returns the fitting error, the camera matrix, list of distortion co-efficients
--   and rotation and translation vectors for each coordinate pair.
calibrateCamera2 ::
     [[((Float, Float, Float), (Float, Float))]]
     -> (Int, Int)
     -> IO (Double, Matrix Float, [[Float]], [[Float]], [[Float]])
calibrateCamera2 views (w,h) = do
    let
        pointCounts :: Matrix Int
        pointCounts  = fromList (1,length views) (map (length) views)
        m = length views
        totalPts = length (concat views)
        objectPoints :: Matrix Float
        objectPoints = fromList (3,totalPts) $ concat [[x,y,z] | ((x,y,z),_) <- concat views]
        imagePoints :: Matrix Float
        imagePoints  = fromList (2,totalPts) $ concat [[x,y]   | (_,(x,y))   <- concat views]
        flags = c'CV_CALIB_FIX_K1
                .|.  c'CV_CALIB_FIX_K1
                .|.  c'CV_CALIB_FIX_K2
                .|.  c'CV_CALIB_FIX_K3
                .|.  c'CV_CALIB_FIX_K4
                .|.  c'CV_CALIB_FIX_K5
                .|.  c'CV_CALIB_FIX_K6
                .|.  c'CV_CALIB_ZERO_TANGENT_DIST

        size = C'CvSize (fromIntegral w) (fromIntegral h)
        cameraMatrix,distCoeffs,rvecs,tvecs :: Matrix Float
        cameraMatrix = emptyMatrix (3,3)
        distCoeffs   = emptyMatrix (1,8)
        rvecs        = emptyMatrix (m,3)
        tvecs        = emptyMatrix (m,3)

    err <- with size $ \c_size ->
     withMatPtr objectPoints $ \c_objectPoints ->
     withMatPtr imagePoints $ \c_imagePoints ->
     withMatPtr pointCounts $ \c_pointCounts ->
     withMatPtr cameraMatrix $ \c_cameraMatrix ->
     withMatPtr distCoeffs $ \c_distCoeffs ->
     withMatPtr rvecs $ \c_rvecs ->
     withMatPtr tvecs $ \c_tvecs ->
      c'wrapCalibrateCamera2 c_objectPoints c_imagePoints c_pointCounts c_size
                             c_cameraMatrix c_distCoeffs c_rvecs c_tvecs flags

    -- print ( objectPoints, imagePoints, pointCounts,cameraMatrix, distCoeffs, rvecs, tvecs )
    return (err, transpose cameraMatrix, toCols distCoeffs, toCols rvecs, toCols tvecs)

stereoRectifyUncalibrated :: Matrix (Float,Float) -> Matrix (Float,Float) -> Matrix Float -> (Int,Int) -> Double
                             -> (Matrix Float, Matrix Float)
stereoRectifyUncalibrated pts1 pts2 fund (w,h) threshold = unsafePerformIO $
    let h1 = emptyMatrix (3,3)
        h2 = emptyMatrix (3,3)
    in withMatPtr pts1     $ \c_pts1 ->
       withMatPtr pts2     $ \c_pts2 ->
       withMatPtr fund     $ \c_fund ->
       withMatPtr h1       $ \c_h1 ->
       withMatPtr h2       $ \c_h2 ->
       with (C'CvSize (fromIntegral w) (fromIntegral h)) $ \c_size -> do
        r <- c'wrapStereoRectifyUncalibrated c_pts1 c_pts2 c_fund c_size c_h1 c_h2 (realToFrac threshold)
        return (h1, h2)

findFundamentalMat :: Matrix Float -> Matrix Float -> CInt -> Double -> Double
                             -> Matrix Float --, Matrix Float)
findFundamentalMat pts1 pts2  method p1 p2 = unsafePerformIO $
    let fund = emptyMatrix (3,3)
    in withMatPtr pts1     $ \c_pts1 ->
       withMatPtr pts2     $ \c_pts2 ->
       withMatPtr fund     $ \c_fund ->
        do
         r <- c'cvFindFundamentalMat c_pts1 c_pts2 c_fund method (realToFrac p1) (realToFrac p2) nullPtr
         return (fund)

stereoCalibrate ::  Matrix (Float,Float,Float) -- ^ Object points
                 -> Matrix (Float,Float)       -- ^ ImagePoints 1
                 -> Matrix (Float,Float)       -- ^ ImagePoints 2
                 -> Matrix (Int)               -- ^ point counts
                 -> Matrix Float               -- ^ camera matrix 1
                 -> Matrix Float               -- ^ camera matrix 2
                 -> Matrix Float               -- ^ dist coeffs 1
                 -> Matrix Float               -- ^ dist coeffs 2
                 -> (Int,Int)                  -- ^ image size
                 -> C'CvTermCriteria           -- ^ stopping condition
                 -> CInt                       -- ^ Flags
                 -> (Double, Matrix Float, Matrix Float ,Matrix Float, Matrix Float)
stereoCalibrate obj im1 im2 counts cam1 cam2 dist1 dist2 (w,h) tc flags =  unsafePerformIO $ do
    r <- CV.Matrix.create (3,3)
    e <- CV.Matrix.create (3,3)
    t <- CV.Matrix.create (3,1)
    f <- CV.Matrix.create (3,3)

    withMatPtr obj     $ \c_obj ->
       withMatPtr im1     $ \c_im1 ->
       withMatPtr im2     $ \c_im2 ->
       withMatPtr counts     $ \c_counts ->
       withMatPtr cam1     $ \c_cam1 ->
       withMatPtr cam2     $ \c_cam2 ->
       withMatPtr dist1     $ \c_dist1 ->
       withMatPtr dist2     $ \c_dist2 ->
       withMatPtr r     $ \c_r ->
       withMatPtr t     $ \c_t ->
       withMatPtr e     $ \c_e ->
       withMatPtr f     $ \c_f ->
       with tc          $ \c_tc ->
       with (CV.Bindings.Types.C'CvSize (fromIntegral w) (fromIntegral h)) $ \c_size -> do
        pr <- c'wrapStereoCalibrate
            c_obj
            c_im1
            c_im2
            c_counts
            c_cam1
            c_dist1
            c_cam2
            c_dist2
            c_size
            c_r
            c_t
            c_e
            c_f
            c_tc
            flags
        return (realToFrac pr,r,t,e,f)


remove obsolete utility

remove obsolete utility

Ignore-this: 759c2e07828e43b1e2ad74c866fbc61d

darcs-hash:20091217131844-9241b-26756e5abbafaf9c5a33f9dd913b9228f1f33560.gz

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Utils
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Utility functions
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Utils
  (
    driverVersion
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C

-- |
-- Return the version number of the installed CUDA driver
--
driverVersion :: IO Int
driverVersion =  resultIfOk =<< cuDriverGetVersion

{# fun unsafe cuDriverGetVersion
  { alloca- `Int' peekIntConv* } -> `Status' cToEnum #}

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
-- |
-- Module    : Foreign.CUDA.Driver.Utils
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Utility functions
--
--------------------------------------------------------------------------------

module Foreign.CUDA.Driver.Utils
  (
    forceEither, driverVersion
  )
  where

#include <cuda.h>
{# context lib="cuda" #}

-- Friends
import Foreign.CUDA.Driver.Error
import Foreign.CUDA.Internal.C2HS

-- System
import Foreign
import Foreign.C

-- |
-- Unwrap an Either construct, throwing an error for non-right values
--
forceEither :: Either String a -> a
forceEither (Left  s) = error s
forceEither (Right r) = r

-- |
-- Return the version number of the installed CUDA driver
--
driverVersion :: IO Int
driverVersion =  resultIfOk =<< cuDriverGetVersion

{# fun unsafe cuDriverGetVersion
  { alloca- `Int' peekIntConv* } -> `Status' cToEnum #}


Remove fixed TODO.

Remove fixed TODO.

-- Copyright (c) 2016, Google Inc.
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--     * Redistributions of source code must retain the above copyright
--       notice, this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of Google Inc. nor the
--       names of its contributors may be used to endorse or promote products
--       derived from this software without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
-- ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-- WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL Google Inc. BE LIABLE FOR ANY
-- DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-- (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
-- LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
-- ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-- (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-- SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--
--------------------------------------------------------------------------------
{-# LANGUAGE ScopedTypeVariables, OverloadedStrings, ExistentialQuantification, RecordWildCards #-}
module Network.Grpc.Core.Call where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad

import qualified Data.ByteString              as B
import qualified Data.ByteString.Lazy         as L
import           Data.Monoid ((<>), Last(..))
import           Data.IORef

import           Foreign.C.Types              as C
import qualified Foreign.Marshal.Alloc        as C
import qualified Foreign.Marshal.Utils        as C
import qualified Foreign.Ptr                  as C
import           Foreign.Ptr (Ptr)
import qualified Foreign.Storable             as C

-- transformers
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.Except
import           Control.Monad.Trans.Reader

import qualified Network.Grpc.CompletionQueue as CQ
import Network.Grpc.Lib.PropagationBits
{#import Network.Grpc.Lib.ByteBuffer#}
{#import Network.Grpc.Lib.Metadata#}
{#import Network.Grpc.Lib.TimeSpec#}
{#import Network.Grpc.Lib.Core#}

#include <grpc/grpc.h>
#include "hs_grpc.h"

{#context lib = "grpc" prefix = "grpc" #}

data Deadline
  = AbsoluteDeadline TimeSpec
  | RelativeDeadline Int --seconds

data ClientContext = ClientContext
  { ccChannel :: Channel
  , ccCQ :: CompletionQueue
  , ccWorker :: CQ.Worker
  }

data CallOptions = CallOptions
  { coDeadline :: Maybe Deadline
  , coParentContext :: Maybe () -- todo
  , coPropagationMask :: Maybe PropagationMask
  , coMetadata :: [Metadata]
  }

instance Monoid CallOptions where
  mempty = CallOptions Nothing Nothing Nothing []
  mappend (CallOptions a b c d) (CallOptions a' b' c' d') =
    CallOptions
      (getLast (Last a <> Last a'))
      (getLast (Last b <> Last b'))
      (getLast (Last c <> Last c'))
      (d <> d')

withAbsoluteDeadline :: TimeSpec -> CallOptions
withAbsoluteDeadline deadline =
  mempty { coDeadline = Just (AbsoluteDeadline deadline) }

withRelativeDeadlineSeconds :: Int -> CallOptions
withRelativeDeadlineSeconds seconds =
  mempty { coDeadline = Just (RelativeDeadline seconds) }

withParentContext :: () -> CallOptions
withParentContext ctx =
  mempty { coParentContext = Just ctx }

withParentContextPropagating :: () -> PropagationMask -> CallOptions
withParentContextPropagating ctx prop =
  mempty { coParentContext   = Just ctx
         , coPropagationMask = Just prop }

withMetadata :: [Metadata] -> CallOptions
withMetadata md =
  mempty { coMetadata = md }

resolveDeadline :: CallOptions -> IO TimeSpec
resolveDeadline co =
  case coDeadline co of
    Nothing -> return gprInfFuture
    Just (AbsoluteDeadline deadline) -> return deadline
    Just (RelativeDeadline seconds) ->
      secondsFromNow (fromIntegral seconds)

newClientContext :: Channel -> IO ClientContext
newClientContext chan = do
  cq <- completionQueueCreate reservedPtr
  cqt <- CQ.startCompletionQueueThread cq
  return (ClientContext chan cq cqt)

destroyClientContext :: ClientContext -> IO ()
destroyClientContext (ClientContext _ cq w) = do
  completionQueueShutdown cq
  CQ.waitWorkerTermination w

type MethodName = B.ByteString
type Arg = B.ByteString

-- | Run the IO function once, cache it in the MVar.
onceMVar :: MVar (Maybe a) -> IO a -> IO a
onceMVar mvar io = modifyMVar mvar $ \x ->
  case x of
    Just x' -> return (x, x')
    Nothing -> do
      x' <- io
      return (Just x', x')

reservedPtr :: Ptr ()
reservedPtr = C.nullPtr

data UnaryResult a = UnaryResult [Metadata] [Metadata] a deriving Show

callUnary :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (UnaryResult L.ByteString))
callUnary ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  bracket (grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline) grpcCallDestroy $ \call0 -> newMVar call0 >>= \mcall -> do
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata (coMetadata co)
    sendCloseOp <- opSendCloseFromClient
    sendMessageOp <- opSendMessage arg
    recvStatusOp <- opRecvStatusOnClient crw
    recvMessageOp <- opRecvMessage
    recvInitialMetadataOp <- opRecvInitialMetadata crw

    res <- callBatch crw [
          OpX sendInitOp
        , OpX sendCloseOp
        , OpX recvInitialMetadataOp
        , OpX recvMessageOp
        , OpX sendMessageOp
        , OpX recvStatusOp
        ]
    case res of
      RpcOk _ -> do
        (RpcStatus trailMD status statusDetails) <- opRead recvStatusOp
        case status of
          StatusOk -> do
            initMD <- opRead recvInitialMetadataOp
            answ <- opRead recvMessageOp
            let answ' = maybe L.empty id answ
            return (RpcOk (UnaryResult initMD trailMD answ'))
          _ -> return (RpcError (StatusError status statusDetails))
      RpcError err -> return (RpcError err)

callDownstream :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (Client B.ByteString L.ByteString))
callDownstream ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)
  sendCloseOp <- opSendCloseFromClient
  sendMessageOp <- opSendMessage arg

  res <- callBatch crw [
        OpX sendInitOp
      , OpX sendCloseOp
      , OpX sendMessageOp
      ]
  case res of
    RpcOk _ -> do
      return (RpcOk (Client crw defaultEncoder defaultDecoder))
    RpcError err -> return (RpcError err)

callUpstream :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callUpstream ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> do
      return (RpcOk (Client crw defaultEncoder defaultDecoder))
    RpcError err -> return (RpcError err)

callBidi :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callBidi ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline >>= newMVar

  crw <- newClientReaderWriter ctx mcall
  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
    RpcError err -> return (RpcError err)

data Client req resp = Client
  { clientCrw :: ClientReaderWriter
  , clientEncoder :: Encoder req
  , clientDecoder :: Decoder resp
  }

type Decoder a = L.ByteString -> IO (RpcReply a)
type Encoder a = a -> IO (RpcReply B.ByteString)

defaultDecoder :: L.ByteString -> IO (RpcReply L.ByteString)
defaultDecoder bs = return (RpcOk bs)

defaultEncoder :: B.ByteString -> IO (RpcReply B.ByteString)
defaultEncoder bs = return (RpcOk bs)

data ClientReaderWriter = ClientReaderWriter {
  context :: ClientContext,
  callMVar_ :: MVar Call,
  initialMDRef :: !(IORef (Maybe [Metadata])),
  trailingMDRef :: !(IORef (Maybe [Metadata])),
  statusFromServer :: !(IORef (Maybe RpcStatus))
}

newClientReaderWriter :: ClientContext -> MVar Call -> IO ClientReaderWriter
newClientReaderWriter ctx mcall = do
  initMD <- newIORef Nothing
  trailMD <- newIORef Nothing
  status <- newIORef Nothing
  return (ClientReaderWriter ctx mcall initMD trailMD status)

data OpX = forall t. OpX (OpT t)

data OpArray = OpArray { opArrPtr :: !GrpcOpPtr
                       , opArrLen :: !CULong
                       , opArrFree :: !(IO ())
                       }

toArray :: [OpX] -> IO OpArray
toArray ops = do
  aptr <- C.mallocBytes (length ops * {#sizeof grpc_op#})
  let ptrs = iterate (`C.plusPtr` {#sizeof grpc_op#}) aptr
      ops' = concatMap (\(OpX op) -> opAdd op) ops
      free = C.free aptr >> sequence_ (map (\(OpX op) -> opFinish op) ops)
      write op p = do
        {#set grpc_op->flags#} p 0
        {#set grpc_op->reserved#} p C.nullPtr
        op p
  zipWithM_ write ops' ptrs
  return $! OpArray aptr (fromIntegral $ length ops) free

callBatch :: ClientReaderWriter -> [OpX] -> IO (RpcReply ())
callBatch crw ops = do
  let ctx = context crw
  CQ.withEvent (ccWorker ctx) $ \eDesc -> do
    arr <- toArray ops
    callStatus <- withMVar (callMVar_ crw) $ \call ->
      grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
    -- print ("callStatus: " ++ show callStatus)
    case callStatus of
      CallOk -> do
        e <- CQ.interruptibleWaitEvent eDesc
        -- print ("event: " ++ show e)
        opArrFree arr -- TODO: We leak if we're interrupted.
        case e of
          QueueOpComplete OpSuccess _ -> return (RpcOk ())
          QueueOpComplete OpError _ -> return (RpcError (Error "callBatch: op error"))
          QueueTimeOut -> return (RpcError DeadlineExceeded)
          QueueShutdown -> return (RpcError (Error "queue shutdown"))
      _ -> do
         -- print callStatus
         return (RpcError (CallErrorStatus callStatus))

data OpT out = Op
  { opAdd    :: [Ptr GrpcOp -> IO ()]
  , opValue  :: IORef out
  , opFinish :: IO ()
  }

opRead :: OpT a -> IO a
opRead op = readIORef (opValue op)

opRecvInitialMetadata :: ClientReaderWriter -> IO (OpT [Metadata])
opRecvInitialMetadata crw = do
  arr <- mallocMetadataArray
  value <- newIORef (error "opRecvInitialMetadata never finished")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvInitialMetadata))
        {#set grpc_op->data.recv_initial_metadata#} p arr
      ]
    finish = do
      mds <- readMetadataArray arr
      writeIORef (initialMDRef crw) (Just mds)
      writeIORef value mds
      freeMetadataArray arr
  return (Op add value finish)

opSendMessage :: B.ByteString -> IO (OpT ())
opSendMessage bs = do
  bb <- fromByteString bs
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendMessage))
        {#set grpc_op->data.send_message#} p bb
      ]
    finish =
      byteBufferDestroy bb
  return (Op add value finish)

opRecvMessage :: IO (OpT (Maybe L.ByteString))
opRecvMessage = do
  bbptr <- C.malloc :: IO (Ptr (Ptr CByteBuffer))
  value <- newIORef Nothing
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvMessage))
        {#set grpc_op->data.recv_message#} p bbptr
      ]
    finish = do
      bb <- C.peek bbptr
      C.free bbptr
      writeIORef value =<< if bb /= C.nullPtr
        then do
          lbs <- toLazyByteString bb
          byteBufferDestroy bb
          return (Just lbs)
        else return Nothing
  return (Op add value finish)

opSendInitialMetadata :: [Metadata] -> IO (OpT ())
opSendInitialMetadata elems = do
  (mdArrPtr, free) <- mallocMetadata elems
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendInitialMetadata))
        {#set grpc_op->data.send_initial_metadata.count#} p (fromIntegral (length elems))
        {#set grpc_op->data.send_initial_metadata.metadata#} p (C.castPtr mdArrPtr)
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.is_set#} p 0
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.level#} p 0
      ]
    finish = do
      free
  return (Op add value finish)

data RpcStatus = RpcStatus [Metadata] StatusCode B.ByteString
  deriving Show

opRecvStatusOnClient :: ClientReaderWriter -> IO (OpT RpcStatus)
opRecvStatusOnClient (ClientReaderWriter{..}) = do
  trailingMetadataArrPtr <- mallocMetadataArray
  statusCodePtr <- C.malloc :: IO (Ptr StatusCodeT)
  ptrPtrStatusStr <- C.new C.nullPtr :: IO (Ptr (Ptr CChar))
  capacityPtr <- C.new 0 :: IO (Ptr SizeT)
  value <- newIORef (error "opRecvStatusOnClient never ran")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvStatusOnClient))
        {#set grpc_op->data.recv_status_on_client.trailing_metadata#} p trailingMetadataArrPtr
        {#set grpc_op->data.recv_status_on_client.status#} p statusCodePtr
        {#set grpc_op->data.recv_status_on_client.status_details#} p ptrPtrStatusStr
        {#set grpc_op->data.recv_status_on_client.status_details_capacity#} p capacityPtr
      ]
    finish = do
      trailingMd <- readMetadataArray trailingMetadataArrPtr
      statusCode <- fmap toStatusCode (C.peek statusCodePtr)
      statusDetails <- B.packCString =<< C.peek ptrPtrStatusStr
      let status = RpcStatus trailingMd statusCode statusDetails
      writeIORef value status
      writeIORef statusFromServer (Just status)
      free
    free = do
      freeMetadataArray trailingMetadataArrPtr
      C.free statusCodePtr
      C.free ptrPtrStatusStr
      C.free capacityPtr
  return (Op add value finish)

opSendCloseFromClient :: IO (OpT ())
opSendCloseFromClient = do
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendCloseFromClient))
      ]
    finish =
      return ()
  return (Op add value finish)

clientWaitForInitialMetadata :: ClientReaderWriter -> IO (RpcReply [Metadata])
clientWaitForInitialMetadata crw@(ClientReaderWriter { .. }) = do
  initMD <- readIORef initialMDRef
  case initMD of
    Just md -> return (RpcOk md)
    Nothing -> do
      recvInitialMetadataOp <- opRecvInitialMetadata crw
      res <- callBatch crw [ OpX recvInitialMetadataOp ]
      case res of
        RpcOk _ -> do
          md <- opRead recvInitialMetadataOp
          writeIORef initialMDRef (Just md)
          return (RpcOk md)
        RpcError err ->
          return (RpcError err)

clientReadInitialMetadata :: ClientReaderWriter -> IO (RpcReply (Maybe [Metadata]))
clientReadInitialMetadata (ClientReaderWriter {..}) = do
  md <- readIORef initialMDRef
  return (RpcOk md)

clientRead :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
clientRead crw = do
  _ <- clientWaitForInitialMetadata crw
  recvMessage crw

recvMessage :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
recvMessage crw = do
  recvMessageOp <- opRecvMessage
  res <- callBatch crw [ OpX recvMessageOp ]
  case res of
    RpcOk _ -> do
      bs <- opRead recvMessageOp
      return (RpcOk bs)
    RpcError err -> do
      putStrLn "recvMessage: callBatch failed"
      return (RpcError err)

clientWaitForStatus :: ClientReaderWriter -> IO (RpcReply RpcStatus)
clientWaitForStatus crw@(ClientReaderWriter{..}) = do
  status <- readIORef statusFromServer
  case status of
    Nothing -> do
      recvStatusOp <- opRecvStatusOnClient crw
      res <- callBatch crw [ OpX recvStatusOp ]
      case res of
        RpcOk _ -> do
          st <- opRead recvStatusOp
          return (RpcOk st)
        RpcError err -> do
          return (RpcError err)
    Just st -> return (RpcOk st)

clientClose :: ClientReaderWriter -> IO ()
clientClose (ClientReaderWriter{..}) = do
  withMVar callMVar_ $ \call ->
    grpcCallDestroy call

clientWrite :: ClientReaderWriter -> B.ByteString -> IO (RpcReply ())
clientWrite crw@(ClientReaderWriter{..}) arg = do
  sendMessageOp <- opSendMessage arg
  callBatch crw [ OpX sendMessageOp ]

clientSendHalfClose :: ClientReaderWriter -> IO (RpcReply ())
clientSendHalfClose crw@(ClientReaderWriter{..}) = do
  sendCloseOp <- opSendCloseFromClient
  callBatch crw [ OpX sendCloseOp ]

clientCloseCall :: ClientReaderWriter -> IO ()
clientCloseCall ClientReaderWriter { callMVar_ = callMVar } =
  withMVar callMVar $ \call ->
  grpcCallDestroy call

type Rpc req resp a = ReaderT (Client req resp) (ExceptT RpcError IO) a

askCrw :: Rpc req resp ClientReaderWriter
askCrw = asks clientCrw

askDecoder :: Rpc req resp (Decoder resp)
askDecoder = asks clientDecoder

askEncoder :: Rpc req resp (Encoder req)
askEncoder = asks clientEncoder

joinReply :: RpcReply a -> Rpc req resp a
joinReply (RpcOk a) = return a
joinReply (RpcError err) = lift (throwE err)

withNewClient :: RpcReply (Client req resp) -> Rpc req resp a -> IO (RpcReply a)
withNewClient r_client m = do
  case r_client of
    RpcOk client -> withClient client m
    RpcError err -> return (RpcError err)

withClient :: Client req resp -> Rpc req resp a -> IO (RpcReply a)
withClient client m = do
  e <- runExceptT (runReaderT m client)
  case e of
    Left err -> return (RpcError err)
    Right a -> return (RpcOk a)

initialMetadata :: Rpc req resp [Metadata]
initialMetadata = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForInitialMetadata crw)

waitForStatus :: Rpc req resp RpcStatus
waitForStatus = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForStatus crw)

receiveMessage :: Rpc req resp (Maybe resp)
receiveMessage = do
  crw <- askCrw
  msg <- joinReply =<< liftIO (clientRead crw)
  case msg of
    Nothing -> return Nothing
    Just x -> do
      decoder <- askDecoder
      liftM Just (joinReply =<< liftIO (decoder x))

receiveAllMessages :: Rpc req resp [resp]
receiveAllMessages = do
  crw <- askCrw
  decoder <- askDecoder
  let go acc = do
        value <- joinReply =<< liftIO (clientRead crw)
        case value of
          Just x -> do
            y <- joinReply =<< (liftIO (decoder x))
            go (y:acc)
          Nothing -> return (reverse acc)
  go []

sendMessage :: req -> Rpc req resp ()
sendMessage o = do
  crw <- askCrw
  encoder <- askEncoder
  x <- joinReply =<< liftIO (encoder o)
  joinReply =<< liftIO (clientWrite crw x)

sendHalfClose :: Rpc req resp ()
sendHalfClose = do
  crw <- askCrw
  joinReply =<< liftIO (clientSendHalfClose crw)

closeCall :: Rpc req resp ()
closeCall = do
  crw <- askCrw
  liftIO (clientClose crw)

data RpcReply a
  = RpcOk a
  | RpcError RpcError
  deriving Show

data RpcError
  = DeadlineExceeded
  | Cancelled
  | Error String
  | CallErrorStatus CallError
  | StatusError StatusCode B.ByteString
  deriving Show

-- Copyright (c) 2016, Google Inc.
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--     * Redistributions of source code must retain the above copyright
--       notice, this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of Google Inc. nor the
--       names of its contributors may be used to endorse or promote products
--       derived from this software without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
-- ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-- WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL Google Inc. BE LIABLE FOR ANY
-- DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-- (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
-- LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
-- ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-- (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-- SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--
--------------------------------------------------------------------------------
{-# LANGUAGE ScopedTypeVariables, OverloadedStrings, ExistentialQuantification, RecordWildCards #-}
module Network.Grpc.Core.Call where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad

import qualified Data.ByteString              as B
import qualified Data.ByteString.Lazy         as L
import           Data.Monoid ((<>), Last(..))
import           Data.IORef

import           Foreign.C.Types              as C
import qualified Foreign.Marshal.Alloc        as C
import qualified Foreign.Marshal.Utils        as C
import qualified Foreign.Ptr                  as C
import           Foreign.Ptr (Ptr)
import qualified Foreign.Storable             as C

-- transformers
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.Except
import           Control.Monad.Trans.Reader

import qualified Network.Grpc.CompletionQueue as CQ
import Network.Grpc.Lib.PropagationBits
{#import Network.Grpc.Lib.ByteBuffer#}
{#import Network.Grpc.Lib.Metadata#}
{#import Network.Grpc.Lib.TimeSpec#}
{#import Network.Grpc.Lib.Core#}

#include <grpc/grpc.h>
#include "hs_grpc.h"

{#context lib = "grpc" prefix = "grpc" #}

data Deadline
  = AbsoluteDeadline TimeSpec
  | RelativeDeadline Int --seconds

data ClientContext = ClientContext
  { ccChannel :: Channel
  , ccCQ :: CompletionQueue
  , ccWorker :: CQ.Worker
  }

data CallOptions = CallOptions
  { coDeadline :: Maybe Deadline
  , coParentContext :: Maybe () -- todo
  , coPropagationMask :: Maybe PropagationMask
  , coMetadata :: [Metadata] -- todo
  }

instance Monoid CallOptions where
  mempty = CallOptions Nothing Nothing Nothing []
  mappend (CallOptions a b c d) (CallOptions a' b' c' d') =
    CallOptions
      (getLast (Last a <> Last a'))
      (getLast (Last b <> Last b'))
      (getLast (Last c <> Last c'))
      (d <> d')

withAbsoluteDeadline :: TimeSpec -> CallOptions
withAbsoluteDeadline deadline =
  mempty { coDeadline = Just (AbsoluteDeadline deadline) }

withRelativeDeadlineSeconds :: Int -> CallOptions
withRelativeDeadlineSeconds seconds =
  mempty { coDeadline = Just (RelativeDeadline seconds) }

withParentContext :: () -> CallOptions
withParentContext ctx =
  mempty { coParentContext = Just ctx }

withParentContextPropagating :: () -> PropagationMask -> CallOptions
withParentContextPropagating ctx prop =
  mempty { coParentContext   = Just ctx
         , coPropagationMask = Just prop }

withMetadata :: [Metadata] -> CallOptions
withMetadata md =
  mempty { coMetadata = md }

resolveDeadline :: CallOptions -> IO TimeSpec
resolveDeadline co =
  case coDeadline co of
    Nothing -> return gprInfFuture
    Just (AbsoluteDeadline deadline) -> return deadline
    Just (RelativeDeadline seconds) ->
      secondsFromNow (fromIntegral seconds)

newClientContext :: Channel -> IO ClientContext
newClientContext chan = do
  cq <- completionQueueCreate reservedPtr
  cqt <- CQ.startCompletionQueueThread cq
  return (ClientContext chan cq cqt)

destroyClientContext :: ClientContext -> IO ()
destroyClientContext (ClientContext _ cq w) = do
  completionQueueShutdown cq
  CQ.waitWorkerTermination w

type MethodName = B.ByteString
type Arg = B.ByteString

-- | Run the IO function once, cache it in the MVar.
onceMVar :: MVar (Maybe a) -> IO a -> IO a
onceMVar mvar io = modifyMVar mvar $ \x ->
  case x of
    Just x' -> return (x, x')
    Nothing -> do
      x' <- io
      return (Just x', x')

reservedPtr :: Ptr ()
reservedPtr = C.nullPtr

data UnaryResult a = UnaryResult [Metadata] [Metadata] a deriving Show

callUnary :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (UnaryResult L.ByteString))
callUnary ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  bracket (grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline) grpcCallDestroy $ \call0 -> newMVar call0 >>= \mcall -> do
    crw <- newClientReaderWriter ctx mcall

    sendInitOp <- opSendInitialMetadata (coMetadata co)
    sendCloseOp <- opSendCloseFromClient
    sendMessageOp <- opSendMessage arg
    recvStatusOp <- opRecvStatusOnClient crw
    recvMessageOp <- opRecvMessage
    recvInitialMetadataOp <- opRecvInitialMetadata crw

    res <- callBatch crw [
          OpX sendInitOp
        , OpX sendCloseOp
        , OpX recvInitialMetadataOp
        , OpX recvMessageOp
        , OpX sendMessageOp
        , OpX recvStatusOp
        ]
    case res of
      RpcOk _ -> do
        (RpcStatus trailMD status statusDetails) <- opRead recvStatusOp
        case status of
          StatusOk -> do
            initMD <- opRead recvInitialMetadataOp
            answ <- opRead recvMessageOp
            let answ' = maybe L.empty id answ
            return (RpcOk (UnaryResult initMD trailMD answ'))
          _ -> return (RpcError (StatusError status statusDetails))
      RpcError err -> return (RpcError err)

callDownstream :: ClientContext -> CallOptions -> MethodName -> Arg -> IO (RpcReply (Client B.ByteString L.ByteString))
callDownstream ctx@(ClientContext chan cq _) co method arg = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)
  sendCloseOp <- opSendCloseFromClient
  sendMessageOp <- opSendMessage arg

  res <- callBatch crw [
        OpX sendInitOp
      , OpX sendCloseOp
      , OpX sendMessageOp
      ]
  case res of
    RpcOk _ -> do
      return (RpcOk (Client crw defaultEncoder defaultDecoder))
    RpcError err -> return (RpcError err)

callUpstream :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callUpstream ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline >>= newMVar
  crw <- newClientReaderWriter ctx mcall

  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> do
      return (RpcOk (Client crw defaultEncoder defaultDecoder))
    RpcError err -> return (RpcError err)

callBidi :: ClientContext -> CallOptions -> MethodName -> IO (RpcReply (Client B.ByteString L.ByteString))
callBidi ctx@(ClientContext chan cq _) co method = do
  deadline <- resolveDeadline co
  mcall <- grpcChannelCreateCall chan C.nullPtr propagateDefaults cq method "localhost" deadline >>= newMVar

  crw <- newClientReaderWriter ctx mcall
  sendInitOp <- opSendInitialMetadata (coMetadata co)

  res <- callBatch crw [ OpX sendInitOp ]
  case res of
    RpcOk _ -> return (RpcOk (Client crw defaultEncoder defaultDecoder))
    RpcError err -> return (RpcError err)

data Client req resp = Client
  { clientCrw :: ClientReaderWriter
  , clientEncoder :: Encoder req
  , clientDecoder :: Decoder resp
  }

type Decoder a = L.ByteString -> IO (RpcReply a)
type Encoder a = a -> IO (RpcReply B.ByteString)

defaultDecoder :: L.ByteString -> IO (RpcReply L.ByteString)
defaultDecoder bs = return (RpcOk bs)

defaultEncoder :: B.ByteString -> IO (RpcReply B.ByteString)
defaultEncoder bs = return (RpcOk bs)

data ClientReaderWriter = ClientReaderWriter {
  context :: ClientContext,
  callMVar_ :: MVar Call,
  initialMDRef :: !(IORef (Maybe [Metadata])),
  trailingMDRef :: !(IORef (Maybe [Metadata])),
  statusFromServer :: !(IORef (Maybe RpcStatus))
}

newClientReaderWriter :: ClientContext -> MVar Call -> IO ClientReaderWriter
newClientReaderWriter ctx mcall = do
  initMD <- newIORef Nothing
  trailMD <- newIORef Nothing
  status <- newIORef Nothing
  return (ClientReaderWriter ctx mcall initMD trailMD status)

data OpX = forall t. OpX (OpT t)

data OpArray = OpArray { opArrPtr :: !GrpcOpPtr
                       , opArrLen :: !CULong
                       , opArrFree :: !(IO ())
                       }

toArray :: [OpX] -> IO OpArray
toArray ops = do
  aptr <- C.mallocBytes (length ops * {#sizeof grpc_op#})
  let ptrs = iterate (`C.plusPtr` {#sizeof grpc_op#}) aptr
      ops' = concatMap (\(OpX op) -> opAdd op) ops
      free = C.free aptr >> sequence_ (map (\(OpX op) -> opFinish op) ops)
      write op p = do
        {#set grpc_op->flags#} p 0
        {#set grpc_op->reserved#} p C.nullPtr
        op p
  zipWithM_ write ops' ptrs
  return $! OpArray aptr (fromIntegral $ length ops) free

callBatch :: ClientReaderWriter -> [OpX] -> IO (RpcReply ())
callBatch crw ops = do
  let ctx = context crw
  CQ.withEvent (ccWorker ctx) $ \eDesc -> do
    arr <- toArray ops
    callStatus <- withMVar (callMVar_ crw) $ \call ->
      grpcCallStartBatch call (opArrPtr arr) (opArrLen arr) (CQ.eventTag eDesc) reservedPtr
    -- print ("callStatus: " ++ show callStatus)
    case callStatus of
      CallOk -> do
        e <- CQ.interruptibleWaitEvent eDesc
        -- print ("event: " ++ show e)
        opArrFree arr -- TODO: We leak if we're interrupted.
        case e of
          QueueOpComplete OpSuccess _ -> return (RpcOk ())
          QueueOpComplete OpError _ -> return (RpcError (Error "callBatch: op error"))
          QueueTimeOut -> return (RpcError DeadlineExceeded)
          QueueShutdown -> return (RpcError (Error "queue shutdown"))
      _ -> do
         -- print callStatus
         return (RpcError (CallErrorStatus callStatus))

data OpT out = Op
  { opAdd    :: [Ptr GrpcOp -> IO ()]
  , opValue  :: IORef out
  , opFinish :: IO ()
  }

opRead :: OpT a -> IO a
opRead op = readIORef (opValue op)

opRecvInitialMetadata :: ClientReaderWriter -> IO (OpT [Metadata])
opRecvInitialMetadata crw = do
  arr <- mallocMetadataArray
  value <- newIORef (error "opRecvInitialMetadata never finished")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvInitialMetadata))
        {#set grpc_op->data.recv_initial_metadata#} p arr
      ]
    finish = do
      mds <- readMetadataArray arr
      writeIORef (initialMDRef crw) (Just mds)
      writeIORef value mds
      freeMetadataArray arr
  return (Op add value finish)

opSendMessage :: B.ByteString -> IO (OpT ())
opSendMessage bs = do
  bb <- fromByteString bs
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendMessage))
        {#set grpc_op->data.send_message#} p bb
      ]
    finish =
      byteBufferDestroy bb
  return (Op add value finish)

opRecvMessage :: IO (OpT (Maybe L.ByteString))
opRecvMessage = do
  bbptr <- C.malloc :: IO (Ptr (Ptr CByteBuffer))
  value <- newIORef Nothing
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvMessage))
        {#set grpc_op->data.recv_message#} p bbptr
      ]
    finish = do
      bb <- C.peek bbptr
      C.free bbptr
      writeIORef value =<< if bb /= C.nullPtr
        then do
          lbs <- toLazyByteString bb
          byteBufferDestroy bb
          return (Just lbs)
        else return Nothing
  return (Op add value finish)

opSendInitialMetadata :: [Metadata] -> IO (OpT ())
opSendInitialMetadata elems = do
  (mdArrPtr, free) <- mallocMetadata elems
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendInitialMetadata))
        {#set grpc_op->data.send_initial_metadata.count#} p (fromIntegral (length elems))
        {#set grpc_op->data.send_initial_metadata.metadata#} p (C.castPtr mdArrPtr)
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.is_set#} p 0
        {#set grpc_op->data.send_initial_metadata.maybe_compression_level.level#} p 0
      ]
    finish = do
      free
  return (Op add value finish)

data RpcStatus = RpcStatus [Metadata] StatusCode B.ByteString
  deriving Show

opRecvStatusOnClient :: ClientReaderWriter -> IO (OpT RpcStatus)
opRecvStatusOnClient (ClientReaderWriter{..}) = do
  trailingMetadataArrPtr <- mallocMetadataArray
  statusCodePtr <- C.malloc :: IO (Ptr StatusCodeT)
  ptrPtrStatusStr <- C.new C.nullPtr :: IO (Ptr (Ptr CChar))
  capacityPtr <- C.new 0 :: IO (Ptr SizeT)
  value <- newIORef (error "opRecvStatusOnClient never ran")
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpRecvStatusOnClient))
        {#set grpc_op->data.recv_status_on_client.trailing_metadata#} p trailingMetadataArrPtr
        {#set grpc_op->data.recv_status_on_client.status#} p statusCodePtr
        {#set grpc_op->data.recv_status_on_client.status_details#} p ptrPtrStatusStr
        {#set grpc_op->data.recv_status_on_client.status_details_capacity#} p capacityPtr
      ]
    finish = do
      trailingMd <- readMetadataArray trailingMetadataArrPtr
      statusCode <- fmap toStatusCode (C.peek statusCodePtr)
      statusDetails <- B.packCString =<< C.peek ptrPtrStatusStr
      let status = RpcStatus trailingMd statusCode statusDetails
      writeIORef value status
      writeIORef statusFromServer (Just status)
      free
    free = do
      freeMetadataArray trailingMetadataArrPtr
      C.free statusCodePtr
      C.free ptrPtrStatusStr
      C.free capacityPtr
  return (Op add value finish)

opSendCloseFromClient :: IO (OpT ())
opSendCloseFromClient = do
  value <- newIORef ()
  let
    add =
      [ \p -> do
        {#set grpc_op->op#} p (fromIntegral (fromEnum OpSendCloseFromClient))
      ]
    finish =
      return ()
  return (Op add value finish)

clientWaitForInitialMetadata :: ClientReaderWriter -> IO (RpcReply [Metadata])
clientWaitForInitialMetadata crw@(ClientReaderWriter { .. }) = do
  initMD <- readIORef initialMDRef
  case initMD of
    Just md -> return (RpcOk md)
    Nothing -> do
      recvInitialMetadataOp <- opRecvInitialMetadata crw
      res <- callBatch crw [ OpX recvInitialMetadataOp ]
      case res of
        RpcOk _ -> do
          md <- opRead recvInitialMetadataOp
          writeIORef initialMDRef (Just md)
          return (RpcOk md)
        RpcError err ->
          return (RpcError err)

clientReadInitialMetadata :: ClientReaderWriter -> IO (RpcReply (Maybe [Metadata]))
clientReadInitialMetadata (ClientReaderWriter {..}) = do
  md <- readIORef initialMDRef
  return (RpcOk md)

clientRead :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
clientRead crw = do
  _ <- clientWaitForInitialMetadata crw
  recvMessage crw

recvMessage :: ClientReaderWriter -> IO (RpcReply (Maybe L.ByteString))
recvMessage crw = do
  recvMessageOp <- opRecvMessage
  res <- callBatch crw [ OpX recvMessageOp ]
  case res of
    RpcOk _ -> do
      bs <- opRead recvMessageOp
      return (RpcOk bs)
    RpcError err -> do
      putStrLn "recvMessage: callBatch failed"
      return (RpcError err)

clientWaitForStatus :: ClientReaderWriter -> IO (RpcReply RpcStatus)
clientWaitForStatus crw@(ClientReaderWriter{..}) = do
  status <- readIORef statusFromServer
  case status of
    Nothing -> do
      recvStatusOp <- opRecvStatusOnClient crw
      res <- callBatch crw [ OpX recvStatusOp ]
      case res of
        RpcOk _ -> do
          st <- opRead recvStatusOp
          return (RpcOk st)
        RpcError err -> do
          return (RpcError err)
    Just st -> return (RpcOk st)

clientClose :: ClientReaderWriter -> IO ()
clientClose (ClientReaderWriter{..}) = do
  withMVar callMVar_ $ \call ->
    grpcCallDestroy call

clientWrite :: ClientReaderWriter -> B.ByteString -> IO (RpcReply ())
clientWrite crw@(ClientReaderWriter{..}) arg = do
  sendMessageOp <- opSendMessage arg
  callBatch crw [ OpX sendMessageOp ]

clientSendHalfClose :: ClientReaderWriter -> IO (RpcReply ())
clientSendHalfClose crw@(ClientReaderWriter{..}) = do
  sendCloseOp <- opSendCloseFromClient
  callBatch crw [ OpX sendCloseOp ]

clientCloseCall :: ClientReaderWriter -> IO ()
clientCloseCall ClientReaderWriter { callMVar_ = callMVar } =
  withMVar callMVar $ \call ->
  grpcCallDestroy call

type Rpc req resp a = ReaderT (Client req resp) (ExceptT RpcError IO) a

askCrw :: Rpc req resp ClientReaderWriter
askCrw = asks clientCrw

askDecoder :: Rpc req resp (Decoder resp)
askDecoder = asks clientDecoder

askEncoder :: Rpc req resp (Encoder req)
askEncoder = asks clientEncoder

joinReply :: RpcReply a -> Rpc req resp a
joinReply (RpcOk a) = return a
joinReply (RpcError err) = lift (throwE err)

withNewClient :: RpcReply (Client req resp) -> Rpc req resp a -> IO (RpcReply a)
withNewClient r_client m = do
  case r_client of
    RpcOk client -> withClient client m
    RpcError err -> return (RpcError err)

withClient :: Client req resp -> Rpc req resp a -> IO (RpcReply a)
withClient client m = do
  e <- runExceptT (runReaderT m client)
  case e of
    Left err -> return (RpcError err)
    Right a -> return (RpcOk a)

initialMetadata :: Rpc req resp [Metadata]
initialMetadata = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForInitialMetadata crw)

waitForStatus :: Rpc req resp RpcStatus
waitForStatus = do
  crw <- askCrw
  joinReply =<< liftIO (clientWaitForStatus crw)

receiveMessage :: Rpc req resp (Maybe resp)
receiveMessage = do
  crw <- askCrw
  msg <- joinReply =<< liftIO (clientRead crw)
  case msg of
    Nothing -> return Nothing
    Just x -> do
      decoder <- askDecoder
      liftM Just (joinReply =<< liftIO (decoder x))

receiveAllMessages :: Rpc req resp [resp]
receiveAllMessages = do
  crw <- askCrw
  decoder <- askDecoder
  let go acc = do
        value <- joinReply =<< liftIO (clientRead crw)
        case value of
          Just x -> do
            y <- joinReply =<< (liftIO (decoder x))
            go (y:acc)
          Nothing -> return (reverse acc)
  go []

sendMessage :: req -> Rpc req resp ()
sendMessage o = do
  crw <- askCrw
  encoder <- askEncoder
  x <- joinReply =<< liftIO (encoder o)
  joinReply =<< liftIO (clientWrite crw x)

sendHalfClose :: Rpc req resp ()
sendHalfClose = do
  crw <- askCrw
  joinReply =<< liftIO (clientSendHalfClose crw)

closeCall :: Rpc req resp ()
closeCall = do
  crw <- askCrw
  liftIO (clientClose crw)

data RpcReply a
  = RpcOk a
  | RpcError RpcError
  deriving Show

data RpcError
  = DeadlineExceeded
  | Cancelled
  | Error String
  | CallErrorStatus CallError
  | StatusError StatusCode B.ByteString
  deriving Show