{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}

module Polysemy.KVStore
  ( -- * Effect
    KVStore (..),

    -- * Actions
    lookupKV,
    lookupOrThrowKV,
    existsKV,
    writeKV,
    deleteKV,
    updateKV,
    modifyKV,

    -- * Interpretations
    runKVStoreAsState,
    runKVStorePure,
    runKVStoreAsKVStore,
    runKVStoreAsKVStoreSem,
  )
where

import qualified Data.Map as M
import Data.Maybe (isJust)
import Polysemy
import Polysemy.Error
import Polysemy.State

-- | Models things like Redis, HTTP GET/POST, etc. Things that are keyed, have
-- a value, and may or may not be there.
data KVStore k v m a where
  LookupKV :: k -> KVStore k v m (Maybe v)
  UpdateKV :: k -> Maybe v -> KVStore k v m ()

makeSem ''KVStore

-- |
--
-- @since 0.1.0.0
writeKV :: Member (KVStore k v) r => k -> v -> Sem r ()
writeKV k = updateKV k . Just
{-# INLINE writeKV #-}

-- |
--
-- @since 0.1.0.0
deleteKV :: forall k v r. Member (KVStore k v) r => k -> Sem r ()
deleteKV k = updateKV k (Nothing @v)
{-# INLINE deleteKV #-}

-- |
--
-- @since 0.1.0.0
lookupOrThrowKV ::
  Members
    '[ KVStore k v,
       Error e
     ]
    r =>
  (k -> e) ->
  k ->
  Sem r v
lookupOrThrowKV f k =
  fromEither . maybe (Left $ f k) Right =<< lookupKV k

-- |
--
-- @since 0.1.0.0
existsKV :: forall k v r. Member (KVStore k v) r => k -> Sem r Bool
existsKV = fmap isJust . lookupKV @k @v

-- |
--
-- @since 0.1.0.0
modifyKV ::
  Member (KVStore k v) r =>
  -- | Default value if the key isn't present
  v ->
  (v -> v) ->
  k ->
  Sem r ()
modifyKV d f k =
  lookupKV k >>= \case
    Just v -> writeKV k $ f v
    Nothing -> writeKV k $ f d

-- | Run a `KVStore` as a `State` effect containing a `Map`.
--
-- @since 0.1.0.0
runKVStoreAsState :: Ord k => Sem (KVStore k v ': r) a -> Sem (State (M.Map k v) ': r) a
runKVStoreAsState = reinterpret $ \case
  LookupKV k -> gets $ M.lookup k
  UpdateKV k v -> modify $ M.alter (const v) k
{-# INLINE runKVStoreAsState #-}

-- | Run a `KVStore` purely as a `Map`.
--
-- @since 0.1.0.0
runKVStorePure ::
  Ord k =>
  M.Map k v ->
  Sem (KVStore k v ': r) a ->
  Sem r (M.Map k v, a)
runKVStorePure m = runState m . runKVStoreAsState
{-# INLINE runKVStorePure #-}

-- | Run a `KVStore` in terms of another `KVStore` by way of pure key and value
-- transformations.
--
-- @since 0.1.1.0
runKVStoreAsKVStore ::
  forall k v k' v' r a.
  -- | A function to transform the key into the interpreted key.
  (k -> k') ->
  -- | A function to transform the value into the interpreted value.
  (v -> v') ->
  -- | A function to transform the interpreted key back into the current value.
  (v' -> v) ->
  Sem (KVStore k v ': r) a ->
  Sem (KVStore k' v' ': r) a
runKVStoreAsKVStore f g h = reinterpret \case
  LookupKV k -> fmap h <$> lookupKV @k' @v' (f k)
  UpdateKV k x -> updateKV @k' @v' (f k) (fmap g x)
{-# INLINE runKVStoreAsKVStore #-}

-- | Run a `KVStore` in terms of another `KVStore` by way of transforming the
-- keys and values with Sem functions.
--
-- @since 0.1.1.0
runKVStoreAsKVStoreSem ::
  forall k v k' v' r a.
  Members '[KVStore k' v'] r =>
  -- | A function to transform the key into the interpreted key.
  (k -> Sem r k') ->
  -- | A function to transform the value into the interpreted value.
  (v -> Sem r v') ->
  -- | A function to transform the interpreted value back into the current value.
  (v' -> Sem r v) ->
  Sem (KVStore k v ': r) a ->
  Sem r a
runKVStoreAsKVStoreSem f g h = interpret \case
  LookupKV k -> f k >>= lookupKV @k' @v' >>= mapM h
  UpdateKV k x -> do
    z <- f k
    z' <- mapM g x
    updateKV @k' @v' z z'
{-# INLINE runKVStoreAsKVStoreSem #-}