{-
  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
import qualified Media.FFMpeg as FF
import qualified Media.FFMpeg.SWScale as S

import System.Environment (getArgs, getProgName)
import System.IO
import Text.Printf (printf)
import Data.Version (showVersion)
import Data.Maybe (fromJust)
import Data.IORef (newIORef, readIORef, modifyIORef)
import Control.Monad (when)
import System.Exit (exitSuccess)
import Control.Concurrent.STM
import Data.ByteString (ByteString, empty, take, drop
                       ,length, append, useAsCStringLen) 
import Foreign
    
-- SDL
import qualified Graphics.UI.SDL as SDL

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"]

data CStreamProps = CStreamProps {
      streamIdx :: Int
    , stream :: FF.Stream
    , codecCtx :: FF.CodecContext
    , codecId :: FF.CodecId
    , decoder :: FF.Codec
    }

instance Show CStreamProps where
    show pr = printf "Stream #%d, codec '%s'\n" (streamIdx pr) (show (codecId pr))

data VStreamProps = VStreamProps {
      width :: Int
    , height :: Int
    , pixFmt :: FF.PixelFormat
    }

instance Show VStreamProps where
    show pr = printf "Format %dx%d '%s'\n" (width pr) (height pr) (show (pixFmt pr))

data AStreamProps = AStreamProps {
      sampleRate :: Int
    , channels :: Int
    , sampleFmt :: FF.SampleFormat
    }

instance Show AStreamProps where
    show pr = printf "Format rate: %d, channels: %d, format: '%s'\n" 
              (sampleRate pr) (channels pr) (show (sampleFmt pr))
                   
data StreamProps = VideoStream CStreamProps VStreamProps 
                 | AudioStream CStreamProps AStreamProps deriving (Show)

commonProps (AudioStream a _) = a
commonProps (VideoStream a _) = a
audioProps (AudioStream _ a) = a
videoProps (VideoStream _ a) = a

requestForAudioData :: 
    FF.CodecContext -> TChan FF.Packet -> TVar ByteString -> Ptr () -> Ptr Word8 -> Word32 -> IO ()
requestForAudioData ctx pc buff _ dst size = do
  buffer <- retrieveBuffer ctx pc buff (fromIntegral size)
  useAsCStringLen buffer $ \(p, s) ->
      copyArray dst (castPtr p :: Ptr Word8) s
              

-- combine together two TChans
retrieveBuffer ctx pc buff s = do
  e <- atomically $ do
         bs <- readTVar buff
         return $ Data.ByteString.length bs
  if e < s
   then do 
     pk <- atomically $ readTChan pc
     buffs <- FF.decodeAudioPacket' ctx pk
     atomically $ do
          bs <- readTVar buff
          writeTVar buff (bs `append` buffs)
     retrieveBuffer ctx pc buff s
   else do
     atomically $ do
          bs <- readTVar buff
          writeTVar buff (Data.ByteString.drop s bs)
          return $ Data.ByteString.take s bs

decodeFile :: String -> IO ()
decodeFile fname = do
  SDL.init [SDL.InitVideo, SDL.InitAudio, SDL.InitTimer]

  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 fname
    FF.findStreamInfo ctx

    printf "File: %s\n" fname

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

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

    let audioStreams = filter (\(_,s) -> maybe False ((==FF.CodecTypeAudio) . FF.getCodecType) (FF.getCodecContext s)) $ zip [0..] streams

    if null videoStreams || null audioStreams 
     then printf "No video or audio streams found\n"
     else do
       vStreamProps <- getVStreamProps $ head videoStreams
       print vStreamProps

       aStreamProps <- getAStreamProps $ head audioStreams
       print aStreamProps

       packetChain <- atomically $ (newTChan :: STM (TChan FF.Packet))
       buffer <- atomically $ (newTVar Data.ByteString.empty)

       fcb <- SDL.mkAudioBufferFillCb $ requestForAudioData (codecCtx $ commonProps aStreamProps) packetChain buffer

       let wspecs = SDL.AudioSpec {
                      SDL.freq = (sampleRate . audioProps) aStreamProps
                    , SDL.format = SDL.AudioS16Sys
                    , SDL.channels = (fromIntegral . channels . audioProps) aStreamProps
                    , SDL.silence = 0
                    , SDL.samples = 4096
                    , SDL.size = 0
                    , SDL.callback = fcb}
           
       (Just specs) <- SDL.openAudio wspecs False
       printf "Output audio specs: %s\n" (show specs)

       let (w, h, pf, cc, vsidx) = ((width . videoProps) vStreamProps
                                   ,(height . videoProps) vStreamProps
                                   ,(pixFmt . videoProps) vStreamProps
                                   ,(codecCtx . commonProps) vStreamProps
                                   ,(streamIdx . commonProps) vStreamProps)

       screen <- SDL.setVideoMode w h 0 []
       rgbSurface <- SDL.createRGBSurface [] w h 24 0x0000FF 0x00FF00 0xFF0000 0x000000
       frame <- FF.allocFrame
       frameRGB <- FF.allocFrame

       surfacePixels <- SDL.surfaceGetPixels rgbSurface >>= FF.castBuffer

       FF.pictureFill frameRGB surfacePixels FF.PixFmtRgb24 w h

       -- packet <- FF.allocPacket

       -- Prepare resizer
       scaler <- S.getContext (w, h, pf) 
                 (w, h, FF.PixFmtRgb24) [S.SwsBicubic]

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

       SDL.pauseAudio False
                 
       let asidx = (streamIdx . commonProps) aStreamProps

       while (FF.readFrame ctx) $ \packet ->
             do
               if (asidx == FF.packetGetStreamIndex packet)
                then do
                  FF.dupPacket packet
                  atomically $ writeTChan packetChain packet
                  -- abufs <- FF.decodeAudioPacket' ((codecCtx . commonProps) aStreamProps)  packet
                  -- FF.cleanPacket packet
                  return ()
                else if vsidx == FF.packetGetStreamIndex packet 
                      then 
                          do
                            frameFinished <- FF.decodeVideo cc frame packet
                            when frameFinished $ 
                                 do
                                   -- Rescale the picture
                                   
                                   -- SDL.lockSurface rgbSurface
                                   
                                   S.scale scaler 
                                        (FF.pictureGetSlice frame) (FF.pictureGetStride frame)
                                        0 h
                                        (FF.pictureGetSlice frameRGB) (FF.pictureGetStride frameRGB)

                                   -- show on screen 
                                   SDL.unlockSurface rgbSurface
                                   SDL.blitSurface rgbSurface Nothing screen Nothing
                                   SDL.flip screen

                                   idx <- readIORef idxRef
                                   modifyIORef idxRef (+1)
                                   printf "\rDecoded frame %05d" (idx)
                                   
                      else return ()

                           

               ev <- SDL.pollEvent
               case ev of
                 SDL.Quit -> do SDL.quit 
                                exitSuccess
                 otherwise -> return ()

       printf "\n"

-- While cycle implementation
while :: (FF.Packet -> IO Bool) -> (FF.Packet -> IO ()) -> IO ()
while cond body = do
  pkt <- FF.allocPacket
  c <- cond pkt
  when c $ body pkt >> while cond body

getStreamProps (sidx, s) = do
  let codecContext = fromJust $ FF.getCodecContext s
  let cId = FF.getCodecId codecContext
  (Just dec) <- FF.findDecoder cId
  FF.openCodec codecContext dec
  return $ CStreamProps {
               streamIdx = sidx
             , stream = s
             , codecCtx = codecContext
             , codecId = cId
             , decoder = dec }

getVStreamProps (sidx, s) = do
  sp <- getStreamProps (sidx, s)
  let [w, h] = map ($ (codecCtx sp)) 
               [FF.getVideoWidth, FF.getVideoHeight]
  let pf = FF.getPixelFormat (codecCtx sp)
  return $ VideoStream sp $ VStreamProps {width = w, height = h, pixFmt = pf}

getAStreamProps (sidx, s) = do
  sp <- getStreamProps (sidx, s)
  let [sr, ch] = map ($ (codecCtx sp))
                 [FF.getAudioSampleRate, FF.getAudioChannels]
  let sf = FF.getAudioSampleFormat (codecCtx sp)
  return $ AudioStream sp $ AStreamProps {sampleRate = sr, channels = ch, sampleFmt = sf}
