{-# LANGUAGE DeriveDataTypeable, ForeignFunctionInterface #-}

{- Loosely based on: v4l-capture.c

/*
 *  V4L2 video capture example
 *
 *  This program can be used and distributed without restrictions.
 *
 *      This program is provided with the V4L2 API
 * see http://linuxtv.org/docs.php for more information
 */

  This version writes YUYV frames on stdout in YUV4MPEG2 format.
  Example usage:
  $ ./v4l2-capture |
    y4mscaler -I ilace=none -I sar=1/1 -O chromass=420mpeg2 |
    mplayer -demuxer y4m - -vo x11

  TODO: this code is hella stinky, needs major work...

-}

import Prelude hiding (catch)
import System.Exit (exitFailure)
import Bindings.Linux.VideoDev2
import Bindings.LibV4L2
import Bindings.MMap
import System.C.IO
import System.Posix.IOCtl (IOControl(..))
import System.Posix.Types (Fd(Fd))
import Bindings.Posix.Sys.Select (c'FD_ZERO, c'FD_SET, c'select)
import qualified Bindings.Posix.Sys.Select as P
import Foreign
import Foreign.C
import Control.Exception
import GHC.Conc (threadDelay)
import GHC.IO.Exception (IOErrorType(ResourceExhausted), ioe_type)
import System.Environment (getArgs)
import System.IO (hPutBuf, stderr, stdout, hFlush, hPutStr)
import Control.Monad (forM, forM_, replicateM, when)
import Data.Typeable (Typeable)

writeFrames :: Bool
writeFrames = True

data Mode = Read | MMap | UserPtr

data V4L2 = V4L2
  { getFd :: IO Fd
  , getWidth :: IO Word32
  , getHeight :: IO Word32
  , nextFrame :: (Ptr Word8 -> Int -> IO ()) -> IO ()
  , finish :: IO ()
  }

data V4L2Exception = V4L2NoCapture | V4L2NoRead | V4L2NoStream | V4L2NoBuffers | V4L2NoMMap
  deriving (Show, Typeable)
instance Exception V4L2Exception

v4l2 :: Mode -> FilePath -> IO V4L2
v4l2 mode device = do
  fdi <- withCString device $ \dev -> c'v4l2_open dev 2050 0 -- (c'O_RDWR .|. c'O_NONBLOCK)  -- FIXME
  let fdfi = fromIntegral fdi
      fd = Fd fdfi
  cap <- ioctl fd C'VIDIOC_QUERYCAP =<< empty
  when (c'v4l2_capability'capabilities cap .&. c'V4L2_CAP_VIDEO_CAPTURE == 0) $ throwIO V4L2NoCapture
  case mode of
    Read -> when (c'v4l2_capability'capabilities cap .&. c'V4L2_CAP_READWRITE == 0) $ throwIO V4L2NoRead
    _ -> when (c'v4l2_capability'capabilities cap .&. c'V4L2_CAP_STREAMING == 0) $ throwIO V4L2NoStream
  (do
    cropcap <- ioctl fd C'VIDIOC_CROPCAP . (\e -> e
      { c'v4l2_cropcap'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE }) =<< empty
    ioctl_ fd C'VIDIOC_S_CROP . (\e -> e
      { c'v4l2_crop'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
      , c'v4l2_crop'c = c'v4l2_cropcap'defrect cropcap }) =<< empty
   ) `catchIO` \_ -> return ()
  pix0 <- return . (\e -> e
    { c'v4l2_pix_format'width = 320
    , c'v4l2_pix_format'height = 240
    , c'v4l2_pix_format'pixelformat = c'V4L2_PIX_FMT_YUYV
    , c'v4l2_pix_format'field = c'V4L2_FIELD_INTERLACED }) =<< empty
  pix1 <- (`u'v4l2_format_u'pix` pix0) =<< empty
  ioctl_ fd C'VIDIOC_S_FMT . (\e -> e
    { c'v4l2_format'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
    , c'v4l2_format'fmt = pix1 }) =<< empty
  fmt <- ioctl fd C'VIDIOC_G_FMT . (\e -> e
    { c'v4l2_format'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE }) =<< empty
  let pix = c'v4l2_format_u'pix . c'v4l2_format'fmt $ fmt
      w = c'v4l2_pix_format'width pix
      h = c'v4l2_pix_format'height pix
      bytes = fromIntegral $ c'v4l2_pix_format'sizeimage pix
  case mode of
    Read -> do
      buffer <- mallocBytes bytes
      return V4L2
        { getFd = return fd
        , getWidth = return w
        , getHeight = return h
        , nextFrame = \f -> do
            throwErrnoIfMinus1_ "read" $ cread (fromIntegral fd) buffer bytes
            f buffer bytes
        , finish = do
            free buffer
            throwErrnoIfMinus1_ "close" $ c'v4l2_close fdi
        }
    MMap -> do
      req <- ioctl fd C'VIDIOC_REQBUFS . (\e -> e
        { c'v4l2_requestbuffers'count = 4
        , c'v4l2_requestbuffers'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
        , c'v4l2_requestbuffers'memory = c'V4L2_MEMORY_MMAP }) =<< empty
      let nbuffers = c'v4l2_requestbuffers'count req
      when (nbuffers < 2) $ throwIO $ V4L2NoBuffers
      buffers <- forM [0 .. nbuffers-1] $ \n -> do
        buf <- ioctl fd C'VIDIOC_QUERYBUF . (\e -> e
          { c'v4l2_buffer'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
          , c'v4l2_buffer'memory = c'V4L2_MEMORY_MMAP
          , c'v4l2_buffer'index = n }) =<< empty
        let len = fromIntegral $ c'v4l2_buffer'length buf
        ptr <- c'mmap nullPtr len (c'PROT_READ .|. c'PROT_WRITE) c'MAP_SHARED fdfi (fromIntegral . c'v4l2_buffer_u'offset . c'v4l2_buffer'u $ buf)
        when (ptr == p'MAP_FAILED) $ throwIO V4L2NoMMap
        return (ptr, len)
      forM_ (buffers `zip` [0 ..]) $ \(_, i) -> do
        ioctl_ fd C'VIDIOC_QBUF . (\e -> e
          { c'v4l2_buffer'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
          , c'v4l2_buffer'memory = c'V4L2_MEMORY_MMAP
          , c'v4l2_buffer'index = i }) =<< empty
        ioctl_ fd C'VIDIOC_STREAMON c'V4L2_BUF_TYPE_VIDEO_CAPTURE
      return V4L2
        { getFd = return fd
        , getWidth = return w
        , getHeight = return h
        , nextFrame = \f -> do
            buf <- ioctl fd C'VIDIOC_DQBUF . (\e -> e
              { c'v4l2_buffer'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
              , c'v4l2_buffer'memory = c'V4L2_MEMORY_MMAP }) =<< empty
            let (ptr, len) = (fst $ buffers !! fromIntegral (c'v4l2_buffer'index buf), c'v4l2_buffer'bytesused buf)
            f (castPtr ptr) (fromIntegral len)
            ioctl_ fd C'VIDIOC_QBUF buf
        , finish = do
            ioctl_ fd C'VIDIOC_STREAMOFF c'V4L2_BUF_TYPE_VIDEO_CAPTURE
            forM_ buffers $ \(ptr, len) -> do
              throwErrnoIfMinus1_ "munmap" $ c'munmap ptr len
            throwErrnoIfMinus1_ "close" $ c'v4l2_close fdi
        }
    UserPtr -> do
      ioctl_ fd C'VIDIOC_REQBUFS . (\e -> e
        { c'v4l2_requestbuffers'count = 4
        , c'v4l2_requestbuffers'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
        , c'v4l2_requestbuffers'memory = c'V4L2_MEMORY_USERPTR }) =<< empty
      buffers <- replicateM 4 (mallocBytes bytes)
      forM_ buffers $ \ptr -> do
        u <- (\e -> u'v4l2_buffer_u'userptr e ({-unsafe-} fromIntegral (ptrToWordPtr ptr))) =<< empty
        ioctl_ fd C'VIDIOC_QBUF . (\e -> e
          { c'v4l2_buffer'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
          , c'v4l2_buffer'memory = c'V4L2_MEMORY_USERPTR
          , c'v4l2_buffer'u = u
          , c'v4l2_buffer'length = fromIntegral bytes }) =<< empty
        ioctl_ fd C'VIDIOC_STREAMON c'V4L2_BUF_TYPE_VIDEO_CAPTURE
      return V4L2
        { getFd = return fd
        , getWidth = return w
        , getHeight = return h
        , nextFrame = \f -> do
            buf <- ioctl fd C'VIDIOC_DQBUF . (\e -> e
              { c'v4l2_buffer'type = c'V4L2_BUF_TYPE_VIDEO_CAPTURE
              , c'v4l2_buffer'memory = c'V4L2_MEMORY_USERPTR }) =<< empty
            let p = wordPtrToPtr . {-unsafe-} fromIntegral . c'v4l2_buffer_u'userptr . c'v4l2_buffer'u $ buf
            f p (fromIntegral $ c'v4l2_buffer'bytesused buf)
            ioctl_ fd C'VIDIOC_QBUF buf
        , finish = do
            -- FIXME
            mapM_ free buffers
            throwErrnoIfMinus1_ "close" $ c'v4l2_close fdi
       }

empty :: Storable a => IO a
empty = do
  alloca $ \p -> do
    _ <- c'memset p 0 (fromIntegral $ sizeOf (unsafePerformIO (peek p)))
    peek p

p'MAP_FAILED :: Ptr a
p'MAP_FAILED = intPtrToPtr c'MAP_FAILED

process_image :: V4L2 -> Ptr Word8 -> Int -> IO ()
process_image v4 p _size = do
  when writeFrames $ do
    w <- getWidth v4
    h <- getHeight v4
    putStrLn "FRAME"
    allocaBytes (fromIntegral $ w * h) $ \y -> do
      forM_ ([ 0, 2 .. fromIntegral $ w * h * 2 -1]`zip`[0..]) $ \(i, j) -> do
        pokeElemOff y j =<< peekElemOff p i
      hPutBuf stdout y (fromIntegral $ w * h)
    allocaBytes (fromIntegral $ (w * h)) $ \u -> do
      forM_ ([ 1, 5 .. fromIntegral $ w * h * 2 -1]`zip`[0..]) $ \(i, j) -> do
        pokeElemOff u j =<< peekElemOff p i
      hPutBuf stdout u (fromIntegral $ (w`div`2) * h)
    allocaBytes (fromIntegral $ (w * h)) $ \v -> do
      forM_ ([ 3, 7 .. fromIntegral $ w * h * 2 -1 ]`zip`[0..]) $ \(i, j) -> do
        pokeElemOff v j =<< peekElemOff p i
      hPutBuf stdout v (fromIntegral $ (w`div`2) * h)
  hFlush stderr
  hPutStr stderr "."
  hFlush stdout

main :: IO ()
main = do
  (device:_) <- (++["/dev/video0"]) `fmap` getArgs
  v <- v4l2 MMap device
  when writeFrames $ do
    w <- getWidth v
    h <- getHeight v
    putStrLn $ "YUV4MPEG2 W" ++ show w ++ " H" ++ show h ++ " C422"
  mainLoopRetry v (process_image v)
  hPutStr stderr "\n"

mainLoopRetry :: V4L2 -> (Ptr Word8 -> Int -> IO ()) -> IO ()
mainLoopRetry v f = do
  fd <- fromIntegral `fmap` getFd v
  alloca $ \fds -> do
    c'FD_ZERO fds
    c'FD_SET fd fds
    with P.C'timeval{ P.c'timeval'tv_sec = 5, P.c'timeval'tv_usec = 0 } $ \tv -> do
      r <- c'select (fd + 1) fds nullPtr nullPtr tv
      e <- getErrno
      if r == -1 && e == eINTR
        then mainLoopRetry v f
        else if r > 0
               then do
                      nextFrame v f `catch` \ex -> case ioe_type ex of
                        ResourceExhausted -> threadDelay 100000 >> mainLoopRetry v f
                        _ -> exitFailure
                      mainLoopRetry v f
               else return ()

foreign import ccall "string.h memset" c'memset :: Ptr a -> CInt -> CSize -> IO (Ptr a)

catchIO :: IO a -> (IOException -> IO a) -> IO a
catchIO = catch

--foreign import ccall "fcntl.h open" 

c_ioctl' :: IOControl req d => Fd -> req -> Ptr d -> IO ()
c_ioctl' f req p = throwErrnoIfMinus1_ "ioctl" $ c'v4l2_ioctl (fromIntegral f) (fromIntegral $ ioctlReq req) (castPtr p)

-- | Calls a ioctl reading the structure after the call
ioctl :: IOControl req d
      => Fd -- ^ The file descriptor 
      -> req -- ^ The request
      -> d -- ^ The data
      -> IO d -- ^ The data after the call
ioctl f req d = with d $ \p -> c_ioctl' f req p >> peek p

-- | Call a ioctl ignoring the result
ioctl_ :: IOControl req d
       => Fd -- ^ The file descriptor
       -> req -- ^ The request
       -> d -- ^ The data
       -> IO ()
ioctl_ f req d = with d $ \p -> c_ioctl' f req p

{-
-- | Call a ioctl with uninitialized data
ioctl' :: IOControl req d
       => Fd -- ^ The file descriptor
       -> req -- ^ The request
       -> IO d -- ^ The data
ioctl' f req = alloca $ \p -> c_ioctl' f req p >> peek p
-}