module Cache where

import Control.Monad (when)
import Data.IORef
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import System.Clock qualified as Clock
import System.Random


data Cache k v = Cache Int (IORef (Map k (v, Clock.TimeSpec))) (IORef StdGen)

cacheNew :: Int -> IO (Cache k v)
cacheNew maxsize = Cache maxsize <$> newIORef Map.empty <*> (newIORef =<< initStdGen)

cacheAdd :: (Ord k, Show k) => Cache k v -> k -> v -> IO ()
cacheAdd cache@(Cache maxsize ref _) key val = do
  now <- Clock.getTime Clock.Monotonic
  sz <- atomicModifyIORef' ref $ \mp ->
    case fst (Map.insertLookupWithKey (\_ _ _ -> (val, now)) key (val, now) mp) of
      Nothing -> (Map.insert key (val, now) mp, Map.size mp + 1)
      Just _ -> (Map.insert key (val, now) mp, 0)
  when (sz > maxsize) $
    cachePrune cache

cacheLookup :: Ord k => Cache k v -> k -> IO (Maybe v)
cacheLookup (Cache _ ref _) key = fmap fst . Map.lookup key <$> readIORef ref

-- Uses 2-random LRU (https://danluu.com/2choices-eviction/)
cachePrune :: Ord k => Cache k v -> IO ()
cachePrune (Cache maxsize ref genref) = do go =<< atomicModifyIORef' genref splitGen
  where
    go gen0 = do
      (gen', sz') <- atomicModifyIORef' ref $ \mp ->
        let sz = Map.size mp
            (idx1, gen1) = uniformR (0, sz - 1) gen0 ; (_, (_, tm1)) = Map.elemAt idx1 mp
            (idx2, gen2) = uniformR (0, sz - 1) gen1 ; (_, (_, tm2)) = Map.elemAt idx2 mp
            oldestIdx = if tm1 < tm2 then idx1 else idx2
        in if sz > maxsize
             then (Map.deleteAt oldestIdx mp, (gen2, sz - 1))
             else (mp, (gen0, sz))
      when (sz' > maxsize) $
        go gen'