{-# LANGUAGE FlexibleContexts, ScopedTypeVariables, TypeSynonymInstances #-}
-- | Convert between FFmpeg frames and JuicyPixels images.
module Codec.FFmpeg.Juicy where
import Codec.Picture
import Codec.FFmpeg.Common
import Codec.FFmpeg.Decode
import Codec.FFmpeg.Encode
import Codec.FFmpeg.Enums
import Codec.FFmpeg.Internal.Linear (V2(..))
import Codec.FFmpeg.Types
import Control.Applicative
import Control.Arrow (first, (&&&))
import Control.Monad ((>=>))
import Control.Monad.Except
import Control.Monad.IO.Class
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe
import Data.Foldable (traverse_)
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as VM
import Foreign.C.Types
import Foreign.Marshal.Array (advancePtr, copyArray)
import Foreign.Ptr (castPtr, Ptr)
import Foreign.Storable (sizeOf)
import Data.Maybe (maybe)


-- | Convert 'AVFrame' to a 'Vector'.
frameToVector :: AVFrame -> IO (Maybe (V.Vector CUChar))
frameToVector = runMaybeT . frameToVectorT


-- | Convert 'AVFrame' to a 'Vector' with the result in the 'MaybeT' transformer.
frameToVectorT :: AVFrame -> MaybeT IO (V.Vector CUChar)
frameToVectorT frame = do

  bufSize <- fromIntegral <$> frameBufferSizeT frame
       
  v <- MaybeT $ do

         v <- VM.new bufSize

         VM.unsafeWith v (frameCopyToBuffer frame)
           >>= return . maybe Nothing (const (Just v))

  lift $ V.unsafeFreeze v


-- | Convert an 'AVFrame' to a 'DynamicImage' with the result in the
-- 'MaybeT' transformer.
--
-- > toJuicyT = MaybeT . toJuicy
toJuicyT :: AVFrame -> MaybeT IO DynamicImage
toJuicyT = MaybeT . toJuicy


-- | Convert an 'AVFrame' to a 'DynamicImage'.
toJuicy :: AVFrame -> IO (Maybe DynamicImage)
toJuicy frame = runMaybeT $ do

  v <- frameToVectorT frame

  MaybeT $ do
  
    w <- fromIntegral <$> getWidth frame
    h <- fromIntegral <$> getHeight frame
    
    let mkImage :: V.Storable (PixelBaseComponent a)
                => (Image a -> DynamicImage) -> Maybe DynamicImage
        mkImage c = Just $ c (Image w h (V.unsafeCast v))
  
    fmt <- getPixelFormat frame
        
    return $ case () of
               _ | fmt == avPixFmtRgb24 -> mkImage ImageRGB8
                 | fmt == avPixFmtGray8 -> mkImage ImageY8
                 | fmt == avPixFmtGray16 -> mkImage ImageY16
                 | otherwise -> Nothing


-- | Convert an 'AVFrame' to an 'Image'.
toJuicyImage :: forall p. JuicyPixelFormat p => AVFrame -> IO (Maybe (Image p))
toJuicyImage frame = runMaybeT $ do
  
  fmt <- lift $ getPixelFormat frame
  guard (fmt == juicyPixelFormat ([] :: [p]))
          
  MaybeT $ do
    
    w <- fromIntegral <$> getWidth frame
    h <- fromIntegral <$> getHeight frame
    
    fmap (Image w h . V.unsafeCast) <$> frameToVector frame


-- | Save an 'AVFrame' to a PNG file on disk assuming the frame could
-- be converted to a 'DynamicImage' using 'toJuicy'.
saveJuicy :: FilePath -> AVFrame -> IO ()
saveJuicy name = toJuicy >=> traverse_ (savePngImage name)


-- | Mapping of @JuicyPixels@ pixel types to FFmpeg pixel formats.
class Pixel a => JuicyPixelFormat a where
  juicyPixelFormat :: proxy a -> AVPixelFormat

instance JuicyPixelFormat Pixel8 where
  juicyPixelFormat _ = avPixFmtGray8

instance JuicyPixelFormat PixelRGB8 where
  juicyPixelFormat _ = avPixFmtRgb24

instance JuicyPixelFormat PixelRGBA8 where
  juicyPixelFormat _ = avPixFmtRgba

-- | Bytes-per-pixel for a JuicyPixels 'Pixel' type.
juicyPixelStride :: forall a proxy. Pixel a => proxy a -> Int
juicyPixelStride _ =
  sizeOf (undefined :: PixelBaseComponent a) * componentCount (undefined :: a)

-- | Read frames from a video stream.
imageReaderT :: forall m p.
                (Functor m, MonadIO m, MonadError String m,
                 JuicyPixelFormat p)
             => InputSource -> m (IO (Maybe (Image p)), IO ())
imageReaderT = fmap (first (runMaybeT . aux toJuicyImage))
            . frameReader (juicyPixelFormat ([] :: [p]))
  where aux g x = MaybeT x >>= MaybeT . g

-- | Read frames from a video stream. Errors are thrown as
-- 'IOException's.
imageReader :: JuicyPixelFormat p
            => InputSource -> IO (IO (Maybe (Image p)), IO ())
imageReader = (>>= either error return) . runExceptT . imageReaderT

-- | Read time stamped frames from a video stream. Time is given in
-- seconds from the start of the stream.
imageReaderTimeT :: forall m p.
                    (Functor m, MonadIO m, MonadError String m,
                     JuicyPixelFormat p)
                 => InputSource -> m (IO (Maybe (Image p, Double)), IO ())
imageReaderTimeT = fmap (first (runMaybeT . aux toJuicyImage))
                 . frameReaderTime (juicyPixelFormat ([] :: [p]))
  where aux g x = do (f,t) <- MaybeT x
                     f' <- MaybeT $ g f
                     return (f', t)

-- | Read time stamped frames from a video stream. Time is given in
-- seconds from the start of the stream. Errors are thrown as
-- 'IOException's.
imageReaderTime :: JuicyPixelFormat p
                => InputSource -> IO (IO (Maybe (Image p, Double)), IO ())
imageReaderTime = (>>= either error return) . runExceptT . imageReaderTimeT

-- | Open a target file for writing a video stream. When the returned
-- function is applied to 'Nothing', the output stream is closed. Note
-- that 'Nothing' /must/ be provided when finishing in order to
-- properly terminate video encoding.
--
-- Support for source images that are of a different size to the
-- output resolution is limited to non-palettized destination formats
-- (i.e. those that are handled by @libswscaler@). Practically, this
-- means that animated gif output is only supported if the source
-- images are of the target resolution.
imageWriter :: forall p. JuicyPixelFormat p
            => EncodingParams -> FilePath -> IO (Maybe (Image p) -> IO ())
imageWriter ep f = (. fmap aux) <$> frameWriter ep f
  where aux img = let w = fromIntegral $ imageWidth img
                      h = fromIntegral $ imageHeight img
                      p = V.unsafeCast $ imageData img
                  in  (juicyPixelFormat ([]::[p]), V2 w h, p)