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

import Data.Id (Id(..), Indexed)
import Data.Map (Map, (!), mapWithKey, member)
import qualified Data.Map as Map (empty, keysSet, member, union)
import qualified Data.Set as Set (delete, toList)
import PDF.Box (Box(..))
import PDF.Object (
      IndirectObjCoordinates(..), Object, Occurrence(..)
    , Structure(..), XRefEntry(..), XRefSection, eofMarker, outputBody
  )
import qualified PDF.Output as Output (line)
import PDF.Output (
      Offset(..), Output(..), Resource(..), byteString, getOffset
    , getOffsets, newLine
  )
import Text.Printf (printf)

data Layer = Layer {
      occurrences :: [Occurrence]
    , objects :: (Indexed Object)
    , 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 ! (ObjectId $ getId 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 Objects = Objects

instance Monad m => Box m Objects Layer (Indexed Object) where
  r Objects = return . objects
  w Objects newObjects layer@(Layer {occurrences, docStructure}) =
    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

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 :: (Indexed Object) -> [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