module Engine.Types.RefCounted where

import RIO

import Control.Monad.Trans.Resource (allocate_)
import GHC.IO.Exception (IOErrorType(UserError), IOException(IOError))
import UnliftIO.Resource (MonadResource, ReleaseKey)
import UnliftIO.Resource qualified as Resource

-- | A 'RefCounted' will perform the specified action when the count reaches 0
data RefCounted = RefCounted
  { rcCount  :: IORef Int
  , rcAction :: IO ()
  }

-- | Create a counter with a value of 1
newRefCounted :: MonadIO m => IO () -> m RefCounted
newRefCounted rcAction = do
  rcCount <- liftIO $ newIORef 1
  pure RefCounted{..}

-- | Decrement the value, the action will be run promptly and in
-- this thread if the counter reached 0.
releaseRefCounted :: MonadIO m => RefCounted -> m ()
releaseRefCounted RefCounted{..} =
  liftIO $ mask \_ ->
    atomicModifyIORef' rcCount (\c -> (c - 1, c - 1)) >>= \case
      n | n < 0 ->
        throwM $ IOError
          Nothing
          UserError
          ""
          "Ref counted value decremented below 0"
          Nothing
          Nothing

      0 ->
        rcAction

      _stillReferenced ->
        pure ()

-- | Increment the counter by 1
takeRefCounted :: MonadIO m => RefCounted -> m ()
takeRefCounted RefCounted{..} =
  liftIO $ atomicModifyIORef' rcCount \c -> (c + 1, ())

-- | Hold a reference for the duration of the 'MonadResource' action
resourceTRefCount :: MonadResource f => RefCounted -> f ()
resourceTRefCount r =
  void $ allocate_ (takeRefCounted r) (releaseRefCounted r)

wrapped :: MonadResource m => m (ReleaseKey, a) -> m (RefCounted, a)
wrapped action = do
  (key, res) <- action
  rc <- newRefCounted $ Resource.release key
  pure (rc, res)