{-
  Testing the ffmpeg video decoding.

  Rewritten from the C code, from the FFMpeg tutorial 
  (http://www.dranger.com/ffmpeg/)
  
  (c) 2009 Vasyl Pasternak
 -}

module Main (main) where

-- FFMpeg imports 
import qualified Media.FFMpeg as FF
import qualified Media.FFMpeg.SWScale as S

-- Common imports
import Data.Version
import System.Environment (getArgs, getProgName)
import Text.Printf (printf)
import Data.Maybe (fromJust)
import Data.IORef (newIORef, readIORef, modifyIORef)
import Control.Monad (when)
import System.IO -- (WriteMode, withBinaryFile, hPutStr, hPutBuf)

main :: IO ()
main = do
  getArgs >>= \args ->
      case args of
        filename:[] -> decodeFile filename
        otherwise -> getProgName >>= \pn -> 
                     (putStrLn . unlines) 
                     ["Usage: \n"
                     ,printf "%s filename\n" pn
                     ,"Where <filename> is the path to the video file"]
                      

-- The main rutine, do all work
decodeFile :: String -> IO ()
decodeFile filename = 
    FF.withFFMpeg $ do 
      -- Print FFMpeg version
      printf "FFMpeg library initialized\n" 
      printf "\tlibAVFormat version: %s\n" (showVersion FF.libAVFormatVersion)
      printf "\tlibAVCodec version: %s\n" (showVersion FF.libAVCodecVersion)

      -- Open file
      ctx <- FF.openInputFile filename
      FF.findStreamInfo ctx
      -- FF.dumpFormat ctx filename
        
      printf "File: %s\n" filename

      -- Get all streams from file
      let streams = FF.getStreams ctx
      printf "Found %d streams\n" (length streams)

      -- Retrive video streams
      let videoStreams = filter (\(_,s) -> maybe False ((==FF.CodecTypeVideo) . FF.getCodecType) (FF.getCodecContext s)) $ zip [0..] streams

      if null videoStreams 
       then do 
         printf "No video streams found. Exiting...\n"
       else do
         let (streamIdx, videoStream) = head videoStreams
         printf "Extracting video information from stream %d\n" (streamIdx :: Int) 
        
         -- Starting decoding video
         let codecContext = fromJust $ FF.getCodecContext videoStream
         let codecId = FF.getCodecId codecContext
         printf "Video codec is '%s'\n" (show codecId)

         (Just decoder) <- FF.findDecoder codecId
         FF.openCodec codecContext decoder

         let (width, height, pixfmt) = (FF.getVideoWidth codecContext, 
                                        FF.getVideoHeight codecContext,
                                        FF.getPixelFormat codecContext)

         printf "Video dimensions are %dx%d\n" width height

         printf "Video pixel format is %s\n" (show pixfmt)


         frame <- FF.allocFrame
         frameRGB <- FF.allocFrame
         
         let pictureSize = FF.pictureGetSize FF.PixFmtRgb24 width height
    
         buffer <- FF.allocBuffer pictureSize
         
         FF.pictureFill frameRGB buffer FF.PixFmtRgb24 width height

         packet <- FF.allocPacket

         -- Prepare resizer
         scaler <- S.getContext (width, height, pixfmt) 
                    (width, height, FF.PixFmtRgb24) [S.SwsBicubic]

         -- use IORef to cound packets
         idxRef <- newIORef (0 :: Int)

         while (FF.readFrame ctx packet) $ 
               do
                 when (streamIdx == FF.packetGetStreamIndex packet) $
                    do
                      frameFinished <- FF.decodeVideo codecContext frame packet
                      when frameFinished $ 
                           do
                             -- Rescale the picture

                             S.scale scaler 
                              (FF.pictureGetSlice frame) (FF.pictureGetStride frame)
                              0 height
                              (FF.pictureGetSlice frameRGB) (FF.pictureGetStride frameRGB)

                             idx <- readIORef idxRef
                             modifyIORef idxRef (+1)
                             printf "\rDecoded frame %05d" (idx)
                             when ((idx `mod` 100) == 0) $ 
                                  writePPM (printf "frame%05d.ppm" idx) width height
                                           buffer pictureSize
                 FF.cleanPacket packet

         printf "\n"


-- While cycle implementation
while :: IO Bool -> IO () -> IO ()
while cond body = do
  c <- cond
  when c $ body >> while cond body

writePPM fname width height buff b = 
    withBinaryFile fname WriteMode $ \h -> do
      hPutStr h (printf "P6\n%d %d\n255\n" width height)
      FF.withBuffer buff $ \buff' -> 
          hPutBuf h buff' b