{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module PDF.Layer (
      AllObjects(..)
    , Layer(..)
    , unify
  ) where

import Control.Monad.Fail (MonadFail(..))
import Data.Map (Map, (!), mapWithKey, member)
import qualified Data.Map as Map (
    empty, insert, keysSet, lookup, member, union
  )
import qualified Data.Set as Set (delete, toList)
import PDF.Box (Box(..))
import PDF.Object (
      IndexedObjects, IndirectObjCoordinates(..), Object, Occurrence(..)
    , Structure(..), XRefEntry(..), XRefSection, eofMarker, outputBody
  )
import qualified PDF.Output as Output (line)
import PDF.Output (
      ObjectId, Offset(..), Output(..), Resource(..), byteString, getOffset
    , getOffsets, newLine
  )
import Prelude hiding (fail)
import Text.Printf (printf)

data Layer = Layer {
      occurrences :: [Occurrence]
    , objects :: IndexedObjects
    , docStructure :: Structure
  } deriving Show

updateXrefs :: XRefSection -> Map Resource Offset -> (XRefSection, Offset)
updateXrefs xRef offsets = (mapWithKey updateEntry xRef, offsets ! StartXRef)
  where
    updateEntry objectId e@(InUse {}) = e {offset = offsets ! (Object objectId)}
    updateEntry _ e = e

instance Output Layer where
  output (Layer {occurrences, objects, docStructure}) =
    fmap (updateXrefs xRef) <$> getOffsets (outputBody (occurrences, objects))
    >>= \(body, (xref, startXRef)) -> mconcat [
          body
        , Output.line "xref"
        , output xref
        , Output.line "trailer"
        , output trailer, newLine
        , Output.line "startxref"
        , Output.line (printf "%d" (getOffset startXRef))
        , byteString eofMarker
      ]
    where
      Structure {xRef, trailer} = docStructure

data AllObjects = AllObjects

instance Monad m => Box m AllObjects Layer (Map ObjectId Object) where
  r AllObjects = return . objects
  w AllObjects layer@(Layer {occurrences, docStructure}) newObjects =
    return $ layer {
          occurrences = keptOccurrences ++ newOccurrences
        , objects = newObjects
        , docStructure = docStructure {
              xRef = (const $ InUse (Offset 0) 0) <$> newObjects
            }
      }
    where
      filterOccurrences c@(Comment _) (occ, newObjIds) = (c:occ, newObjIds)
      filterOccurrences i@(Indirect (IndirectObjCoordinates {objectId})) (occ, newObjIds)
        | Map.member objectId newObjects = (i:occ, Set.delete objectId newObjIds)
        | otherwise = (occ, newObjIds)
      (keptOccurrences, newObjectIds) =
        foldr filterOccurrences ([], Map.keysSet newObjects) occurrences
      makeOccurrence objectId =
        Indirect (IndirectObjCoordinates {objectId, versionNumber = 0})
      newOccurrences = makeOccurrence <$> Set.toList newObjectIds

instance MonadFail m => Box m ObjectId Layer Object where
  r objectId = maybe (fail errorMsg) return . Map.lookup objectId . objects
    where
      errorMsg = "Unknown " ++ show objectId
  w objectId layer@(Layer {objects}) object
    | Map.member objectId objects =
        return $ layer {objects = Map.insert objectId object objects}
    | otherwise = fail $ "Unknown " ++ show objectId

emptyLayer :: Layer
emptyLayer = Layer {
      docStructure = Structure {xRef = Map.empty, trailer = Map.empty}
    , objects = Map.empty
    , occurrences = []
  }

unify :: [Layer] -> Layer
unify = foldl complete emptyLayer
  where
    complete tmpLayer older =
      let mergedObjects = Map.union (objects tmpLayer) (objects older) in
      Layer {
          docStructure =
            unifyDocStructure (docStructure tmpLayer) (docStructure older)
        , objects = mergedObjects
        , occurrences = 
            unifyOccurrences mergedObjects (occurrences tmpLayer) (occurrences older)
      }

unifyDocStructure :: Structure -> Structure -> Structure
unifyDocStructure new old = Structure {
      xRef = Map.union (xRef new) (xRef old)
    , trailer = Map.union (trailer new) (trailer old)
  }

unifyOccurrences :: IndexedObjects -> [Occurrence] -> [Occurrence] -> [Occurrence]
unifyOccurrences objects new = foldr addOlder new
  where
    addOlder occurrence@(Comment _) existing = occurrence : existing
    addOlder occurrence@(Indirect indirect) existing =
      if objectId indirect `member` objects
      then occurrence : existing
      else existing