{-# OPTIONS -fglasgow-exts #-}
module Main where

import Control.Monad
import Data.Word

import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable

import System.Environment

import qualified AVCodec as AVCodec
import qualified AVFormat as AV

import DrawFrame
import GLX

main =
    do [filePath] <- getArgs
       AV.registerAll
       let filename = filePath -- "/root/media/liz.3gp" -- "fifteen.3gp"  -- "/root/media/circus.avi" --  
       allocaBytes (sizeOf (undefined :: Ptr (Ptr ()))) $ \avFormatContextPtr ->
           do ret <- AV.openInputFile avFormatContextPtr filename Nothing 0 Nothing
              when (ret /= 0) (error $ "openInputFile failed: " ++ show ret)
              avFormatContext <- peek ((castPtr avFormatContextPtr) :: Ptr (Ptr ())) >>= return . AV.AVFormatContext . castPtr
              ret <- AV.findStreamInfo avFormatContext
              when (ret < 0) (error $ "findStreamInfo failed: " ++ show ret)
              AV.dumpFormat avFormatContext 0 filename False
              streams <- liftM (zip [0..]) $ AV.streams avFormatContext 
--              mapM (\strm -> AV.codec strm >>= AVCodec.codecName) streams >>= print
              videoStreams <- filterM (\(_, stream) -> AV.codec stream >>= AVCodec.codecType  >>= \ct -> return (ct == AVCodec.codecTypeVideo)) streams
              when (length videoStreams < 1) (error $ filename ++" does not appear to contain any video streams.")
              let videoStream = (head videoStreams)
              codecCtx <- AV.codec (snd videoStream)
              decoder <- maybe (error "no decoder found for video") return =<< AVCodec.findDecoder =<< AVCodec.codecId codecCtx
              -- skipped part about truncated
              ret <- AVCodec.open codecCtx decoder
              when (ret < 0) (error $ "failed to open decoder.")
              -- skipped hack for incorrect frame rate
              frame@(AVCodec.AVFrame framePtr) <- AVCodec.allocFrame
              when (framePtr == nullPtr) (error $ "allocFrame failed")
              -- allocate RGB frame
              frameRGB@(AVCodec.AVFrame framePtr) <- AVCodec.allocFrame
              when (framePtr == nullPtr) (error $ "allocFrame failed")
              w <- AVCodec.pictureWidth codecCtx
              h <- AVCodec.pictureHeight codecCtx
              fmt <- AVCodec.pixFmt codecCtx
              print ("codecCtx size", w, h)
              let numBytes = AVCodec.pictureGetSize AVCodec.pixFmtRgb24 w h
              buffer <- mallocBytes (fromIntegral numBytes)
              AVCodec.pictureFill (AVCodec.castAVFrameToAVPicture frameRGB) buffer AVCodec.pixFmtRgb24 w h
              AV.allocaPacket $ \packet ->
                  let renderLoop r bytesRemaining rawData t dpy win =
                          do (finished, bytesRemaining', rawData') <- getNextFrame avFormatContext codecCtx (fst videoStream) frame packet bytesRemaining rawData
--                             print ("finished frame", finished)
                             AVCodec.imgConvert (AVCodec.castAVFrameToAVPicture frameRGB) AVCodec.pixFmtRgb24 (AVCodec.castAVFrameToAVPicture frame) fmt w h
                             mapM_ (\r' -> glDrawFrame r' t frameRGB numBytes w h >> swapBuffers dpy win) [r, r + 0.4 .. r + (0.4 * 1)]
                             renderLoop (r + 0.8) bytesRemaining' rawData' t dpy win
                  in
                    xInit $ \dpy win -> renderLoop 0.0 0 nullPtr dpy win
              -- free some stuff 
              return ()

getNextFrame :: AV.AVFormatContext -> AVCodec.AVCodecContext -> Int -> AVCodec.AVFrame -> AV.AVPacket -> Int -> Ptr Word8 -> IO (Bool, Int, Ptr Word8)
getNextFrame fmtCtx codecCtx streamIndex frame packet 0 rawData =
    do getNextPacket streamIndex fmtCtx packet
       bytesRemaining <- AV.packetSize packet
       rawData <- AV.packetData packet
       getNextFrame fmtCtx codecCtx streamIndex frame packet bytesRemaining rawData
getNextFrame fmtCtx codecCtx streamIndex frame packet bytesRemaining rawData =
    alloca $ \finishedPtr ->
           do bytesDecoded <- AVCodec.decodeVideo codecCtx frame finishedPtr rawData bytesRemaining
              when (bytesDecoded < 0) (error "Error decoding frame") -- FIXME: this is not really fatal
              finished <- peek finishedPtr :: IO Int
              if finished /= 0
                 then return (True, bytesRemaining - bytesDecoded, rawData `plusPtr` bytesDecoded)
                 else getNextFrame fmtCtx codecCtx streamIndex frame packet (bytesRemaining - bytesDecoded) (rawData `plusPtr` bytesDecoded)

getNextPacket :: Int -> AV.AVFormatContext -> AV.AVPacket -> IO ()
getNextPacket streamIndex fmtCtx packet =
    do pdata <- AV.packetData packet
       when (pdata /= nullPtr) (AV.freePacket packet)
       ret <- AV.readPacket fmtCtx packet
       when (ret >= 0) $
            do si <- AV.packetStreamIndex packet
               if si == streamIndex
                  then return ()
                  else getNextPacket streamIndex fmtCtx packet
       
--     do (bytesDecoded, frameFinished) <- AVCodec.decodeVideo codecCtx frame 

{-
assert :: String -> Int -> IO ()
assert msg 0 = return ()
assert msg n = error $ msg ++ show n
-}