{-# LANGUAGE FlexibleInstances #-}
module PNGparser(readPNG,getPNG,parsePNG,parsePNG',sizeOfPNG,sizeOfPNG') where
import Control.Monad(guard)
import Data.Traversable
import Data.Char(isLower)
import qualified Codec.Compression.Zlib as Zlib
import qualified Data.ByteString.Lazy as BS
import qualified Data.ByteString.Lazy.Char8 as CS
import ByteStream(FromBytes(..))
import Utils2(pieces)
import PNG
import PNGchunks
import PNGfilter(undo_filter)
import PNGinterlace(adam7passesPre)
import Byte
import ParsOps1

default (Int)

readPNG path = parsePNG =<< BS.readFile path
getPNG = parsePNG =<< BS.getContents

parsePNG' s = parsePNG (CS.pack s)

instance FromBytes idata => FromBytes (PNG idata) where
  fromBytes bs = traverse fromBytes =<< parsePNG bs

parsePNG bs = do PNGstream chunks <- parseM pngStreamP bs
                 parseM pngP chunks

pngP = png <$> oneChunkP "IHDR" ihdrP
           <*> many ancillaryP
           <*> optional noPLTE (oneChunkP "PLTE" plteP)
           <*> many ancillaryP
           <*> (zlib_uncompress <$> some (oneChunkP "IDAT" theRest))
           <*  many ancillaryP
           <*  oneChunkP "IEND" (return [])
  where
    png ihdr as1 plte as2 idata = PNG ihdr plte bkgd trns (BS.concat rawlines)
      where
        rawlines = concatMap (map BS.pack.undo_filter b.map BS.unpack) scanlines
        b = (bpp ihdr+7) `quot` 8 -- bytes per pixel, rounded up

        scanlines =
          case interlaceMethod ihdr of
            NoInterlace -> [pieces (pitch ihdr+1) idata] -- all in one pass
            Adam7 -> adam7passesPre ihdr idata

        as = as1++as2
        ct = colorType ihdr

        bkgd = case [parse (bkgdP ct) bs | Chunk _ "bKGD" bs _<-as] of
                 Just bkgd:_ -> bkgd
                 _ -> noBKGD

        trns = case [parse (trnsP ct) bs | Chunk _ "tRNS" bs _<-as] of
                 Just bkgd:_ -> bkgd
                 _ -> noTRNS

-- | Parse just enough of a PNG file to obtain the size
-- (width and height in pixels) of the PNG image.
sizeOfPNG' s = sizeOfPNG (CS.pack s)
-- | Parse just enough of a PNG file to obtain the size
-- (width and height in pixels) of the PNG image.
sizeOfPNG bs =
  do Chunk _ "IHDR" bs _ <- parseM (litsP png_signature *> chunkP <* theRest) bs
     ihdr <- parseM ihdrP bs
     return (width ihdr,height ihdr)

pngStreamP = PNGstream <$ litsP png_signature <*> some chunkP

ancillaryP =
  do chunk@(Chunk len (c:_) bs crc) <- token
     guard (isLower c)
     return chunk

oneChunkP t p =
  do Chunk len t' bs crc <- token
     guard (t'==t)
     parseM p bs

parseM p = either fail return . parseE p

chunkP = do len <- be32P
            Chunk (toEnum len) <$> charsP 4 <*> bytesP len <*> w32P

ihdrP = IHDR <$> be32P <*> be32P <*> depthP <*> colTyP <*> ebP <*> ebP <*> ebP
  where
    depthP = do d <- w8P
                guard (d `elem` [1,2,4,8,16])
                return d

plteP = PLTE <$> many rgbP
rgbP = RGB <$> w8P <*> w8P <*> w8P
rgb16P = RGB <$> w16P <*> w16P <*> w16P

colTyP =
  do b <- byteP
     case b of
       0 -> return GreyScale
       2 -> return Truecolor
       3 -> return IndexedColor
       4 -> return GreyScaleWithAlpha
       6 -> return TruecolorWithAlpha
       _ -> failP "ColorType"

bkgdP ct =
  case ct of
    IndexedColor       -> IxBg        <$> w8P
    GreyScale          -> GreyBg      <$> w16P
    GreyScaleWithAlpha -> GreyBg      <$> w16P
    Truecolor          -> TruecolorBg <$> rgb16P
    TruecolorWithAlpha -> TruecolorBg <$> rgb16P

trnsP ct =
  case ct of
    IndexedColor      -> IxTrans        <$> many w8P
    GreyScale         -> GreyTrans      <$> w16P
    Truecolor         -> TruecolorTrans <$> rgb16P
    _                 -> failP "Unexpected tRNS chunk"

--------------------------------------------------------------------------------

-- | Parse 16-bit and 32-bit numeric values stored in big-endian order
be16P,be32P :: Num a => Parser BS.ByteString a
be32P = fromIntegral <$> w32P
be16P = fromIntegral <$> w16P

w16P = join8  <$> w8P   <*> w8P
w32P = join16 <$> be16P <*> be16P

-- | Parse a value of an enumeration type that is represented as one byte
ebP :: Enum a => Parser BS.ByteString a
ebP = toEnum <$> byteP

--------------------------------------------------------------------------------

byteP :: Num a => Parser BS.ByteString a
byteP = fromIntegral <$> w8P

w8P = get BS.uncons

charsP n = CS.unpack <$> bytesP n

--bytesP n = CS.pack <$> tokens n
bytesP n = get (splitAt' n)
splitAt' n bs = case BS.splitAt n' bs of
                  p@(bs1,_) | BS.length bs1 == n' -> Just p
                            | otherwise -> Nothing
  where
    n' = fromIntegral n

litsP pre = get (fmap ((,)()) . BS.stripPrefix pre)

--------------------------------------------------------------------------------

zlib_uncompress = Zlib.decompress. BS.concat