{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DeriveDataTypeable #-}
-- | Detect the image size without opening the image itself.
--
-- Amazingly, reimplementing this wheel is the recommended way of
-- not processing huge images according to gd's FAQ.  We hope to
-- be at least as restrictive as gd itself is, otherwise some
-- malicious image could get past this code and blow on gd's
-- hand.
--
-- This code has been ressurected by the now-deprecated
-- @imagesize-conduit@ library by Michael Snoyman.
module Graphics.ThumbnailPlus.ImageSize
  ( Size (..)
  , FileFormat (..)
  , sinkImageInfo
  ) where

import qualified Data.ByteString.Lazy as L
import Data.Conduit
import qualified Data.Conduit.Binary as CB
import qualified Data.ByteString as S
import Data.ByteString.Char8 ()
import Data.ByteString.Lazy.Char8 ()
import qualified Data.Typeable as T
import Control.Applicative ((<$>), (<*>))


data Size = Size { width :: Int, height :: Int }
  deriving (Show, Eq, Ord, Read, T.Typeable)


data FileFormat = GIF | PNG | JPG
  deriving (Show, Eq, Ord, Read, Enum, T.Typeable)


-- | Find out the size of an image.  Also returns the file format
-- that parsed correctly.  Note that this function does not
-- verify that the file is indeed in the format that it returns,
-- since it looks only at a small part of the header.
sinkImageInfo :: Monad m => Consumer S.ByteString m (Maybe (Size, FileFormat))
sinkImageInfo = start id
  where
    start front = await >>= maybe (return Nothing) (pushHeader front)

    pushHeader front bs'
      | S.length bs >= 11 && S.take 5 (S.drop 6 bs) ==
        S.pack [0x4A, 0x46, 0x49, 0x46, 0x00] =
          leftover (S.drop 4 bs) >> jpg
      | S.length bs >= 6 && S.take 6 bs `elem` gifs =
        leftover (S.drop 6 bs) >> gif
      | S.length bs >= 8 && S.take 8 bs == S.pack [137, 80, 78, 71, 13, 10, 26, 10] =
        leftover (S.drop 8 bs) >> png
      | S.length bs < 11 = start $ S.append bs
      | otherwise = leftover bs >> return Nothing
      where
      bs = front bs'

    gifs = ["GIF87a", "GIF89a"]
    gif = do
      b <- CB.take 4
      let go x y = fromIntegral x + (fromIntegral y) * 256
      return $ case L.unpack b of
        [w1, w2, h1, h2] -> Just (Size (go w1 w2) (go h1 h2), GIF)
        _ -> Nothing

    png = do
      CB.drop 4
      hdr <- CB.take 4
      if hdr == "IHDR"
        then do
          mw <- getInt 4 0
          mh <- getInt 4 0
          return $ (\w h -> (Size w h, PNG)) <$> mw <*> mh
        else return Nothing

    jpg = do
      mi <- getInt 2 0
      case mi of
        Nothing -> return Nothing
        Just i -> do
          CB.drop $ i - 2
          jpgFrame

    jpgFrame = do
      mx <- CB.head
      case mx of
        Just 255 -> do
          my <- CB.head
          case my of
            Just 0xC0 -> do
              _  <- CB.take 3
              mh <- getInt 2 0
              mw <- getInt 2 0
              return $ (\w h -> (Size w h, JPG)) <$> mw <*> mh
            Just _ -> jpg
            Nothing -> return Nothing
        _ -> return Nothing

getInt :: (Monad m, Integral i)
     => Int
     -> i
     -> Consumer S.ByteString m (Maybe i)
getInt 0 i = return $ Just i
getInt len i = do
  mx <- CB.head
  case mx of
    Nothing -> return Nothing
    Just x -> getInt (len - 1) (i * 256 + fromIntegral x)