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

import Control.Monad.Except (ExceptT(..), runExceptT)
import Control.Monad.Fail (MonadFail(..))
import Control.Monad.Reader (ReaderT, runReaderT)
import Control.Monad.State (MonadState(..))
import Data.Id (Id, IdMap)
import qualified Data.Id as Id (insert, lookup)
import Data.Map (Map)
import qualified Data.Map as Map (insert, lookup)
import Data.OrderedMap (OrderedMap)
import qualified Data.OrderedMap as OrderedMap (lookup, set)
import Prelude hiding (fail)

runRO :: MonadState s m => ReaderT s m a -> m a
runRO ro = get >>= runReaderT ro

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 = return . Map.insert k a

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 MonadFail m => Box m (Id k) (IdMap k a) a where
  r k = maybe (fail "Unknown key") return . Id.lookup k
  w k a = return . Id.insert k a

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