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

import Control.Monad.Except (MonadError(..))
import Control.Monad.Reader (ReaderT)
import Data.Foldable (foldl')
import Data.Id (Id(..), Indexed, keysSet, mapWithKey, member)
import qualified Data.Id as Id (empty, insert, lookup, union)
import Data.Map (Map, (!))
import qualified Data.Map as Map (empty, lookup, union)
import qualified Data.IntSet as IntSet (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
type LayerReader m = ReaderT Layer m

updateXRefs :: XRefSection -> Map Resource Offset -> (XRefSection, Offset)
updateXRefs xRef offsets = (mapWithKey updateEntry xRef, offsets ! StartXRef)
  where
    updateEntry objectId e@(InUse {offset}) =
      case Map.lookup (ObjectId $ getId objectId) offsets of
        Nothing -> Free {nextFree = Id $ getOffset offset, generation = 65535}
        Just newOffset -> e {offset = newOffset}
    updateEntry _ e = e

instance Output Layer where
  output (Layer {occurrences, objects, docStructure}) = do
    (body, savedOffsets) <- getOffsets (outputBody (occurrences, objects))
    let (newXRef, startXRef) = updateXRefs xRef savedOffsets
    mconcat [
          return body
        , Output.line "xref"
        , output newXRef
        , 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)
        | member objectId newObjects = (i:occ, IntSet.delete (getId objectId) newObjIds)
        | otherwise = (occ, newObjIds)
      (keptOccurrences, newObjectIds) =
        foldr filterOccurrences ([], keysSet newObjects) occurrences
      makeOccurrence objectId =
        Indirect (IndirectObjCoordinates {objectId, versionNumber = 0})
      newOccurrences = (makeOccurrence . Id) <$> IntSet.toList newObjectIds

instance MonadError String m => Box m (Id Object) Layer Object where
  r objectId =
    maybe (throwError "Unknown key") return . Id.lookup objectId . objects
  w objectId a layer@(Layer {objects})
    | member objectId objects = return $
      layer {objects = Id.insert objectId a objects}
    | otherwise = throwError "Unknown key"

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

unify :: [Layer] -> Layer
unify = foldl' complete emptyLayer
  where
    complete tmpLayer older =
      let mergedObjects = Id.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 = Id.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