-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

module Utils.LRU
  ( LRUTick
  , LRU, initLRU
  , insertLRU, lookupLRU, touchLRU, timeoutLRU
  , LRUHandle, initLRUHandle
  , insertLRUHandle, lookupLRUHandle, touchLRUHandle, timeoutLRUHandle
  , readLRUHandle
  , lruStoreSize
  , getLRUWeight
  , describeLRU
  ) where

import ClassyPrelude

import Data.OrdPSQ (OrdPSQ)
import qualified Data.OrdPSQ as OrdPSQ

import Control.Lens

-- https://jaspervdj.be/posts/2015-02-24-lru-cache.html


newtype LRUTick = LRUTick { _getLRUTick :: Word64 }
  deriving (Eq, Ord, Show)
  deriving newtype (NFData)

makeLenses ''LRUTick

data LRU k t w v = LRU
  { lruStore :: !(OrdPSQ k (t, LRUTick) (v, w))
  , lruWeight :: !w
  , lruMaximumWeight :: !w
  }

instance (NFData k, NFData t, NFData w, NFData v) => NFData (LRU k t w v) where
  rnf LRU{..} = rnf lruStore
          `seq` rnf lruWeight
          `seq` rnf lruMaximumWeight

describeLRU :: Show w
            => LRU k t w v
            -> String
describeLRU LRU{..} = intercalate ", "
  [ "lruStore: "         <> show (OrdPSQ.size lruStore)
  , "lruWeight: "        <> show lruWeight
  , "lruMaximumWeight: " <> show lruMaximumWeight
  ]

lruStoreSize :: LRU k t w v -> Int
lruStoreSize = OrdPSQ.size . lruStore

getLRUWeight :: LRU k t w v -> w
getLRUWeight = lruWeight

initialLRUTick, maximumLRUTick :: LRUTick
initialLRUTick = LRUTick 0
maximumLRUTick = LRUTick maxBound

initLRU :: forall k t w v.
           Integral w
        => w -- ^ @lruMaximumWeight@
        -> (LRU k t w v, LRUTick)
initLRU lruMaximumWeight
  | lruMaximumWeight < 0 = error "initLRU given negative maximum weight"
  | otherwise = (, initialLRUTick) LRU
    { lruStore = OrdPSQ.empty
    , lruWeight = 0
    , lruMaximumWeight
    }

insertLRU :: forall k t w v.
             ( Ord k, Ord t
             , Integral w
             )
          => k
          -> t
          -> Maybe (v, w)
          -> LRU k t w v
          -> LRUTick -> (LRU k t w v, LRUTick)
insertLRU k t newVal oldLRU@LRU{..} now
  | later <= initialLRUTick = uncurry (insertLRU k t newVal) $ initLRU lruMaximumWeight
  | Just (_, w) <- newVal, w > lruMaximumWeight = (oldLRU, now)
  | Just (_, w) <- newVal = (, later) $ 
      let (lruStore', lruWeight') = evictToSize (lruMaximumWeight - w) lruStore lruWeight
          (fromMaybe 0 . preview (_Just . _2 . _2) -> oldWeight, lruStore'')
            = OrdPSQ.alter (, ((t, now), ) <$> newVal) k lruStore'
       in oldLRU
            { lruStore = lruStore''
            , lruWeight = lruWeight' - oldWeight + w
            }
  | Just (_, (_, w), lruStore') <- OrdPSQ.deleteView k lruStore = (, now) oldLRU
      { lruStore = lruStore'
      , lruWeight = lruWeight - w
      }
  | otherwise = (oldLRU, now)
  where
    later :: LRUTick
    later = over getLRUTick succ now

    evictToSize :: w -> OrdPSQ k (t, LRUTick) (v, w) -> w -> (OrdPSQ k (t, LRUTick) (v, w), w)
    evictToSize tSize c cSize
      | cSize <= tSize = (c, cSize)
      | Just (_, _, (_, w'), c') <- OrdPSQ.minView c = evictToSize tSize c' (cSize - w')
      | otherwise = error "evictToSize: cannot reach required size through eviction"

lookupLRU :: forall k t w v.
             Ord k
          => k
          -> LRU k t w v
          -> Maybe (v, w)
lookupLRU k LRU{..} = view _2 <$> OrdPSQ.lookup k lruStore

touchLRU :: forall k t w v.
            ( Ord k, Ord t
            , Integral w
            )
         => k
         -> t
         -> LRU k t w v
         -> LRUTick -> ((LRU k t w v, LRUTick), Maybe (v, w))
touchLRU k t oldLRU@LRU{..} now
  | (Just (_, v), _) <- altered
  , later <= initialLRUTick = (, Just v) . uncurry (insertLRU k t $ Just v) $ initLRU lruMaximumWeight
  | (Just (_, v), lruStore') <- altered = ((oldLRU{ lruStore = lruStore' }, later), Just v)
  | otherwise = ((oldLRU, now), Nothing)
  where 
    altered = OrdPSQ.alter (\oldVal -> (oldVal, over _1 (max (t, later)) <$> oldVal)) k lruStore
    
    later :: LRUTick
    later = over getLRUTick succ now

