{-# LANGUAGE GADTs #-}
module Data.HMemDb.KeyBackends
    (
     PreKeyBack(PreKeyBack),
     KeyBack(KeyBack),
     deleteFromKeys,
     insertIntoKeys,
     modifyInKeys
    ) where
import Control.Concurrent.STM (STM, TVar, modifyTVar', readTVar, writeTVar)
import Control.Monad (guard, void)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Cont (ContT(ContT, runContT))
import Control.Monad.Trans.Maybe (runMaybeT)
import Data.Foldable (for_)
import Data.Map (Map)
import qualified Data.Map as M (delete, insert, lookup, update)
import Data.HMemDb.MapTVar (MS, liftMaybe)
import Data.HMemDb.RefContainer (RefContainer(insRef, delRef))
import Data.HMemDb.References (Ref(refIndex))
brackets :: Monad m => m a -> (a -> m r) -> ContT r m ()
brackets before after =
    do a <- lift before
       ContT $ \f -> f () >> after a
data PreKeyBack s r a where
    PreKeyBack ::
        Ord i => (a -> i)
        -> TVar (Map Integer i)
        -> TVar (Map i (s (Ref r)))
        -> PreKeyBack s r a
data KeyBack r a where KeyBack :: RefContainer s => PreKeyBack s r a -> KeyBack r a
insertIntoKey :: RefContainer s => a -> Ref r -> PreKeyBack s r a -> ContT () MS ()
insertIntoKey new ref (PreKeyBack g ii tv) =
    brackets before after where
        before =
            do let i = g new
               mp <- lift $ readTVar tv
               s <- liftMaybe $ insRef ref $ M.lookup i mp
               return (M.insert (refIndex ref) i, M.insert i s mp)
        after ~(makeNewImp, newMp) =
            lift $ do
              modifyTVar' ii makeNewImp
              writeTVar tv $! newMp
insertIntoKeys :: a -> Ref r -> [KeyBack r a] -> MS ()
insertIntoKeys a ref ks =
    runContT (for_ ks $ \(KeyBack pk) -> insertIntoKey a ref pk) return
deleteFromKey :: RefContainer s => Ref r -> PreKeyBack s r a -> STM ()
deleteFromKey ref (PreKeyBack _ ii tv) =
    void $ runMaybeT $
    do imp <- lift $ readTVar ii
       i <- liftMaybe $ M.lookup (refIndex ref) imp
       mp <- lift $ readTVar tv
       s <- liftMaybe $ M.lookup i mp
       lift $ do
         writeTVar ii $! M.delete (refIndex ref) imp
         writeTVar tv $! M.update (const $ delRef ref s) i mp
deleteFromKeys :: Ref r -> [KeyBack r a] -> STM ()
deleteFromKeys ref ks = for_ ks $ \(KeyBack pk) -> deleteFromKey ref pk
modifyInKey :: RefContainer s => a -> Ref r -> PreKeyBack s r a -> ContT () STM ()
modifyInKey new ref (PreKeyBack g ii tv) =
    brackets before after where
        before =
            runMaybeT $ do
              let i = refIndex ref
              imp <- lift $ readTVar ii
              oldI <- liftMaybe $ M.lookup i imp
              mp <- lift $ readTVar tv
              let newI = g new
              guard $ newI /= oldI
              oldS <- liftMaybe $ M.lookup oldI mp
              let tempMp = M.update (const $ delRef ref oldS) oldI mp
              newS <- liftMaybe $ insRef ref $ M.lookup newI tempMp
              return (M.insert i newI imp, M.insert newI newS tempMp)
        after Nothing = return ()
        after (Just ~(newImp, newMp)) =
            do writeTVar ii $! newImp
               writeTVar tv $! newMp
modifyInKeys :: a -> Ref r -> [KeyBack r a] -> STM ()
modifyInKeys a ref ks =
    runContT (for_ ks $ \(KeyBack pk) -> modifyInKey a ref pk) return