Hufflepdf/src/PDF/Object/Navigation.hs

{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
module PDF.Object.Navigation (
      Clear(..)
    , PDFContent
    , Raw(..)
    , (./)
    , (//)
    , (>./)
    , (>//)
    , castObject
    , getDictionary
    , getKey
    , objectById
    , catalog
  ) where

import Codec.Compression.Zlib (compress, decompress)
import Control.Monad.Reader (MonadReader(..))
import Control.Monad.Fail (MonadFail(..))
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as Lazy (ByteString, fromStrict, toStrict)
import Data.Id (Id, at)
import qualified Data.Map as Map (lookup)
import PDF.Box (Box(..))
import PDF.Layer (Layer(..))
import PDF.Object (
      Dictionary, DirectObject(..), IndirectObjCoordinates(..)
    , Name(..), Object(..), Structure(..)
  )
import Prelude hiding (fail)
import Text.Printf (printf)

type PDFContent m = (MonadReader Layer m, MonadFail m)
type Component = String

getDictionary :: PDFContent m => Object -> m Dictionary
getDictionary (Direct (Dictionary aDict)) = return aDict
getDictionary (Direct (Reference ref)) =
  objectById (objectId ref) >>= getDictionary
getDictionary (Stream {header}) = return header
getDictionary obj = expected "dictionary : " obj

expected :: (MonadFail m, Show a) => String -> a -> m b
expected name = fail . printf "Not a %s: %s" name . show

getKey :: PDFContent m => String -> Object -> m DirectObject
getKey key object = getDictionary object >>= catchMaybe . Map.lookup (Name key)
  where
    errorMessage =
      printf "Key %s not found in object %s" key (show object)
    catchMaybe = maybe (fail errorMessage) return

objectById :: PDFContent m => (Id Object) -> m Object
objectById objectId = do
  layer <- ask
  return (objects layer `at` objectId)

(./) :: PDFContent m => m Object -> Component -> m Object
(./) object key = (object >>= getKey key >>= castObject)

castObject :: PDFContent m => DirectObject -> m Object
castObject (Reference (IndirectObjCoordinates {objectId})) =
  objectById objectId
castObject directObject = return $ Direct directObject

(//) :: PDFContent m => m Object -> [Component] -> m Object
(//) object [] = object
(//) object (key:keys) = object ./ key // keys

(>./) :: PDFContent m => Object -> Component -> m Object
(>./) object = (return object ./)

(>//) :: PDFContent m => Object -> [Component] -> m Object
(>//) object = (return object //)

catalog :: PDFContent m => m Object
catalog = Direct . Dictionary . trailer . docStructure <$> ask

data Clear = Clear
data Raw = Raw

onLazy :: (Lazy.ByteString -> Lazy.ByteString) -> ByteString -> ByteString
onLazy f = Lazy.toStrict . f . Lazy.fromStrict

contains :: String -> DirectObject -> Bool
contains needle (NameObject (Name n)) = needle == n
contains needle (Array directObjects) = oneOf directObjects (contains needle)
  where
    oneOf [] _ = False
    oneOf (x:xs) p = p x || oneOf xs p
contains _ _ = False

instance MonadFail m => Box m Clear Object ByteString where
  r Clear (Stream {header, streamContent}) = return $
    case Map.lookup (Name "Filter") header of
      Just directObject
        | contains "FlateDecode" directObject -> onLazy decompress streamContent
      _ -> streamContent
  r _ obj = expected "stream" obj

  w Clear streamContent obj@(Stream {header}) = return $
    case Map.lookup (Name "Filter") header of
      Just directObject
        | contains "FlateDecode" directObject ->
          obj {streamContent = onLazy compress streamContent}
      _ -> obj {streamContent}
  w _ _ obj = expected "stream" obj

instance MonadFail m => Box m Raw Object ByteString where
  r Raw (Stream {streamContent}) = return streamContent
  r _ obj = expected "stream" obj

  w Raw streamContent obj@(Stream {}) = return $ obj {streamContent}
  w _ _ obj = expected "stream" obj