module Cache where

import Control.Concurrent.STM
import Control.Monad (forM_)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map

import ZNC


-- value, previous (more recently accessed), next (less recently accessed)
data LRUNode a = LRUNode a (TVar (Maybe (LRUNode a))) (TVar (Maybe (LRUNode a)))

-- head (most recent), last (least recent)
data LRUChain a = LRUChain (TVar (Maybe (LRUNode a))) (TVar (Maybe (LRUNode a)))

lruchainNew :: IO (LRUChain a)
lruchainNew = LRUChain <$> newTVarIO Nothing <*> newTVarIO Nothing

lruchainAdd :: LRUChain a -> a -> STM (LRUNode a)
lruchainAdd (LRUChain hd tl) x = do
  mhdnode <- readTVar hd
  case mhdnode of
    Nothing -> do  -- chain empty
      node <- LRUNode x <$> newTVar Nothing <*> newTVar Nothing
      writeTVar hd (Just node)
      writeTVar tl (Just node)
      return node
    Just hdnode -> do
      node <- LRUNode x <$> newTVar Nothing <*> newTVar (Just hdnode)
      writeTVar hd (Just node)
      return node

lruchainBump :: LRUChain a -> LRUNode a -> STM ()
lruchainBump (LRUChain hd _) node@(LRUNode _ prevvar nextvar) = do
  mprev <- readTVar prevvar
  case mprev of
    Nothing -> return ()  -- already most recent
    Just (LRUNode _ _ prevnextvar) -> do
      -- remove node from chain where it is now
      mnext <- readTVar nextvar
      writeTVar prevnextvar mnext
      case mnext of
        Nothing -> return ()
        Just (LRUNode _ nextprevvar _) -> writeTVar nextprevvar mprev
      -- add node to chain at head
      writeTVar nextvar =<< readTVar hd
      writeTVar hd (Just node)

lruchainEvict :: LRUChain a -> STM (Maybe a)
lruchainEvict (LRUChain _ tl) = do
  mtlnode <- readTVar tl
  case mtlnode of
    Nothing -> return Nothing
    Just (LRUNode x prevvar _) -> do
      mprev <- readTVar prevvar
      writeTVar tl mprev
      case mprev of
        Nothing -> return ()
        Just (LRUNode _ _ prevnextvar) -> writeTVar prevnextvar Nothing
      return (Just x)

-- maxsize, elements in cache, linked list
data LRU a = LRU Int (TVar (Map a (LRUNode a))) (LRUChain a)

lruNew :: Ord a => Int -> IO (LRU a)
lruNew maxsize = LRU maxsize <$> newTVarIO mempty <*> lruchainNew

-- | If a value got evicted, it is returned
lruBump :: Ord a => LRU a -> a -> IO (Maybe a)
lruBump (LRU maxsize mapvar chain) value = atomically $ do
  mp <- readTVar mapvar
  case Map.lookup value mp of
    Nothing -> do
      node <- lruchainAdd chain value
      writeTVar mapvar $! Map.insert value node mp
      if Map.size mp >= maxsize
        then do
          mret <- lruchainEvict chain
          forM_ mret $ \ret -> writeTVar mapvar $! Map.delete ret mp
          return mret
        else return Nothing
    Just node -> do
      lruchainBump chain node
      return Nothing