timeoutLRU :: forall k t w v.
              ( Ord k, Ord t
              , Integral w
              )
           => t
           -> LRU k t w v
           -> LRU k t w v
timeoutLRU t oldLRU@LRU{..} = oldLRU
  { lruStore = lruStore'
  , lruWeight = lruWeight - evictedWeight
  }
  where
    (evicted, lruStore') = OrdPSQ.atMostView (t, maximumLRUTick) lruStore
    evictedWeight = sumOf (folded . _3 . _2) evicted

newtype LRUHandle k t w v = LRUHandle { _getLRUHandle :: IORef (LRU k t w v, LRUTick) }
  deriving (Eq)

initLRUHandle :: forall k t w v m.
                 ( MonadIO m
                 , Integral w
                 )
              => w -- ^ @lruMaximumWeight@
              -> m (LRUHandle k t w v)
initLRUHandle maxWeight = fmap LRUHandle . newIORef $ initLRU maxWeight

insertLRUHandle :: forall k t w v m.
                   ( MonadIO m
                   , Ord k, Ord t
                   , Integral w
                   , NFData k, NFData t, NFData w, NFData v
                   )
                => LRUHandle k t w v
                -> k
                -> t
                -> (v, w)
                -> m ()
insertLRUHandle (LRUHandle lruVar) k t newVal
  = modifyIORef' lruVar $ force . uncurry (insertLRU k t $ Just newVal)

lookupLRUHandle :: forall k t w v m.
                   ( MonadIO m
                   , Ord k
                   )
                => LRUHandle k t w v
                -> k
                -> m (Maybe (v, w))
lookupLRUHandle (LRUHandle lruVar) k
  = views _1 (lookupLRU k) <$> readIORef lruVar

touchLRUHandle :: forall k t w v m.
                  ( MonadIO m
                  , Ord k, Ord t
                  , Integral w
                  , NFData k, NFData t, NFData w, NFData v
                  )
               => LRUHandle k t w v
               -> k
               -> t
               -> m (Maybe (v, w))
touchLRUHandle (LRUHandle lruVar) k t = do
  oldLRU <- readIORef lruVar
  let (newLRU, touched) = uncurry (touchLRU k t) oldLRU
  force newLRU `seq` writeIORef lruVar newLRU
  return touched

timeoutLRUHandle :: forall k t w v m.
                    ( MonadIO m
                    , Ord k, Ord t
                    , Integral w
                    , NFData k, NFData t, NFData w, NFData v
                    )
                 => LRUHandle k t w v
                 -> t
                 -> m ()
timeoutLRUHandle (LRUHandle lruVar) t
  = modifyIORef' lruVar $ force . over _1 (timeoutLRU t)

readLRUHandle :: MonadIO m
              => LRUHandle k t w v
              -> m (LRU k t w v, LRUTick)
readLRUHandle (LRUHandle lruVar) = readIORef lruVar