content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE TypeSynonymInstances #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  -- | Module implementing a basic png export, no filtering is applyed, but
  -- export at least valid images.
  module Codec.Picture.Png.Internal.Export( PngSavable( .. )
                                 , PngPaletteSaveable( .. )
                                 , writePng
                                 , encodeDynamicPng
                                 , writeDynamicPng
                                 ) where
  #if !MIN_VERSION_base(4,8,0)
  import Data.Monoid( mempty )
  #endif
  
  import Control.Monad( forM_ )
  import Control.Monad.ST( ST, runST )
  import Data.Bits( unsafeShiftR, (.&.) )
  import Data.Binary( encode )
  #if !MIN_VERSION_base(4,11,0)
  import Data.Monoid( (<>) )
  #endif
  import Data.Word(Word8, Word16)
  import qualified Codec.Compression.Zlib as Z
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy as Lb
  
  import qualified Data.Vector.Storable as VS
  import qualified Data.Vector.Storable.Mutable as M
  
  import Codec.Picture.Types
  import Codec.Picture.Png.Internal.Type
  import Codec.Picture.Png.Internal.Metadata
  import Codec.Picture.Metadata( Metadatas )
  import Codec.Picture.VectorByteConversion( blitVector, toByteString )
  
  -- | Encode a paletted image into a png if possible.
  class PngPaletteSaveable a where
    -- | Encode a paletted image as a color indexed 8-bit PNG.
    -- the palette must have between 1 and 256 values in it.
    -- Accepts `PixelRGB8` and `PixelRGBA8` as palette pixel type
    encodePalettedPng :: Image a -> Image Pixel8 -> Either String Lb.ByteString
    encodePalettedPng = encodePalettedPngWithMetadata mempty
  
    -- | Equivalent to 'encodePalettedPng' but allow writing of metadatas.
    -- See `encodePngWithMetadata` for the details of encoded metadatas
    -- Accepts `PixelRGB8` and `PixelRGBA8` as palette pixel type
    encodePalettedPngWithMetadata :: Metadatas -> Image a -> Image Pixel8 -> Either String Lb.ByteString
  
  instance PngPaletteSaveable PixelRGB8 where
    encodePalettedPngWithMetadata metas pal img
        | w <= 0 || w > 256 || h /= 1 = Left "Invalid palette"
        | VS.any isTooBig $ imageData img =
            Left "Image contains indexes absent from the palette"
        | otherwise = Right $ genericEncodePng (Just pal) Nothing PngIndexedColor metas img
          where w = imageWidth pal
                h = imageHeight pal
                isTooBig v = fromIntegral v >= w
  
  instance PngPaletteSaveable PixelRGBA8 where
    encodePalettedPngWithMetadata metas pal img
        | w <= 0 || w > 256 || h /= 1 = Left "Invalid palette"
        | VS.any isTooBig $ imageData img =
            Left "Image contains indexes absent from the palette"
        | otherwise = Right $ genericEncodePng (Just opaquePalette) (Just alphaPal) PngIndexedColor metas img
      where
        w = imageWidth pal
        h = imageHeight pal
        opaquePalette = dropAlphaLayer pal
        alphaPal = imageData $ extractComponent PlaneAlpha pal
        isTooBig v = fromIntegral v >= w
  
  -- | Encode an image into a png if possible.
  class PngSavable a where
      -- | Transform an image into a png encoded bytestring, ready
      -- to be written as a file.
      encodePng :: Image a -> Lb.ByteString
      encodePng = encodePngWithMetadata mempty
  
      -- | Encode a png using some metadatas. The following metadata keys will
      -- be stored in a `tEXt` field :
      --
      --  * 'Codec.Picture.Metadata.Title'
      --  * 'Codec.Picture.Metadata.Description'
      --  * 'Codec.Picture.Metadata.Author'
      --  * 'Codec.Picture.Metadata.Copyright'
      --  * 'Codec.Picture.Metadata.Software'
      --  * 'Codec.Picture.Metadata.Comment'
      --  * 'Codec.Picture.Metadata.Disclaimer'
      --  * 'Codec.Picture.Metadata.Source'
      --  * 'Codec.Picture.Metadata.Warning'
      --  * 'Codec.Picture.Metadata.Unknown' using the key present in the constructor.
      -- 
      -- the followings metadata will be stored in the `gAMA` chunk.
      --
      --  * 'Codec.Picture.Metadata.Gamma'
      --
      -- The followings metadata will be stored in a `pHYs` chunk
      --
      --  * 'Codec.Picture.Metadata.DpiX'
      --  * 'Codec.Picture.Metadata.DpiY' 
      encodePngWithMetadata :: Metadatas -> Image a -> Lb.ByteString
  
  preparePngHeader :: Image a -> PngImageType -> Word8 -> PngIHdr
  preparePngHeader (Image { imageWidth = w, imageHeight = h }) imgType depth = PngIHdr
    { width             = fromIntegral w
    , height            = fromIntegral h
    , bitDepth          = depth
    , colourType        = imgType
    , compressionMethod = 0
    , filterMethod      = 0
    , interlaceMethod   = PngNoInterlace
    }
  
  -- | Helper function to directly write an image as a png on disk.
  writePng :: (PngSavable pixel) => FilePath -> Image pixel -> IO ()
  writePng path img = Lb.writeFile path $ encodePng img
  
  endChunk :: PngRawChunk
  endChunk = mkRawChunk iENDSignature mempty
  
  prepareIDatChunk :: Lb.ByteString -> PngRawChunk
  prepareIDatChunk = mkRawChunk iDATSignature
  
  genericEncode16BitsPng :: forall px. (Pixel px, PixelBaseComponent px ~ Word16)
                         => PngImageType -> Metadatas -> Image px -> Lb.ByteString
  genericEncode16BitsPng imgKind metas
                   image@(Image { imageWidth = w, imageHeight = h, imageData = arr }) =
    encode PngRawImage { header = hdr
                       , chunks = encodeMetadatas metas 
                                <> [ prepareIDatChunk imgEncodedData
                                   , endChunk
                                   ]
                       }
      where hdr = preparePngHeader image imgKind 16
            zero = B.singleton 0
            compCount = componentCount (undefined :: px)
  
            lineSize = compCount * w
            blitToByteString vec = blitVector vec 0 (lineSize * 2)
            encodeLine line = blitToByteString $ runST $ do
                finalVec <- M.new $ lineSize * 2 :: ST s (M.STVector s Word8)
                let baseIndex = line * lineSize
                forM_ [0 ..  lineSize - 1] $ \ix -> do
                    let v = arr `VS.unsafeIndex` (baseIndex + ix)
                        high = fromIntegral $ (v `unsafeShiftR` 8) .&. 0xFF
                        low = fromIntegral $ v .&. 0xFF
  
                    (finalVec `M.unsafeWrite` (ix * 2 + 0)) high
                    (finalVec `M.unsafeWrite` (ix * 2 + 1)) low
  
                VS.unsafeFreeze finalVec
  
            imgEncodedData = Z.compress . Lb.fromChunks
                          $ concat [[zero, encodeLine line] | line <- [0 .. h - 1]]
  
  preparePalette :: Palette -> PngRawChunk
  preparePalette pal = PngRawChunk
    { chunkLength = fromIntegral $ imageWidth pal * 3
    , chunkType   = pLTESignature
    , chunkCRC    = pngComputeCrc [pLTESignature, binaryData]
    , chunkData   = binaryData
    }
    where binaryData = Lb.fromChunks [toByteString $ imageData pal]
  
  preparePaletteAlpha :: VS.Vector Pixel8 -> PngRawChunk
  preparePaletteAlpha alphaPal = PngRawChunk
    { chunkLength = fromIntegral $ VS.length alphaPal
    , chunkType   = tRNSSignature
    , chunkCRC    = pngComputeCrc [tRNSSignature, binaryData]
    , chunkData   = binaryData
    }
    where binaryData = Lb.fromChunks [toByteString alphaPal]
  
  type PaletteAlpha = VS.Vector Pixel8
  
  genericEncodePng :: forall px. (Pixel px, PixelBaseComponent px ~ Word8)
                   => Maybe Palette
                   -> Maybe PaletteAlpha
                   -> PngImageType -> Metadatas -> Image px
                   -> Lb.ByteString
  genericEncodePng palette palAlpha imgKind metas
                   image@(Image { imageWidth = w, imageHeight = h, imageData = arr }) =
    encode PngRawImage { header = hdr
                       , chunks = encodeMetadatas metas
                                <> paletteChunk
                                <> transpChunk
                                <> [ prepareIDatChunk imgEncodedData
                                   , endChunk
                                   ]}
    where
      hdr = preparePngHeader image imgKind 8
      zero = B.singleton 0
      compCount = componentCount (undefined :: px)
  
      paletteChunk = case palette of
        Nothing -> []
        Just p -> [preparePalette p]
  
      transpChunk = case palAlpha of
        Nothing -> []
        Just p -> [preparePaletteAlpha p]
  
      lineSize = compCount * w
      encodeLine line = blitVector arr (line * lineSize) lineSize
      imgEncodedData = Z.compress
          . Lb.fromChunks
          $ concat [[zero, encodeLine line] | line <- [0 .. h - 1]]
  
  instance PngSavable PixelRGBA8 where
    encodePngWithMetadata = genericEncodePng Nothing Nothing PngTrueColourWithAlpha
  
  instance PngSavable PixelRGB8 where
    encodePngWithMetadata = genericEncodePng Nothing Nothing PngTrueColour
  
  instance PngSavable Pixel8 where
    encodePngWithMetadata = genericEncodePng Nothing Nothing PngGreyscale
  
  instance PngSavable PixelYA8 where
    encodePngWithMetadata = genericEncodePng Nothing Nothing PngGreyscaleWithAlpha
  
  instance PngSavable PixelYA16 where
    encodePngWithMetadata = genericEncode16BitsPng PngGreyscaleWithAlpha
  
  instance PngSavable Pixel16 where
    encodePngWithMetadata = genericEncode16BitsPng PngGreyscale
  
  instance PngSavable PixelRGB16 where
    encodePngWithMetadata = genericEncode16BitsPng PngTrueColour
  
  instance PngSavable PixelRGBA16 where
    encodePngWithMetadata = genericEncode16BitsPng PngTrueColourWithAlpha
  
  -- | Write a dynamic image in a .png image file if possible.
  -- The same restriction as encodeDynamicPng apply.
  writeDynamicPng :: FilePath -> DynamicImage -> IO (Either String Bool)
  writeDynamicPng path img = case encodeDynamicPng img of
          Left err -> return $ Left err
          Right b  -> Lb.writeFile path b >> return (Right True)
  
  -- | Encode a dynamic image in PNG if possible, supported images are:
  --
  --   * 'ImageY8'
  --
  --   * 'ImageY16'
  --
  --   * 'ImageYA8'
  --
  --   * 'ImageYA16'
  --
  --   * 'ImageRGB8'
  --
  --   * 'ImageRGB16'
  --
  --   * 'ImageRGBA8'
  --
  --   * 'ImageRGBA16'
  --
  encodeDynamicPng :: DynamicImage -> Either String Lb.ByteString
  encodeDynamicPng (ImageRGB8 img) = Right $ encodePng img
  encodeDynamicPng (ImageRGBA8 img) = Right $ encodePng img
  encodeDynamicPng (ImageY8 img) = Right $ encodePng img
  encodeDynamicPng (ImageY16 img) = Right $ encodePng img
  encodeDynamicPng (ImageYA8 img) = Right $ encodePng img
  encodeDynamicPng (ImageYA16 img) = Right $ encodePng img
  encodeDynamicPng (ImageRGB16 img) = Right $ encodePng img
  encodeDynamicPng (ImageRGBA16 img) = Right $ encodePng img
  encodeDynamicPng _ = Left "Unsupported image format for PNG export"
  