content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE CPP #-}
  
  -- | A good explanation of the JPEG format, including diagrams, is given at:
  -- <https://github.com/corkami/formats/blob/master/image/jpeg.md>
  --
  -- The full spec (excluding EXIF): https://www.w3.org/Graphics/JPEG/itu-t81.pdf
  module Codec.Picture.Jpg.Internal.Types( MutableMacroBlock
                                , createEmptyMutableMacroBlock
                                , printMacroBlock
                                , printPureMacroBlock
                                , DcCoefficient
                                , JpgImage( .. )
                                , JpgComponent( .. )
                                , JpgFrameHeader( .. )
                                , JpgFrame( .. )
                                , JpgFrameKind( .. )
                                , JpgScanHeader( .. )
                                , JpgQuantTableSpec( .. )
                                , JpgHuffmanTableSpec( .. )
                                , JpgImageKind( .. )
                                , JpgScanSpecification( .. )
                                , JpgColorSpace( .. )
                                , AdobeTransform( .. )
                                , JpgAdobeApp14( .. )
                                , JpgJFIFApp0( .. )
                                , JFifUnit( .. )
                                , TableList( .. )
                                , RestartInterval( .. )
                                , getJpgImage
                                , calculateSize
                                , dctBlockSize
                                , parseECS
                                , parseECS_simple
                                , skipUntilFrames
                                , skipFrameMarker
                                , parseFrameOfKind
                                , parseFrames
                                , parseFrameKinds
                                , parseToFirstFrameHeader
                                ) where
  
  
  #if !MIN_VERSION_base(4,8,0)
  import Control.Applicative( pure, (<*>), (<$>) )
  #endif
  
  import Control.DeepSeq( NFData(..) )
  import Control.Monad( when, replicateM, forM, forM_, unless )
  import Control.Monad.ST( ST )
  import Data.Bits( (.|.), (.&.), unsafeShiftL, unsafeShiftR )
  import Data.List( partition )
  import Data.Maybe( maybeToList )
  import GHC.Generics( Generic )
  
  #if !MIN_VERSION_base(4,11,0)
  import Data.Monoid( (<>) )
  #endif
  
  import Foreign.Storable ( Storable )
  import Data.Vector.Unboxed( (!) )
  import qualified Data.Vector as V
  import qualified Data.Vector.Unboxed as VU
  import qualified Data.Vector.Storable as VS
  import qualified Data.Vector.Storable.Mutable as M
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Char8 as BC
  import qualified Data.ByteString.Lazy as L
  import qualified Data.ByteString.Unsafe as BU
  
  import Data.Int( Int16, Int64 )
  import Data.Word(Word8, Word16 )
  import Data.Binary( Binary(..) )
  
  import Data.Binary.Get( Get
                        , getWord8
                        , getWord16be
                        , getByteString
                        , skip
                        , bytesRead
                        , lookAhead
                        , ByteOffset
                        , getLazyByteString
                        )
  import qualified Data.Binary.Get.Internal as GetInternal
  
  import Data.Binary.Put( Put
                        , putWord8
                        , putWord16be
                        , putLazyByteString
                        , putByteString
                        , runPut
                        )
  
  import Codec.Picture.InternalHelper
  import Codec.Picture.Jpg.Internal.DefaultTable
  import Codec.Picture.Tiff.Internal.Types
  import Codec.Picture.Tiff.Internal.Metadata( exifOffsetIfd )
  import Codec.Picture.Metadata.Exif
  
  import Text.Printf
  
  -- | Type only used to make clear what kind of integer we are carrying
  -- Might be transformed into newtype in the future
  type DcCoefficient = Int16
  
  -- | Macroblock that can be transformed.
  type MutableMacroBlock s a = M.STVector s a
  
  data JpgFrameKind =
        JpgBaselineDCTHuffman
      | JpgExtendedSequentialDCTHuffman
      | JpgProgressiveDCTHuffman
      | JpgLosslessHuffman
      | JpgDifferentialSequentialDCTHuffman
      | JpgDifferentialProgressiveDCTHuffman
      | JpgDifferentialLosslessHuffman
      | JpgExtendedSequentialArithmetic
      | JpgProgressiveDCTArithmetic
      | JpgLosslessArithmetic
      | JpgDifferentialSequentialDCTArithmetic
      | JpgDifferentialProgressiveDCTArithmetic
      | JpgDifferentialLosslessArithmetic
      | JpgQuantizationTable
      | JpgHuffmanTableMarker
      | JpgStartOfScan
      | JpgEndOfImage
      | JpgAppSegment Word8
      | JpgExtensionSegment Word8
  
      | JpgRestartInterval
      | JpgRestartIntervalEnd Word8
      deriving (Eq, Show, Generic)
  instance NFData JpgFrameKind
  
  data JpgFrame =
        JpgAppFrame        !Word8 B.ByteString
      | JpgAdobeAPP14      !JpgAdobeApp14
      | JpgJFIF            !JpgJFIFApp0
      | JpgExif            ![ImageFileDirectory]
      | JpgExtension       !Word8 B.ByteString
      | JpgQuantTable      ![JpgQuantTableSpec]
      | JpgHuffmanTable    ![(JpgHuffmanTableSpec, HuffmanPackedTree)]
      | JpgScanBlob        !JpgScanHeader !L.ByteString -- ^ The @ByteString@ is the ECS (Entropy-Coded Segment), typically the largest part of compressed image data.
      | JpgScans           !JpgFrameKind !JpgFrameHeader
      | JpgIntervalRestart !Word16
      deriving (Eq, Show, Generic)
  instance NFData JpgFrame
  
  data JpgColorSpace
    = JpgColorSpaceYCbCr
    | JpgColorSpaceYCC
    | JpgColorSpaceY
    | JpgColorSpaceYA
    | JpgColorSpaceYCCA
    | JpgColorSpaceYCCK
    | JpgColorSpaceCMYK
    | JpgColorSpaceRGB
    | JpgColorSpaceRGBA
    deriving (Eq, Show, Generic)
  instance NFData JpgColorSpace
  
  data AdobeTransform
    = AdobeUnknown    -- ^ Value 0
    | AdobeYCbCr      -- ^ value 1
    | AdobeYCck       -- ^ value 2
    deriving (Eq, Show, Generic)
  instance NFData AdobeTransform
  
  data JpgAdobeApp14 = JpgAdobeApp14
    { _adobeDctVersion :: !Word16
    , _adobeFlag0      :: !Word16
    , _adobeFlag1      :: !Word16
    , _adobeTransform  :: !AdobeTransform
    }
    deriving (Eq, Show, Generic)
  instance NFData JpgAdobeApp14
  
  -- | Size: 1
  data JFifUnit
    = JFifUnitUnknown   -- ^ 0
    | JFifPixelsPerInch -- ^ 1
    | JFifPixelsPerCentimeter -- ^ 2
    deriving (Eq, Show, Generic)
  instance NFData JFifUnit
  
  instance Binary JFifUnit where
    put v = putWord8 $ case v of
      JFifUnitUnknown -> 0
      JFifPixelsPerInch -> 1
      JFifPixelsPerCentimeter -> 2
    get = do
      v <- getWord8
      pure $ case v of
        0 -> JFifUnitUnknown
        1 -> JFifPixelsPerInch
        2 -> JFifPixelsPerCentimeter
        _ -> JFifUnitUnknown
  
  data JpgJFIFApp0 = JpgJFIFApp0
    { _jfifUnit      :: !JFifUnit
    , _jfifDpiX      :: !Word16
    , _jfifDpiY      :: !Word16
    , _jfifThumbnail :: !(Maybe {- (Image PixelRGB8) -} Int)
    }
    deriving (Eq, Show, Generic)
  instance NFData JpgJFIFApp0
  
  instance Binary JpgJFIFApp0 where
    get = do
      sig <- getByteString 5
      when (sig /= BC.pack "JFIF\0") $
          fail "Invalid JFIF signature"
      major <- getWord8
      minor <- getWord8
      when (major /= 1 && minor > 2) $
          fail "Unrecognize JFIF version"
      unit <- get
      dpiX <- getWord16be
      dpiY <- getWord16be
      w <- getWord8
      h <- getWord8
      let pxCount = 3 * w * h
      img <- case pxCount of
        0 -> return Nothing
        _ -> return Nothing
      return $ JpgJFIFApp0
          { _jfifUnit      = unit
          , _jfifDpiX      = dpiX
          , _jfifDpiY      = dpiY
          , _jfifThumbnail = img
          }
  
  
    put jfif = do
      putByteString $ BC.pack "JFIF\0" -- 5
      putWord8 1                       -- 1 6
      putWord8 2                       -- 1 7
      put $ _jfifUnit jfif             -- 1 8
      putWord16be $ _jfifDpiX jfif     -- 2 10
      putWord16be $ _jfifDpiY jfif     -- 2 12
      putWord8 0                       -- 1 13
      putWord8 0                       -- 1 14
  
  {-Thumbnail width (tw) 	1 	Horizontal size of embedded JFIF thumbnail in pixels-}
  {-Thumbnail height (th) 	1 	Vertical size of embedded JFIF thumbnail in pixels-}
  {-Thumbnail data 	3 × tw × th 	Uncompressed 24 bit RGB raster thumbnail-}
  
  instance Binary AdobeTransform where
    put v = case v of
      AdobeUnknown -> putWord8 0
      AdobeYCbCr -> putWord8 1
      AdobeYCck -> putWord8 2
  
    get = do
      v <- getWord8
      pure $ case v of
        0 -> AdobeUnknown
        1 -> AdobeYCbCr
        2 -> AdobeYCck
        _ -> AdobeUnknown
  
  instance Binary JpgAdobeApp14 where
    get = do
      let sig = BC.pack "Adobe"
      fileSig <- getByteString 5
      when (fileSig /= sig) $
         fail "Invalid Adobe APP14 marker"
      version <- getWord16be
      when (version /= 100) $
         fail $ "Invalid Adobe APP14 version " ++ show version
      JpgAdobeApp14 version
                    <$> getWord16be
                    <*> getWord16be <*> get
  
    put (JpgAdobeApp14 v f0 f1 t) = do
      putByteString $ BC.pack "Adobe"
      putWord16be v
      putWord16be f0
      putWord16be f1
      put t
  
  
  data JpgFrameHeader = JpgFrameHeader
      { jpgFrameHeaderLength   :: !Word16
      , jpgSamplePrecision     :: !Word8
      , jpgHeight              :: !Word16
      , jpgWidth               :: !Word16
      , jpgImageComponentCount :: !Word8
      , jpgComponents          :: ![JpgComponent]
      }
      deriving (Eq, Show, Generic)
  instance NFData JpgFrameHeader
  
  
  instance SizeCalculable JpgFrameHeader where
      calculateSize hdr = 2 + 1 + 2 + 2 + 1
                        + sum [calculateSize c | c <- jpgComponents hdr]
  
  data JpgComponent = JpgComponent
      { componentIdentifier       :: !Word8
        -- | Stored with 4 bits
      , horizontalSamplingFactor  :: !Word8
        -- | Stored with 4 bits
      , verticalSamplingFactor    :: !Word8
      , quantizationTableDest     :: !Word8
      }
      deriving (Eq, Show, Generic)
  instance NFData JpgComponent
  
  instance SizeCalculable JpgComponent where
      calculateSize _ = 3
  
  data JpgImage = JpgImage { jpgFrame :: [JpgFrame] }
      deriving (Eq, Show, Generic)
  instance NFData JpgImage
  
  data JpgScanSpecification = JpgScanSpecification
      { componentSelector :: !Word8
        -- | Encoded as 4 bits
      , dcEntropyCodingTable :: !Word8
        -- | Encoded as 4 bits
      , acEntropyCodingTable :: !Word8
  
      }
      deriving (Eq, Show, Generic)
  instance NFData JpgScanSpecification
  
  instance SizeCalculable JpgScanSpecification where
      calculateSize _ = 2
  
  data JpgScanHeader = JpgScanHeader
      { scanLength :: !Word16
      , scanComponentCount :: !Word8
      , scans :: [JpgScanSpecification]
  
        -- | (begin, end)
      , spectralSelection    :: (Word8, Word8)
  
        -- | Encoded as 4 bits
      , successiveApproxHigh :: !Word8
  
        -- | Encoded as 4 bits
      , successiveApproxLow :: !Word8
      }
      deriving (Eq, Show, Generic)
  instance NFData JpgScanHeader
  
  instance SizeCalculable JpgScanHeader where
      calculateSize hdr = 2 + 1
                        + sum [calculateSize c | c <- scans hdr]
                        + 2
                        + 1
  
  data JpgQuantTableSpec = JpgQuantTableSpec
      { -- | Stored on 4 bits
        quantPrecision     :: !Word8
  
        -- | Stored on 4 bits
      , quantDestination   :: !Word8
  
      , quantTable         :: MacroBlock Int16
      }
      deriving (Eq, Show, Generic)
  instance NFData JpgQuantTableSpec
  
  class SizeCalculable a where
      calculateSize :: a -> Int
  
  -- | Type introduced only to avoid some typeclass overlapping
  -- problem
  newtype TableList a = TableList [a]
  
  instance (SizeCalculable a, Binary a) => Binary (TableList a) where
      put (TableList lst) = do
          putWord16be . fromIntegral $ sum [calculateSize table | table <- lst] + 2
          mapM_ put lst
  
      get = TableList <$> (getWord16be >>= \s -> innerParse (fromIntegral s - 2))
        where innerParse :: Int -> Get [a]
              innerParse 0    = return []
              innerParse size = do
                  onStart <- fromIntegral <$> bytesRead
                  table <- get
                  onEnd <- fromIntegral <$> bytesRead
                  (table :) <$> innerParse (size - (onEnd - onStart))
  
  instance SizeCalculable JpgQuantTableSpec where
      calculateSize table =
          1 + (fromIntegral (quantPrecision table) + 1) * 64
  
  instance Binary JpgQuantTableSpec where
      put table = do
          let precision = quantPrecision table
          put4BitsOfEach precision (quantDestination table)
          forM_ (VS.toList $ quantTable table) $ \coeff ->
              if precision == 0 then putWord8 $ fromIntegral coeff
                               else putWord16be $ fromIntegral coeff
  
      get = do
          (precision, dest) <- get4BitOfEach
          coeffs <- replicateM 64 $ if precision == 0
                  then fromIntegral <$> getWord8
                  else fromIntegral <$> getWord16be
          return JpgQuantTableSpec
              { quantPrecision = precision
              , quantDestination = dest
              , quantTable = VS.fromListN 64 coeffs
              }
  
  data JpgHuffmanTableSpec = JpgHuffmanTableSpec
      { -- | 0 : DC, 1 : AC, stored on 4 bits
        huffmanTableClass       :: !DctComponent
        -- | Stored on 4 bits
      , huffmanTableDest        :: !Word8
  
      , huffSizes :: !(VU.Vector Word8)
      , huffCodes :: !(V.Vector (VU.Vector Word8))
      }
      deriving (Eq, Show, Generic)
  instance NFData JpgHuffmanTableSpec
  
  instance SizeCalculable JpgHuffmanTableSpec where
      calculateSize table = 1 + 16 + sum [fromIntegral e | e <- VU.toList $ huffSizes table]
  
  instance Binary JpgHuffmanTableSpec where
      put table = do
          let classVal = if huffmanTableClass table == DcComponent
                            then 0 else 1
          put4BitsOfEach classVal $ huffmanTableDest table
          mapM_ put . VU.toList $ huffSizes table
          forM_ [0 .. 15] $ \i ->
              when (huffSizes table ! i /= 0)
                   (let elements = VU.toList $ huffCodes table V.! i
                    in mapM_ put elements)
  
      get = do
          (huffClass, huffDest) <- get4BitOfEach
          sizes <- replicateM 16 getWord8
          codes <- forM sizes $ \s ->
              VU.replicateM (fromIntegral s) getWord8
          return JpgHuffmanTableSpec
              { huffmanTableClass =
                  if huffClass == 0 then DcComponent else AcComponent
              , huffmanTableDest = huffDest
              , huffSizes = VU.fromListN 16 sizes
              , huffCodes = V.fromListN 16 codes
              }
  
  instance Binary JpgImage where
      put (JpgImage { jpgFrame = frames }) =
          putWord8 0xFF >> putWord8 0xD8 >> mapM_ putFrame frames
              >> putWord8 0xFF >> putWord8 0xD9
  
      -- | Consider using `getJpgImage` instead for a non-semi-lazy implementation.
      get = do
          skipUntilFrames
          frames <- parseFramesSemiLazy
          -- let endOfImageMarker = 0xD9
          {-checkMarker commonMarkerFirstByte endOfImageMarker-}
          return JpgImage { jpgFrame = frames }
  
  -- | Like `get` from `instance Binary JpgImage`, but without the legacy
  -- semi-lazy implementation.
  getJpgImage :: Get JpgImage
  getJpgImage = do
      skipUntilFrames
      frames <- parseFrames
      return JpgImage { jpgFrame = frames }
  
  skipUntilFrames :: Get ()
  skipUntilFrames = do
      let startOfImageMarker = 0xD8
      checkMarker commonMarkerFirstByte startOfImageMarker
      eatUntilCode
  
  eatUntilCode :: Get ()
  eatUntilCode = do
      code <- getWord8
      unless (code == 0xFF) eatUntilCode
  
  takeCurrentFrame :: Get B.ByteString
  takeCurrentFrame = do
      size <- getWord16be
      getByteString (fromIntegral size - 2)
  
  putFrame :: JpgFrame -> Put
  putFrame (JpgAdobeAPP14 adobe) =
      put (JpgAppSegment 14) >> putWord16be 14 >> put adobe
  putFrame (JpgJFIF jfif) =
      put (JpgAppSegment 0) >> putWord16be (14+2) >> put jfif
  putFrame (JpgExif exif) = putExif exif
  putFrame (JpgAppFrame appCode str) =
      put (JpgAppSegment appCode) >> putWord16be (fromIntegral $ B.length str) >> put str
  putFrame (JpgExtension appCode str) =
      put (JpgExtensionSegment appCode) >> putWord16be (fromIntegral $ B.length str) >> put str
  putFrame (JpgQuantTable tables) =
      put JpgQuantizationTable >> put (TableList tables)
  putFrame (JpgHuffmanTable tables) =
      put JpgHuffmanTableMarker >> put (TableList $ map fst tables)
  putFrame (JpgIntervalRestart size) =
      put JpgRestartInterval >> put (RestartInterval size)
  putFrame (JpgScanBlob hdr blob) =
      put JpgStartOfScan >> put hdr >> putLazyByteString blob
  putFrame (JpgScans kind hdr) =
      put kind >> put hdr
  
  --------------------------------------------------
  ----            Serialization instances
  --------------------------------------------------
  commonMarkerFirstByte :: Word8
  commonMarkerFirstByte = 0xFF
  
  checkMarker :: Word8 -> Word8 -> Get ()
  checkMarker b1 b2 = do
      rb1 <- getWord8
      rb2 <- getWord8
      when (rb1 /= b1 || rb2 /= b2)
           (fail "Invalid marker used")
  
  -- | Simpler implementation of `parseECS` to allow an easier understanding
  -- of the logic, and to provide a comparison for correctness.
  parseECS_simple :: Get L.ByteString
  parseECS_simple = do
      -- There's no efficient way in `binary` to parse byte-by-byte while assembling a
      -- resulting ByteString (without using `.Internal` modules, which is what
      --  `parseECS` does), so instead first compute the length of the content
      -- byte-by-byte inside a `lookAhead` (not advancing the parser offset), and
      -- then efficiently take that long a ByteString (advancing the parser offset).
      --
      -- This is still slow compared to `parseECS` because parser functions
      -- (`getWord8`) are used repeatedly, instead of plain loops over ByteString contents.
      -- The slowdown is ~2x on GHC 8.10.7 on an Intel Core i7-7500U.
      n <- lookAhead getContentLength
      getLazyByteString n
    where
      getContentLength :: Get ByteOffset
      getContentLength = do
          bytesReadBeforeContent <- bytesRead
          let loop :: Word8 -> Get ByteOffset
              loop !v = do
                  vNext <- getWord8
                  let isReset = 0xD0 <= vNext && vNext <= 0xD7
                  let vIsSegmentMarker = v == 0xFF && vNext /= 0 && not isReset
                  if not vIsSegmentMarker
                      then loop vNext
                      else do
                          bytesReadAfterContentPlus2 <- bytesRead -- "plus 2" because we've also read the segment marker (0xFF and `vNext`)
                          let !contentLength = (bytesReadAfterContentPlus2 - 2) - bytesReadBeforeContent
                          return contentLength
  
          v_first <- getWord8
          loop v_first
  
  -- Replace by `Data.ByteString.dropEnd` once we require `bytestring >= 0.11.1.0`.
  bsDropEnd :: Int -> B.ByteString -> B.ByteString
  bsDropEnd n bs
      | n <= 0    = bs
      | n >= len  = B.empty
      | otherwise = B.take (len - 1) bs
    where
      len = B.length bs
  {-# INLINE bsDropEnd #-}
  
  -- | Parses a Scan's ECS (Entropy-Coded Segment, the largest part of compressed image data)
  -- from the `Get` stream.
  --
  -- When this function is called, the parser's offset should be
  -- immediately behind the SOS tag.
  --
  -- As described on e.g. https://www.ccoderun.ca/programming/2017-01-31_jpeg/,
  --
  -- > To find the next segment after the SOS, you must keep reading until you
  -- > find a 0xFF bytes which is not immediately followed by 0x00 (see "byte stuffing")
  -- > [or a reset marker's byte: 0xD0 through 0xD7].
  -- > Normally, this will be the EOI segment that comes at the end of the file.
  --
  -- where the 0xFF is the next segment's marker.
  -- See https://github.com/corkami/formats/blob/master/image/jpeg.md#entropy-coded-segment
  -- for more details.
  --
  -- This function returns the ECS, not including the next segment's
  -- marker on its trailing end.
  parseECS :: Get L.ByteString
  parseECS = do
      -- For a simpler but slower implementation of this function, see
      -- `parseECS_simple`.
  
      v_first <- getWord8
      -- TODO: Compare with what `scan` from `binary-parsers` does.
      --       Probably we cannot use it because it does not allow us to set the parser state
      --       to be _before_ the segment marker which would be convenient to not have to
      --       make a special case the function that calls this function.
      --       But `scan` works on pointers into the bytestring chunks. Why, for performance?
      --       I've asked on https://github.com/winterland1989/binary-parsers/issues/7
      --       If that is for performance, we may want to replicate the same thing here.
      --
      --       An orthogonal idea is to use `Data.ByteString.elemIndex` to fast-forward
      --       to the next 0xFF using `memchr`, but the `unsafe` call to `memchr` might
      --       have too much overhead, since 0xFF bytes appear statistically every 256 bytes.
      --       See https://stackoverflow.com/questions/14519905/how-much-does-it-cost-for-haskell-ffi-to-go-into-c-and-back
  
      -- `withInputChunks` allows us to work on chunks of ByteStrings,
      -- reducing the number of higher-overhead `Get` functions called.
      -- It also allows to easily assemble the ByteString to return,
      -- which may be cross-chunk.
      -- `withInputChunks` terminates when we return a
      --     Right (consumed :: ByteString, unconsumed :: ByteString)
      -- from `consumeChunk`, setting the `Get` parser's offset to just before `unconsumed`.
      -- Because the segment marker we seek may be the 2 bytes across chunk boundaries,
      -- we need to keep a reference to the previous chunk (initialised as `B.empty`),
      -- so that we can set `consumed` properly, because this function is supposed
      -- to not consume the start of the segment marker (see code dropping the last
      -- byte of the previous chunk below).
      GetInternal.withInputChunks
          (v_first, B.empty)
          consumeChunk
          (         L.fromChunks . (B.singleton v_first :)) -- `v_first` also belongs to the returned BS
          (return . L.fromChunks . (B.singleton v_first :)) -- `v_first` also belongs to the returned BS
    where
      consumeChunk :: GetInternal.Consume (Word8, B.ByteString) -- which is: (Word8, B.ByteString) -> B.ByteString -> Either (Word8, B.ByteString) (B.ByteString, B.ByteString)
      consumeChunk (!v_chunk_start, !prev_chunk) !chunk
          -- If `withInputChunks` hands us an empty chunk (which `binary` probably
          -- won't do, but since that's not documented, handle it anyway) then skip over it,
          -- so that we always remember the last `prev_chunk` that actually has data in it,
          -- since we `bsDropEnd 1 prev_chunk` in the `case` below.
          | B.null chunk = Left (v_chunk_start, prev_chunk)
          | otherwise = loop v_chunk_start 0
            where
              loop :: Word8 -> Int -> Either (Word8, B.ByteString) (B.ByteString, B.ByteString)
              loop !v !offset_in_chunk
                  | offset_in_chunk >= B.length chunk = Left (v, chunk)
                  | otherwise =
                      let !vNext = BU.unsafeIndex chunk offset_in_chunk -- bounds check is done above
                          !isReset = 0xD0 <= vNext && vNext <= 0xD7
                          !vIsSegmentMarker = v == 0xFF && vNext /= 0 && not isReset
                      in
                          if not vIsSegmentMarker
                              then loop vNext (offset_in_chunk+1)
                              else
                                  -- Set the parser state to _before_ the segment marker.
                                  -- The first case, where the segment marker's 2 bytes are exactly
                                  -- at the chunk boundary, requires us to allocate a new BS with
                                  -- `B.cons`; luckily this case should be rare.
                                  let (!consumed, !unconsumed) = case () of
                                       () | offset_in_chunk == 0 -> (bsDropEnd 1 prev_chunk, v `B.cons` chunk) -- segment marker starts at `v`, which is the last byte of the previous chunk
                                          | offset_in_chunk == 1 -> (B.empty, chunk) -- segment marker starts exactly at `chunk`
                                          | otherwise            -> B.splitAt (offset_in_chunk - 1) chunk -- segment marker starts at `v`, which is 1 before `vNext` (which is at `offset_in_chunk`)
                                  in Right $! (consumed, unconsumed)
  
  
  
  parseAdobe14 :: B.ByteString -> Maybe JpgFrame
  parseAdobe14 str = case runGetStrict get str of
      Left _err -> Nothing
      Right app14 -> Just $! JpgAdobeAPP14 app14
  
  -- | Parse JFIF or JFXX information. Right now only JFIF.
  parseJF__ :: B.ByteString -> Maybe JpgFrame
  parseJF__ str = case runGetStrict get str of
      Left _err -> Nothing
      Right jfif -> Just $! JpgJFIF jfif
  
  parseExif :: B.ByteString -> Maybe JpgFrame
  parseExif str
    | exifHeader `B.isPrefixOf` str =
        let
          tiff = B.drop (B.length exifHeader) str
        in
          case runGetStrict (getP tiff) tiff of
              Left _err -> Nothing
              Right (_hdr :: TiffHeader, []) -> Nothing
              Right (_hdr :: TiffHeader, ifds : _) -> Just $! JpgExif ifds
    | otherwise = Nothing
    where
      exifHeader = BC.pack "Exif\0\0"
  
  putExif :: [ImageFileDirectory] -> Put
  putExif ifds = putAll where
    hdr = TiffHeader
      { hdrEndianness = EndianBig
      , hdrOffset = 8
      }
  
    ifdList = case partition (isInIFD0 . ifdIdentifier) ifds of
      (ifd0, []) -> [ifd0]
      (ifd0, ifdExif) -> [ifd0 <> pure exifOffsetIfd, ifdExif]
  
    exifBlob = runPut $ do
      putByteString $ BC.pack "Exif\0\0"
      putP BC.empty (hdr, ifdList)
  
    putAll = do
      put (JpgAppSegment 1)
      putWord16be . fromIntegral $ L.length exifBlob + 2
      putLazyByteString exifBlob
  
  skipFrameMarker :: Get ()
  skipFrameMarker = do
      word <- getWord8
      when (word /= 0xFF) $ do
          readedData <- bytesRead
          fail $ "Invalid Frame marker (" ++ show word
                  ++ ", bytes read : " ++ show readedData ++ ")"
  
  -- | Parses a single frame.
  --
  -- Returns `Nothing` when we encounter a frame we want to skip.
  parseFrameOfKind :: JpgFrameKind -> Get (Maybe JpgFrame)
  parseFrameOfKind kind = do
      case kind of
          JpgEndOfImage -> return Nothing
          JpgAppSegment 0 -> parseJF__ <$> takeCurrentFrame
          JpgAppSegment 1 -> parseExif <$> takeCurrentFrame
          JpgAppSegment 14 -> parseAdobe14 <$> takeCurrentFrame
          JpgAppSegment c -> Just . JpgAppFrame c <$> takeCurrentFrame
          JpgExtensionSegment c -> Just . JpgExtension c <$> takeCurrentFrame
          JpgQuantizationTable ->
              (\(TableList quants) -> Just $! JpgQuantTable quants) <$> get
          JpgRestartInterval ->
              (\(RestartInterval i) -> Just $! JpgIntervalRestart i) <$> get
          JpgHuffmanTableMarker ->
              (\(TableList huffTables) -> Just $!
                      JpgHuffmanTable [(t, packHuffmanTree . buildPackedHuffmanTree $ huffCodes t) | t <- huffTables])
                      <$> get
          JpgStartOfScan -> do
              scanHeader <- get
              ecs <- parseECS
              return $! Just $! JpgScanBlob scanHeader ecs
          _ -> Just . JpgScans kind <$> get
  
  
  -- | Parse a list of `JpgFrame`s.
  --
  -- This function has various quirks; consider the below with great caution
  -- when using this function.
  --
  -- While @data JpgFrame = ... | JpgScanBlob !...` itself has strict fields,
  --
  -- This function is written in such a way that that it can construct
  -- the @[JpgFrame]@ "lazily" such that the expensive byte-by-byte traversal
  -- in `parseECS` to create a `JpgScanBlob` can be avoided if only
  -- list elements before that `JpgScanBlob` are evaluated.
  --
  -- That means the user can write code such as
  --
  -- > let mbFirstScan =
  -- >       case runGetOrFail (get @JPG.JpgImage) hugeImageByteString of -- (`get @JPG.JpgImage` uses `parseFramesSemiLazy`)
  -- >         Right (_restBs, _offset, res) ->
  -- >           find (\frame -> case frame of { JPG.JpgScans{} -> True; _ -> False }) (JPG.jpgFrame res)
  --
  -- with the guarantee that only the bytes before the ECS (large compressed image data)
  -- will be inspected, assuming that indeed there is at least 1 `JpgScan` in front
  -- of the `JpgScanBlob` that contains the ECS.
  --
  -- This guarantee can be useful to e.g. quickly read just the image
  -- dimensions (width, height) without traversing the large data.
  --
  -- Also note that this `Get` parser does not correctly maintain the parser byte offset
  -- (`Data.Binary.Get.bytesRead`), because as soon as a `JpgStartOfScan` is returned,
  -- it uses `Data.Binary.Get.getRemainingLazyBytes` to provide:
  --
  -- 1. the laziness described above, and
  -- 2. the ability to ignore any parser failure after the first successfully-parsed
  --    `JpgScanBlob` (it is debatable whether this behaviour is a desirable behaviour of this
  --    library, but it is historically so and existing exposed functions do not break
  --    this for backwards compatibility with existing uses of this library).
  --    This fact also means that even `parseNextFrameStrict` cannot maintain
  --    correct parser byte offsets.
  --
  -- Further note that if you are reading a huge JPEG image from disk strictly,
  -- this will already incur a full traversal (namely creation) of the `hugeImageByteString`.
  -- Thus, `parseNextFrameLazy` only provides any benefit if you:
  --
  -- - read the image from disk using lazy IO (not recommended!) such as via
  --   `Data.ByteString.Lazy.readFile`,
  -- - or do something similar, such as creating the `hugeImageByteString` via @mmap()@
  --
  -- This function is called "semi lazy" because only the first `JpgScanBlob` returned
  -- in the `[JpgFrame]` is returned lazily; frames of other types, or multiple
  -- `JpgScanBlob`s, are confusingly not dealt with lazily.
  --
  -- If as a caller you do not want to deal with any of these quirks,
  -- and use proper strict IO and/or via `Data.Binary.Get`'s incremental input interface:
  --
  -- - If you want the whole `[JpgFrame]`: use `parseFrames`.
  -- - If you want parsing to terminate early as in the example shown above,
  --   use in combination with just the right amount of `get :: Get JpgFrameKind`,
  --   `parseFrameOfKind`, and `skipFrameMarker`.
  parseFramesSemiLazy :: Get [JpgFrame]
  parseFramesSemiLazy = do
      kind <- get
      case kind of
          -- The end-of-image case needs to be here because `_ ->` default case below
          -- unconditionally uses `skipFrameMarker` which does not exist after `JpgEndOfImage`.
          JpgEndOfImage -> pure []
          JpgStartOfScan -> do
              scanHeader <- get
              remainingBytes <- getRemainingLazyBytes
              -- It is after the above `getRemainingLazyBytes` that the `Get` parser lazily succeeds,
              -- allowing consumers of `parseFramesSemiLazy` evaluate all `[JpgFrame]` list elements
              -- until (excluding) the cons-cell around the `JpgScanBlob ...` we construct below.
  
              return $ case runGet parseECS remainingBytes of
                  Left _ ->
                      -- Construct invalid `JpgScanBlob` even when the compressed JPEG
                      -- data is truncated or otherwise invalid, because that's what JuicyPixels's
                      -- `parseFramesSemiLazy` function did in the past, for backwards compat.
                      [JpgScanBlob scanHeader remainingBytes]
                  Right ecs ->
                      JpgScanBlob scanHeader ecs
                      :
                      -- TODO Why `drop 1` instead of `runGet (skipFrameMarker *> parseFramesSemiLazy) remainingBytes` that would check that the dropped 1 Byte is really a frame marker?
                      case runGet parseFramesSemiLazy (L.drop (L.length ecs + 1) remainingBytes) of
                          -- After we've encountered the first scan blob containing encoded image data,
                          -- we accept anything else after to fail parsing, ignoring that failure,
                          -- end emitting no further frames.
                          -- TODO: Explain why JuicyPixel chose to use this logic, insteaed of failing.
                          Left _ -> []
                          Right remainingFrames -> remainingFrames
          _ -> do
              mbFrame <- parseFrameOfKind kind
              skipFrameMarker
              remainingFrames <- parseFramesSemiLazy
              return $ maybeToList mbFrame ++ remainingFrames
  
  -- | Parse a list of `JpgFrame`s.
  parseFrames :: Get [JpgFrame]
  parseFrames = do
      kind <- get
      case kind of
          JpgEndOfImage -> pure []
          _ -> do
              mbFrame <- parseFrameOfKind kind
              skipFrameMarker
              remainingFrames <- parseFrames
              return $ maybeToList mbFrame ++ remainingFrames
  
  -- | Parse a list of `JpgFrameKind`s with their corresponding offsets and lengths
  -- (not counting the segment and frame markers into the lengths).
  --
  -- Useful for debugging.
  parseFrameKinds :: Get [(JpgFrameKind, Int64, Int64)]
  parseFrameKinds = do
      kindMarkerOffset :: Int64 <- bytesRead
      kind <- get
      case kind of
          JpgEndOfImage -> pure [(JpgEndOfImage, kindMarkerOffset, 0)]
          _ -> do
              parserOffsetBefore <- bytesRead
              _ <- parseFrameOfKind kind
              parserOffsetAfter <- bytesRead
              let !segmentLengthWithoutMarker = parserOffsetAfter - parserOffsetBefore
              skipFrameMarker
              remainingKinds <- parseFrameKinds
              return $ (kind, kindMarkerOffset, segmentLengthWithoutMarker):remainingKinds
  
  -- | Parses forward, returning the first scan header encountered.
  --
  -- Should be used after `skipUntilFrames`.
  --
  -- Fails parsing when an SOS segment marker (`JpgStartOfScan`, resulting
  -- in `JpgScanBlob`) is encountered before an SOF segment marker (that
  -- results in `JpgScans` carrying the `JpgFrameHeader`).
  parseToFirstFrameHeader :: Get (Maybe JpgFrameHeader)
  parseToFirstFrameHeader = do
      kind <- get
      case kind of
          JpgEndOfImage -> return Nothing
          JpgStartOfScan -> fail "parseToFirstFrameHeader: Encountered SOS frame marker before frame header that tells its dimensions"
          _ -> do
              mbFrame <- parseFrameOfKind kind
              case mbFrame of
                  Nothing -> continueSearching
                  Just frame -> case frame of
                      JpgScans _ frameHeader -> return $ Just $! frameHeader
                      _ -> continueSearching
    where
      continueSearching = do
          skipFrameMarker
          parseToFirstFrameHeader
  
  buildPackedHuffmanTree :: V.Vector (VU.Vector Word8) -> HuffmanTree
  buildPackedHuffmanTree = buildHuffmanTree . map VU.toList . V.toList
  
  secondStartOfFrameByteOfKind :: JpgFrameKind -> Word8
  secondStartOfFrameByteOfKind = aux
    where
      aux JpgBaselineDCTHuffman = 0xC0
      aux JpgExtendedSequentialDCTHuffman = 0xC1
      aux JpgProgressiveDCTHuffman = 0xC2
      aux JpgLosslessHuffman = 0xC3
      aux JpgDifferentialSequentialDCTHuffman = 0xC5
      aux JpgDifferentialProgressiveDCTHuffman = 0xC6
      aux JpgDifferentialLosslessHuffman = 0xC7
      aux JpgExtendedSequentialArithmetic = 0xC9
      aux JpgProgressiveDCTArithmetic = 0xCA
      aux JpgLosslessArithmetic = 0xCB
      aux JpgHuffmanTableMarker = 0xC4
      aux JpgDifferentialSequentialDCTArithmetic = 0xCD
      aux JpgDifferentialProgressiveDCTArithmetic = 0xCE
      aux JpgDifferentialLosslessArithmetic = 0xCF
      aux JpgEndOfImage = 0xD9
      aux JpgQuantizationTable = 0xDB
      aux JpgStartOfScan = 0xDA
      aux JpgRestartInterval = 0xDD
      aux (JpgRestartIntervalEnd v) = v
      aux (JpgAppSegment a) = (a + 0xE0)
      aux (JpgExtensionSegment a) = a
  
  data JpgImageKind = BaseLineDCT | ProgressiveDCT
  
  instance Binary JpgFrameKind where
      put v = putWord8 0xFF >> put (secondStartOfFrameByteOfKind v)
      get = do
          -- no lookahead :(
          {-word <- getWord8-}
          word2 <- getWord8
          case word2 of
              0xC0 -> return JpgBaselineDCTHuffman
              0xC1 -> return JpgExtendedSequentialDCTHuffman
              0xC2 -> return JpgProgressiveDCTHuffman
              0xC3 -> return JpgLosslessHuffman
              0xC4 -> return JpgHuffmanTableMarker
              0xC5 -> return JpgDifferentialSequentialDCTHuffman
              0xC6 -> return JpgDifferentialProgressiveDCTHuffman
              0xC7 -> return JpgDifferentialLosslessHuffman
              0xC9 -> return JpgExtendedSequentialArithmetic
              0xCA -> return JpgProgressiveDCTArithmetic
              0xCB -> return JpgLosslessArithmetic
              0xCD -> return JpgDifferentialSequentialDCTArithmetic
              0xCE -> return JpgDifferentialProgressiveDCTArithmetic
              0xCF -> return JpgDifferentialLosslessArithmetic
              0xD9 -> return JpgEndOfImage
              0xDA -> return JpgStartOfScan
              0xDB -> return JpgQuantizationTable
              0xDD -> return JpgRestartInterval
              a | a >= 0xF0 -> return $! JpgExtensionSegment a
                | a >= 0xE0 -> return $! JpgAppSegment (a - 0xE0)
                | a >= 0xD0 && a <= 0xD7 -> return $! JpgRestartIntervalEnd a
                | otherwise -> fail ("Invalid frame marker (" ++ show a ++ ")")
  
  put4BitsOfEach :: Word8 -> Word8 -> Put
  put4BitsOfEach a b = put $ (a `unsafeShiftL` 4) .|. b
  
  get4BitOfEach :: Get (Word8, Word8)
  get4BitOfEach = do
      val <- get
      return ((val `unsafeShiftR` 4) .&. 0xF, val .&. 0xF)
  
  newtype RestartInterval = RestartInterval Word16
  
  instance Binary RestartInterval where
      put (RestartInterval i) = putWord16be 4 >> putWord16be i
      get = do
          size <- getWord16be
          when (size /= 4) (fail "Invalid jpeg restart interval size")
          RestartInterval <$> getWord16be
  
  instance Binary JpgComponent where
      get = do
          ident <- getWord8
          (horiz, vert) <- get4BitOfEach
          quantTableIndex <- getWord8
          return JpgComponent
              { componentIdentifier = ident
              , horizontalSamplingFactor = horiz
              , verticalSamplingFactor = vert
              , quantizationTableDest = quantTableIndex
              }
      put v = do
          put $ componentIdentifier v
          put4BitsOfEach (horizontalSamplingFactor v) $ verticalSamplingFactor v
          put $ quantizationTableDest v
  
  instance Binary JpgFrameHeader where
      get = do
          beginOffset <- fromIntegral <$> bytesRead
          frmHLength <- getWord16be
          samplePrec <- getWord8
          h <- getWord16be
          w <- getWord16be
          compCount <- getWord8
          components <- replicateM (fromIntegral compCount) get
          endOffset <- fromIntegral <$> bytesRead
          when (beginOffset - endOffset < fromIntegral frmHLength)
               (skip $ fromIntegral frmHLength - (endOffset - beginOffset))
          return JpgFrameHeader
              { jpgFrameHeaderLength = frmHLength
              , jpgSamplePrecision = samplePrec
              , jpgHeight = h
              , jpgWidth = w
              , jpgImageComponentCount = compCount
              , jpgComponents = components
              }
  
      put v = do
          putWord16be $ jpgFrameHeaderLength v
          putWord8    $ jpgSamplePrecision v
          putWord16be $ jpgHeight v
          putWord16be $ jpgWidth v
          putWord8    $ jpgImageComponentCount v
          mapM_ put   $ jpgComponents v
  
  instance Binary JpgScanSpecification where
      put v = do
          put $ componentSelector v
          put4BitsOfEach (dcEntropyCodingTable v) $ acEntropyCodingTable v
  
      get = do
          compSel <- get
          (dc, ac) <- get4BitOfEach
          return JpgScanSpecification {
              componentSelector = compSel
            , dcEntropyCodingTable = dc
            , acEntropyCodingTable = ac
            }
  
  instance Binary JpgScanHeader where
      get = do
          thisScanLength <- getWord16be
          compCount <- getWord8
          comp <- replicateM (fromIntegral compCount) get
          specBeg <- get
          specEnd <- get
          (approxHigh, approxLow) <- get4BitOfEach
  
          return JpgScanHeader {
              scanLength = thisScanLength,
              scanComponentCount = compCount,
              scans = comp,
              spectralSelection = (specBeg, specEnd),
              successiveApproxHigh = approxHigh,
              successiveApproxLow = approxLow
          }
  
      put v = do
          putWord16be $ scanLength v
          putWord8 $ scanComponentCount v
          mapM_ put $ scans v
          putWord8 . fst $ spectralSelection v
          putWord8 . snd $ spectralSelection v
          put4BitsOfEach (successiveApproxHigh v) $ successiveApproxLow v
  
  {-# INLINE createEmptyMutableMacroBlock #-}
  -- | Create a new macroblock with the good array size
  createEmptyMutableMacroBlock :: (Storable a, Num a) => ST s (MutableMacroBlock s a)
  createEmptyMutableMacroBlock = M.replicate 64 0
  
  printMacroBlock :: (Storable a, PrintfArg a)
                  => MutableMacroBlock s a -> ST s String
  printMacroBlock block = pLn 0
      where pLn 64 = return "===============================\n"
            pLn i = do
                v <- block `M.unsafeRead` i
                vn <- pLn (i+1)
                return $ printf (if i `mod` 8 == 0 then "\n%5d " else "%5d ") v ++ vn
  
  printPureMacroBlock :: (Storable a, PrintfArg a) => MacroBlock a -> String
  printPureMacroBlock block = pLn 0
      where pLn 64 = "===============================\n"
            pLn i = str ++ pLn (i + 1)
              where str | i `mod` 8 == 0 = printf "\n%5d " v
                        | otherwise = printf "%5d" v
                    v = block VS.! i
  
  
  {-# INLINE dctBlockSize #-}
  dctBlockSize :: Num a => a
  dctBlockSize = 8
  