{-# LANGUAGE AllowAmbiguousTypes #-}
{-# 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,
  )
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 #-}