{- |
Module      :  Codec.PNM
Copyright   :  (c) Claude Heiland-Allen 2012
License     :  BSD3

Maintainer  :  claude@mathr.co.uk
Stability   :  provisional
Portability :  portable

PNM image format header parsing and pretty printing.

PNM is PBM + PGM + PPM, a family of lowest common denominator image file formats.

References: <http://netpbm.sourceforge.net/doc/pnm.html>.

-}
module Codec.PNM
  ( PNM(..), pnmRasterBytes, pnmPretty, Parse(..), pnmParse
  , pnms, unpnms, onPNMs, onPNMs'
  ) where

import Prelude hiding ((++), (**), concat, splitAt)
import qualified Prelude as P

import Data.ByteString.Lazy (ByteString, uncons, empty, concat, splitAt)
import Data.Int (Int64)
import Codec.PNM.Parse

-- | PNM image headers.
data PNM
  = PBM{ pnmPlain :: Bool, pnmWidth, pnmHeight            :: Integer }
  | PGM{ pnmPlain :: Bool, pnmWidth, pnmHeight, pnmMaxVal :: Integer }
  | PPM{ pnmPlain :: Bool, pnmWidth, pnmHeight, pnmMaxVal :: Integer }
  deriving (Eq, Ord, Show)

-- | Compute the raster size in bytes for binary PNM images.
pnmRasterBytes :: PNM -> Maybe Integer
pnmRasterBytes pnm
  | pnmPlain pnm = Nothing
  | otherwise = Just $ case pnm of
      PBM _ w h -> ((w + 7) `div` 8) * h
      PGM _ w h m
        | m < 256 -> w * h
        | otherwise -> 2 * w * h
      PPM _ w h m
        | m < 256 -> 3 * w * h
        | otherwise -> 6 * w * h

-- | Pretty-print a PNM image header without any comments.
pnmPretty :: PNM -> ByteString
pnmPretty (PBM True  w h  ) = str $ "P1\n" ** show w ** " " ** show h ** "\n"
pnmPretty (PBM False w h  ) = str $ "P4\n" ** show w ** " " ** show h ** "\n"
pnmPretty (PGM True  w h m) = str $ "P2\n" ** show w ** " " ** show h ** "\n" ** show m ** "\n"
pnmPretty (PGM False w h m) = str $ "P5\n" ** show w ** " " ** show h ** "\n" ** show m ** "\n"
pnmPretty (PPM True  w h m) = str $ "P3\n" ** show w ** " " ** show h ** "\n" ** show m ** "\n"
pnmPretty (PPM False w h m) = str $ "P6\n" ** show w ** " " ** show h ** "\n" ** show m ** "\n"

(**) :: String -> String -> String
(**) = (P.++)

-- | Parse a PNM image header.
pnmParse :: ByteString -> Parse PNM
pnmParse s = case uncons s of
  Nothing -> Empty
  Just (p, ps)
    | p == 80{-P-} -> case uncons ps of
        Just (d, _)
          | d == d1 -> p1 s
          | d == d2 -> p2 s
          | d == d3 -> p3 s
          | d == d4 -> p4 s
          | d == d5 -> p5 s
          | d == d6 -> p6 s
        _ -> Wrong s
  _ -> Wrong s

p1, p2, p3, p4, p5, p6 :: ByteString -> Parse PNM

(p1, p4) = (pbm "P1" (PBM True), pbm "P4" (PBM False))
  where
    pbm magic ctor s0 = case string (str magic) s0 of
      Empty -> Empty
      Parse () ps s1 -> case number s1 of
        Parse w ws s2 -> case number s2 of
          Parse h hs s3 -> case oneSpace s3 of
            Parse () ss s4 -> Parse (ctor w h) (ps ++ ws ++ hs ++ ss) s4
            _ -> Wrong s0
          _ -> Wrong s0
        _ -> Wrong s0
      _ -> Wrong s0

(p2, p3, p5, p6) = (ppm "P2" (PGM True), ppm "P3" (PPM True), ppm "P5" (PGM False), ppm "P6" (PPM False))
  where
    ppm magic ctor s0 = case string (str magic) s0 of
      Empty -> Empty
      Parse () ps s1 -> case number s1 of
        Parse w ws s2 -> case number s2 of
          Parse h hs s3 -> case number s3 of
            Parse m ms s4 | m < 65536 -> case oneSpace s4 of
              Parse () ss s5 -> Parse (ctor w h m) (ps ++ ws ++ hs ++ ms ++ ss) s5
              _ -> Wrong s0
            _ -> Wrong s0
          _ -> Wrong s0
        _ -> Wrong s0
      _ -> Wrong s0

-- |  Parse a sequence of binary PNM images.
--
--    Malformed input (including huge raster sizes or plain images)
--    is treated as end-of-image-stream.
--
pnms :: ByteString -> [(PNM, ByteString)]
pnms s = case pnmParse s of
  Parse pnm _ rest
    | Nothing < bytes && bytes <= Just (toInteger (maxBound :: Int64)) ->
        (pnm, raster) : pnms next
    where
      ~(raster, next) = splitAt (fromInteger bytes') rest
      bytes@(~(Just bytes')) = pnmRasterBytes pnm
  _ -> []  -- FIXME what to do about huge rasters or malformed input?

-- |  Pretty-print a sequence of binary PNM images.
--
--    The precondition that the raster is of the correct length is not
--    checked, so malformed output is possible.
--
unpnms :: [(PNM, ByteString)] -> ByteString
unpnms = concat . map (\ ~(pnm, raster) -> pnmPretty pnm ++ raster)

-- |  Process a sequence of binary PNM images.
--
--    Malformed input (including huge raster sizes or plain images)
--    is treated as end-of-image-stream.
--
--    The precondition that the raster is of the correct length is not
--    checked, so malformed output is possible.
--
--    Header comments are not preserved.
--
onPNMs :: (PNM -> ByteString -> (PNM, ByteString)) -> ByteString -> ByteString
onPNMs f = unpnms . map (\ ~(pnm, raster) -> f pnm raster) . pnms

-- |  Process a sequence of binary PNM images.
--
--    Malformed input (including huge raster sizes or plain images)
--    is treated as end-of-image-stream.
--
--    The precondition that the raster is of the correct length is not
--    checked, so malformed output is possible.
--
--    Header comments are preserved.  Assuming well-formed input:
--
--    > onPNMs' (\_ r -> r) = id
--
onPNMs' :: (PNM -> ByteString -> ByteString) -> ByteString -> ByteString
onPNMs' f s = case pnmParse s of
  Parse pnm header rest
    | Nothing < bytes && bytes <= Just (toInteger (maxBound :: Int64)) ->
        header ++ raster' ++ onPNMs' f next
    where
      raster' = f pnm raster
      ~(raster, next) = splitAt (fromInteger bytes') rest
      bytes@(~(Just bytes')) = pnmRasterBytes pnm
  _ -> empty  -- FIXME what to do about huge rasters or malformed input?