{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
module PDF.Box (
      Box(..)
    , I(..)
    , at
    , atAll
    , edit
  ) where

import Control.Monad.Fail (MonadFail(..))
import Control.Monad.State (MonadState(..))
import Data.Map (Map)
import qualified Data.Map as Map (insert, lookup, member)
import Prelude hiding (fail)

newtype I = I Int

class Monad m => Box m i a b | m a i -> b where
  r :: i -> a -> m b
  w :: i -> b -> a -> m a

at :: Box m i a b => i -> (b -> m b) -> a -> m a
at i f a = r i a >>= f >>= flip (w i) a

atAll :: (Traversable t, Monad m, Box m i a (t b)) => i -> (b -> m b) -> a -> m a
atAll i f = at i $ (mapM f)

edit :: MonadState a m => (a -> m a) -> m ()
edit f = get >>= f >>= put

instance MonadFail m => Box m I [a] a where
  r (I i) = getAt i
    where
      getAt _ [] = fail $ "Index out of bounds " ++ show i
      getAt 0 (x:_) = return x
      getAt k (_:xs) = getAt (k+1) xs
  w (I i) newX = setAt i
    where
      setAt _ [] = fail $ "Index out of bounds " ++ show i
      setAt 0 (_:xs) = return (newX:xs)
      setAt k (x:xs) = (x:)<$>(setAt (k-1) xs)

instance (Ord k, MonadFail m) => Box m k (Map k a) a where
  r k = maybe (fail "Unknown key") return . Map.lookup k
  w k a aMap
    | Map.member k aMap =
        return $ Map.insert k a aMap
    | otherwise = fail "Unknown key"