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

import Control.Monad.Except (ExceptT(..), runExceptT)
import Control.Monad.Fail (MonadFail(..))
import Control.Monad.State (MonadState(..))
--import Control.Monad.Reader (MonadReader(..))
import Data.Map (Map)
import qualified Data.Map as Map (insert, lookup, member)
import Data.OrderedMap (OrderedMap)
import qualified Data.OrderedMap as OrderedMap (lookup, set)
import Prelude hiding (fail)

{-
instance (Monad m, MonadState s m) => MonadReader s m where
  ask = get
  local f m = do
    backup <- get
    put $ f backup
    result <- m
    put backup
    return result
-}

newtype Index = Index Int
newtype Maybe_ x = Maybe_ x
newtype Either_ b x = Either_ x

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 Index [a] a where
  r (Index i) [] = fail $ "Index out of bounds " ++ show i
  r (Index 0) (x:_) = return x
  r (Index i) (_:xs) = r (Index (i-1)) xs

  w (Index i) _ [] = fail $ "Index out of bounds " ++ show i
  w (Index 0) newX (_:xs) = return (newX:xs)
  w (Index i) newX (x:xs) = (x:) <$> w (Index (i-1)) newX 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"

instance (Ord k, MonadFail m) => Box m k (OrderedMap k a) a where
  r k = maybe (fail "Unknown key") return . OrderedMap.lookup k
  w k a orderedMap = r k orderedMap >> return (OrderedMap.set k a orderedMap)

instance (Monad m, Box Maybe i a b) => Box m (Maybe_ i) a (Maybe b) where
  r (Maybe_ i) = return . r i
  w (Maybe_ i) (Just b) a =
    return . maybe a id $ w i b a
  w _ _ obj = return obj

instance (Monad m, Box (ExceptT e m) i a b) => Box m (Either_ e i) a (Either e b) where
  r (Either_ i) a = runExceptT (r i a)
  w (Either_ i) (Right b) a =
    either (const a) id <$> runExceptT (w i b a :: ExceptT e m a)
  w _ _ a = return a