{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
module Arbor.LruCache
    ( lookup
    , makeCache
    , Z.CacheConfig(..)
    , Z.Cache(..)
    , retrieveData
    , evictData
    , entries
    ) where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.Function      ((&))
import Data.Maybe
import Prelude            hiding (lookup)

-- import qualified Arbor.LruCache.Internal.PriorityQueue as PQ
import qualified Arbor.LruCache.Internal.PriorityQueue as PQ
import qualified Arbor.LruCache.Type                   as Z
import qualified Control.Concurrent.STM                as STM
import qualified Data.Map.Strict                       as M

lookup :: Ord k => k -> Z.Cache k v -> IO v
lookup k cache = do
  newTmv <- STM.newTVarIO Nothing
  let maxInFlight       = cache & Z.config & Z.maxRequestsInFlight
  let evict             = cache & Z.evict
  let tRequestsInFlight = cache & Z.requestsInFlight
  let retrieve          = cache & Z.retrieve
  let tEntries          = cache & Z.entries
  let tOccupancy        = cache & Z.occupancy

  join $ STM.atomically $ do
    es <- STM.readTVar tEntries
    case M.lookup k es of
      Just tmv -> do
        registerForEviction k cache
        return $ STM.atomically $ do
          registerForEviction k cache
          STM.readTVar tmv >>= maybe STM.retry return

      Nothing -> do
        requestsInFlight <- STM.readTVar tRequestsInFlight
        if requestsInFlight >= maxInFlight
          then STM.retry
          else do
            STM.writeTVar tRequestsInFlight (requestsInFlight + 1)
            STM.writeTVar tEntries (M.insert k newTmv es)
            return $ do
              v <- catch (retrieve k) $ \(e :: SomeException) -> do
                STM.atomically $ do
                  entries2 <- STM.readTVar tEntries
                  forM_ (M.lookup k entries2) $ \tv -> STM.writeTVar tv (throw e)
                  STM.modifyTVar' tRequestsInFlight pred
                  STM.writeTVar tEntries (M.delete k entries2)
                throw e

              kvsForEviction <- STM.atomically $ do
                STM.writeTVar newTmv (Just v)
                STM.modifyTVar' tRequestsInFlight pred
                STM.modifyTVar' tOccupancy succ

                registerForEviction k cache
                takeEvictionsDue cache

              forM_ kvsForEviction $ uncurry evict

              return v

registerForEviction :: Eq k => k -> Z.Cache k v -> STM.STM ()
registerForEviction k cache = do
  let tEvictionQueue    = cache & Z.evictionQueue
  let tEvictionPriority = cache & Z.evictionPriority

  STM.modifyTVar' tEvictionPriority (+1)
  evictionPriority <- STM.readTVar tEvictionPriority
  STM.modifyTVar' tEvictionQueue (PQ.insert evictionPriority k)

takeEvictionsDue :: Ord k => Z.Cache k v -> STM.STM [(k, v)]
takeEvictionsDue cache = do
  let maxOccupancy      = cache & Z.config & Z.maxOccupancy
  let tEntries          = cache & Z.entries
  let tOccupancy        = cache & Z.occupancy
  let tEvictionQueue    = cache & Z.evictionQueue

  evictionQueue <- STM.readTVar tEvictionQueue
  occupancy <- STM.readTVar tOccupancy

  if occupancy > maxOccupancy
    then case PQ.take (occupancy - maxOccupancy) evictionQueue of
      (ks, evictionQueue') -> do
        STM.writeTVar tEvictionQueue evictionQueue'
        STM.writeTVar tOccupancy maxOccupancy

        removeEvictionsFromEntries ks tEntries

    else return []

removeEvictionsFromEntries :: Ord k => [k] -> STM.TVar (M.Map k (STM.TVar (Maybe v))) -> STM.STM [(k, v)]
removeEvictionsFromEntries ks tEntries = do
  es <- STM.readTVar tEntries

  let kmtmvs = (\k -> (k, M.lookup k es)) <$> ks

  mkvs <- forM kmtmvs $ \(k, mtmv) -> case mtmv of
    Just tmv -> do
      mv <- STM.readTVar tmv
      return ((k, ) <$> mv)
    Nothing -> return Nothing

  let kvs = catMaybes mkvs

  STM.writeTVar tEntries (foldl (flip M.delete) es (fst <$> kvs))

  return kvs

entries :: Ord k => Z.Cache k v -> IO (M.Map k (Maybe v))
entries cache = do
  let tEntries          = cache & Z.entries

  STM.atomically $ do
    m <- STM.readTVar tEntries
    kvs <- forM (M.toList m) $ \(k, tmv) -> do
      mv <- STM.readTVar tmv
      return (k, mv)

    return (M.fromList kvs)

makeCache :: Z.CacheConfig -> (k -> IO v) -> (k -> v -> IO ()) -> IO (Z.Cache k v)
makeCache config retrieve evict = do
  tRequestsInFlight <- STM.newTVarIO 0
  tEntries          <- STM.newTVarIO M.empty
  tOccupancy        <- STM.newTVarIO 0
  tEvictionQueue    <- STM.newTVarIO PQ.empty
  tEvictionPriority <- STM.newTVarIO 0

  return Z.Cache
    { Z.config            = config
    , Z.requestsInFlight  = tRequestsInFlight
    , Z.entries           = tEntries
    , Z.evictionQueue     = tEvictionQueue
    , Z.evictionPriority  = tEvictionPriority
    , Z.occupancy         = tOccupancy
    , Z.retrieve          = retrieve
    , Z.evict             = evict
    }

retrieveData :: (String, Int) -> IO String
retrieveData (s, d) = do
  threadDelay d
  putStrLn $ "Retrieved " ++ show (s, d)
  return $ "Got: " ++ show (s, d)

evictData :: (String, Int) -> String -> IO ()
evictData k v = putStrLn $ "Evicting " ++ show (k, v)