content: {-# LANGUAGE FlexibleContexts         #-}
  {-# LANGUAGE ForeignFunctionInterface #-}
  {-# LANGUAGE RecordWildCards          #-}
  -- | Video decoding API. Includes FFI declarations for the underlying
  -- FFmpeg functions, wrappers for these functions that wrap error
  -- condition checking, and high level Haskellized interfaces.
  module Codec.FFmpeg.Decode where
  
  import           Codec.FFmpeg.AudioStream
  import           Codec.FFmpeg.Common
  import           Codec.FFmpeg.Enums
  import           Codec.FFmpeg.Scaler
  import           Codec.FFmpeg.Types
  import           Control.Arrow             (first)
  import           Control.Monad             (void, when)
  import           Control.Monad.Except
  import           Control.Monad.IO.Class    (MonadIO(liftIO))
  import           Control.Monad.Trans.Maybe
  import           Foreign.C.String
  import           Foreign.C.Types
  import           Foreign.Marshal.Alloc     (alloca, free, mallocBytes)
  import           Foreign.Marshal.Array     (advancePtr)
  import           Foreign.Marshal.Utils     (with)
  import           Foreign.Ptr
  import           Foreign.Storable
  
  -- * FFI Declarations
  
  foreign import ccall "avformat_open_input"
    avformat_open_input :: Ptr AVFormatContext -> CString -> Ptr AVInputFormat
                        -> Ptr AVDictionary -> IO CInt
  
  foreign import ccall "avformat_find_stream_info"
    avformat_find_stream_info :: AVFormatContext -> Ptr () -> IO CInt
  
  foreign import ccall "av_find_best_stream"
    av_find_best_stream :: AVFormatContext -> AVMediaType -> CInt -> CInt
                        -> Ptr AVCodec -> CInt -> IO CInt
  
  foreign import ccall "avcodec_find_decoder"
     avcodec_find_decoder :: AVCodecID -> IO AVCodec
  
  foreign import ccall "avcodec_find_decoder_by_name"
    avcodec_find_decoder_by_name :: CString -> IO AVCodec
  
  foreign import ccall "avpicture_get_size"
    avpicture_get_size :: AVPixelFormat -> CInt -> CInt -> IO CInt
  
  foreign import ccall "av_malloc"
    av_malloc :: CSize -> IO (Ptr ())
  
  foreign import ccall "av_read_frame"
    av_read_frame :: AVFormatContext -> AVPacket -> IO CInt
  
  foreign import ccall "avcodec_decode_audio4"
    decode_audio :: AVCodecContext -> AVFrame -> Ptr CInt -> AVPacket
                 -> IO CInt
  
  foreign import ccall "avcodec_decode_video2"
    decode_video :: AVCodecContext -> AVFrame -> Ptr CInt -> AVPacket
                 -> IO CInt
  foreign import ccall "avformat_close_input"
    close_input :: Ptr AVFormatContext -> IO ()
  
  foreign import ccall "av_dict_set"
    av_dict_set :: Ptr AVDictionary -> CString -> CString -> CInt -> IO CInt
  
  foreign import ccall "av_find_input_format"
    av_find_input_format :: CString -> IO (Ptr AVInputFormat)
  
  foreign import ccall "av_format_set_video_codec"
    av_format_set_video_codec :: AVFormatContext -> AVCodec -> IO ()
  
  dictSet :: Ptr AVDictionary -> String -> String -> IO ()
  dictSet d k v = do
    r <- withCString k $ \k' -> withCString v $ \v' ->
           av_dict_set d k' v' 0
    when (r < 0) $
      stringError r >>= \err ->
         error $ "av_dict_set failed("++ err ++"): "++k++" => "++v
  
  -- * FFmpeg Decoding Interface
  
  -- | Open the first video input device enumerated by FFMPEG.
  openCamera :: (MonadIO m, MonadError String m) => String -> CameraConfig -> m AVFormatContext
  openCamera cam cfg =
    wrapIOError . alloca $ \ctx ->
      withCString cam $ \cstr ->
        do avPtr <- mallocAVFormatContext
           setupCamera avPtr cam
           poke ctx avPtr
           fmt <- case format cfg of
                    Just "mjpeg" -> withCString "v4l2" av_find_input_format
                    Just f -> withCString f av_find_input_format
                    Nothing -> pure nullPtr
           r <- alloca $ \dict -> do
                  setConfig dict cfg
                  avformat_open_input ctx cstr fmt dict
           when (r /= 0) $
             stringError r >>= \err ->
               fail ("ffmpeg failed opening file: " ++ err)
           peek ctx
    where
      run :: (a -> IO b) -> Maybe a -> IO ()
      run _ Nothing  = return ()
      run f (Just x) = void (f x)
  
      setConfig :: Ptr AVDictionary -> CameraConfig -> IO ()
      setConfig dict (CameraConfig {..}) =
        do run (dictSet dict "framerate" . show) framerate
           run (\(w,h) -> dictSet dict "video_size" (show w ++ "x" ++ show h)) resolution
  
      setupCamera :: AVFormatContext -> String -> IO ()
      setupCamera avfc c =
        do setCamera avfc
           setFilename avfc c
           when (format cfg == Just "mjpeg") $ do
             mjpeg <- avcodec_find_decoder avCodecIdMjpeg
             setVideoCodecID avfc avCodecIdMjpeg
             av_format_set_video_codec avfc mjpeg
  
  openInput :: (MonadIO m, MonadError String m) => InputSource -> m AVFormatContext
  openInput ipt =
    case ipt of
      File fileName -> openFile fileName
      Camera cam cf -> openCamera cam cf
  
  -- | Open an input media file.
  openFile :: (MonadIO m, MonadError String m) => String -> m AVFormatContext
  openFile filename =
    wrapIOError . alloca $ \ctx ->
      withCString filename $ \cstr ->
        do poke (castPtr ctx) nullPtr
           r <- avformat_open_input ctx cstr nullPtr nullPtr
           when (r /= 0) (stringError r >>= \s ->
                            fail $ "ffmpeg failed opening file: " ++ s)
           peek ctx
  
  -- | @AVFrame@ is a superset of @AVPicture@, so we can upcast an
  -- 'AVFrame' to an 'AVPicture'.
  frameAsPicture :: AVFrame -> AVPicture
  frameAsPicture = AVPicture . getPtr
  
  -- | Find a codec given by name.
  findDecoder :: (MonadIO m, MonadError String m) => String -> m AVCodec
  findDecoder name =
    do r <- liftIO $ withCString name avcodec_find_decoder_by_name
       when (getPtr r == nullPtr)
            (throwError $ "Unsupported codec: " ++ show name)
       return r
  
  -- | Read packets of a media file to get stream information. This is
  -- useful for file formats with no headers such as MPEG.
  checkStreams :: (MonadIO m, MonadError String m) => AVFormatContext -> m ()
  checkStreams ctx =
    do r <- liftIO $ avformat_find_stream_info ctx nullPtr
       when (r < 0) (throwError "Couldn't find stream information")
  
  -- | Searches for a video stream in an 'AVFormatContext'. If one is
  -- found, returns the index of the stream in the container, and its
  -- associated 'AVCodecContext' and 'AVCodec'.
  findVideoStream :: (MonadIO m, MonadError String m)
                  => AVFormatContext
                  -> m (CInt, AVCodecContext, AVCodec, AVStream)
  findVideoStream fmt = do
    wrapIOError . alloca $ \codec -> do
        poke codec (AVCodec nullPtr)
        i <- av_find_best_stream fmt avmediaTypeVideo (-1) (-1) codec 0
        when (i < 0) (fail "Couldn't find a video stream")
        cod <- peek codec
        streams <- getStreams fmt
        vidStream <- peek (advancePtr streams (fromIntegral i))
        ctx <- getCodecContext vidStream
        return (i, ctx, cod, vidStream)
  
  findAudioStream :: (MonadIO m, MonadError String m)
                  => AVFormatContext
                  -> m (CInt, AVCodecContext, AVCodec, AVStream)
  findAudioStream fmt = do
    wrapIOError . alloca $ \codec -> do
        poke codec (AVCodec nullPtr)
        i <- av_find_best_stream fmt avmediaTypeAudio (-1) (-1) codec 0
        when (i < 0) (fail "Couldn't find audio stream")
        cod <- peek codec
        streams <- getStreams fmt
        audioStream <- peek (advancePtr streams (fromIntegral i))
        ctx <- getCodecContext audioStream
        return (i, ctx, cod, audioStream)
  
  -- | Find a registered decoder with a codec ID matching that found in
  -- the given 'AVCodecContext'.
  getDecoder :: (MonadIO m, MonadError String m)
             => AVCodecContext -> m AVCodec
  getDecoder ctx = do p <- liftIO $ getCodecID ctx >>= avcodec_find_decoder
                      when (getPtr p == nullPtr) (throwError "Unsupported codec")
                      return p
  
  -- | Initialize the given 'AVCodecContext' to use the given
  -- 'AVCodec'. **NOTE**: This function is not thread safe!
  openCodec :: (MonadIO m, MonadError String m)
            => AVCodecContext -> AVCodec -> m AVDictionary
  openCodec ctx cod =
    wrapIOError . alloca $ \dict -> do
      poke dict (AVDictionary nullPtr)
      r <- open_codec ctx cod dict
      when (r < 0) (fail "Couldn't open decoder")
      peek dict
  
  -- | Return the next frame of a stream.
  read_frame_check :: AVFormatContext -> AVPacket -> IO ()
  read_frame_check ctx pkt = do r <- av_read_frame ctx pkt
                                when (r < 0) (fail "Frame read failed")
  
  -- | Read frames of the given 'AVPixelFormat' from a video stream.
  frameReader :: (MonadIO m, MonadError String m)
              => AVPixelFormat -> InputSource -> m (IO (Maybe AVFrame), IO ())
  frameReader dstFmt ipt =
    do inputContext <- openInput ipt
       checkStreams inputContext
       (vidStreamIndex, ctx, cod, _vidStream) <- findVideoStream inputContext
       _ <- openCodec ctx cod
       prepareReader inputContext vidStreamIndex dstFmt ctx
  
  -- | Read RGB frames with the result in the 'MaybeT' transformer.
  --
  -- > frameReaderT = fmap (first MaybeT) . frameReader
  frameReaderT :: (Functor m, MonadIO m, MonadError String m)
               => InputSource -> m (MaybeT IO AVFrame, IO ())
  frameReaderT = fmap (first MaybeT) . frameReader avPixFmtRgb24
  
  -- | Read time stamped frames of the given 'AVPixelFormat' from a
  -- video stream. Time is given in seconds from the start of the
  -- stream.
  frameReaderTime :: (MonadIO m, MonadError String m)
                  => AVPixelFormat -> InputSource
                  -> m (IO (Maybe (AVFrame, Double)), IO ())
  frameReaderTime dstFmt src =
    do inputContext <- openInput src
       checkStreams inputContext
       (vidStreamIndex, ctx, cod, vidStream) <- findVideoStream inputContext
       _ <- openCodec ctx cod
       (reader, cleanup) <- prepareReader inputContext vidStreamIndex dstFmt ctx
       AVRational num den <- liftIO $ getTimeBase vidStream
       let (numl, dend) = (fromIntegral num, fromIntegral den)
           frameTime' frame =
             do n <- getPts frame
                return $ fromIntegral (n * numl) / dend
           readTS = do frame <- reader
                       case frame of
                         Nothing -> return Nothing
                         Just f -> do t <- frameTime' f
                                      return $ Just (f, t)
       return (readTS, cleanup)
  
  frameAudioReader :: (MonadIO m, MonadError String m)
                   => InputSource -> m (AudioStream, IO (Maybe AVFrame), IO ())
  frameAudioReader fileName = do
    inputContext <- openInput fileName
    checkStreams inputContext
    (audioStreamIndex, ctx, cod, audioStream) <- findAudioStream inputContext
    openCodec ctx cod
    as <- liftIO $ do
      bitrate <- getBitRate ctx
      samplerate <- getSampleRate ctx
      channelLayout <- getChannelLayout ctx
      sampleFormat <- getSampleFormat ctx
      channels <- getChannels ctx
      codecId <- getCodecID cod
      return $ AudioStream
          { asBitRate = bitrate
          , asSampleRate = samplerate
          , asSampleFormat = sampleFormat
          , asChannelLayout = channelLayout
          , asChannelCount = channels
          , asCodec = codecId
          }
    (readFrame, finalize) <- prepareAudioReader inputContext audioStreamIndex ctx
    return (as, readFrame, finalize)
  
  -- | Read time stamped RGB frames with the result in the 'MaybeT'
  -- transformer.
  --
  -- > frameReaderT = fmap (first MaybeT) . frameReader
  frameReaderTimeT :: (Functor m, MonadIO m, MonadError String m)
                   => InputSource -> m (MaybeT IO (AVFrame, Double), IO ())
  frameReaderTimeT = fmap (first MaybeT) . frameReaderTime avPixFmtRgb24
  
  prepareAudioReader :: (MonadIO m, MonadError String m)
                     => AVFormatContext -> CInt -> AVCodecContext
                     -> m (IO (Maybe AVFrame), IO ())
  prepareAudioReader fmtCtx audStream codCtx =
    wrapIOError $ do
      frame <- frame_alloc_check
      pkt <- av_packet_alloc
      let cleanup = do with frame av_frame_free
                       _ <- codec_close codCtx
                       with fmtCtx close_input
                       free (getPtr pkt)
          getFrame = do
            ret <- avcodec_receive_frame codCtx frame
            if ret < 0
              then do
                r <- av_read_frame fmtCtx pkt
                if r < 0
                  then return Nothing
                  else do
                    whichStream <- getStreamIndex pkt
                    if whichStream == audStream
                      then do
                        runWithError "Error sending packet" (avcodec_send_packet codCtx pkt)
                        getFrame
                      else free_packet pkt >> getFrame
              else return $ Just frame
      return (getFrame `catchError` const (return Nothing), cleanup)
  
  -- | Construct an action that gets the next available frame, and an
  -- action to release all resources associated with this video stream.
  prepareReader :: (MonadIO m, MonadError String m)
                => AVFormatContext -> CInt -> AVPixelFormat -> AVCodecContext
                -> m (IO (Maybe AVFrame), IO ())
  prepareReader fmtCtx vidStream dstFmt codCtx =
    wrapIOError $
    do fRaw <- frame_alloc_check
       fRgb <- frame_alloc_check
  
       w <- getWidth codCtx
       h <- getHeight codCtx
       fmt <- getPixelFormat codCtx
  
       setWidth fRgb w
       setHeight fRgb h
       setPixelFormat fRgb dstFmt
  
       frame_get_buffer_check fRgb 32
  
       sws <- swsInit (ImageInfo w h fmt) (ImageInfo w h dstFmt) swsBilinear
  
       pkt <- AVPacket <$> mallocBytes packetSize
       let cleanup = do with fRgb av_frame_free
                        with fRaw av_frame_free
                        _ <- codec_close codCtx
                        with fmtCtx close_input
                        free (getPtr pkt)
           getFrame = do
             read_frame_check fmtCtx pkt
             whichStream <- getStreamIndex pkt
             if whichStream == vidStream
             then do
               fin <- alloca $ \finished -> do
                        _ <- decode_video codCtx fRaw finished pkt
                        peek finished
               if fin > 0
               then do
                 -- Some streaming codecs require a final flush with
                 -- an empty packet
                 -- fin' <- alloca $ \fin2 -> do
                 --           free_packet pkt
                 --           (#poke AVPacket, data) pkt nullPtr
                 --           (#poke AVPacket, size) pkt (0::CInt)
                 --           decode_video codCtx fRaw fin2 pkt
                 --           peek fin2
  
                 _ <- swsScale sws fRaw fRgb
  
                 -- Copy the raw frame's timestamp to the RGB frame
                 getPktPts fRaw >>= setPts fRgb
  
                 free_packet pkt
                 return $ Just fRgb
               else free_packet pkt >> getFrame
             else free_packet pkt >> getFrame
       return (getFrame `catchError` const (return Nothing), cleanup)
  