module Ribosome.Data.Ribo(
  Ribo,
  state,
  inspect,
  modify,
  name,
  lockOrSkip,
) where

import Control.Concurrent.STM.TVar (modifyTVar)
import qualified Control.Lens as Lens (view, over, at)
import qualified Data.Map.Strict as Map (insert)
import UnliftIO (finally)
import UnliftIO.STM (TVar, TMVar, atomically, readTVarIO, newTMVarIO, tryTakeTMVar, tryPutTMVar)
import Neovim (Neovim, ask)
import Ribosome.Data.Ribosome (Ribosome(Ribosome), Locks)
import qualified Ribosome.Data.Ribosome as Ribosome (_locks, locks)

type Ribo e = Neovim (Ribosome e)

state :: Ribo (TVar e) e
state = do
  Ribosome _ _ t <- ask
  readTVarIO t

inspect :: (e -> a) -> Ribo (TVar e) a
inspect f = fmap f state

modify :: (e -> e) -> Ribo (TVar e) ()
modify f = do
  Ribosome _ _ t <- ask
  atomically $ modifyTVar t f

name :: Ribo e String
name = do
  Ribosome n _ _ <- ask
  return n

getLocks :: Ribo e Locks
getLocks = do
  Ribosome _ intTv _ <- ask
  int <- readTVarIO intTv
  return $ Ribosome.locks int

inspectLocks :: (Locks -> a) -> Ribo e a
inspectLocks f = fmap f getLocks

modifyLocks :: (Locks -> Locks) -> Ribo e ()
modifyLocks f = do
  Ribosome _ intTv _ <- ask
  atomically $ modifyTVar intTv $ Lens.over Ribosome._locks f

getOrCreateLock :: String -> Ribo e (TMVar ())
getOrCreateLock key = do
  currentLock <- inspectLocks $ Lens.view $ Lens.at key
  case currentLock of
    Just tv -> return tv
    Nothing -> do
      tv <- newTMVarIO ()
      modifyLocks $ Map.insert key tv
      getOrCreateLock key

lockOrSkip :: String -> Ribo e () -> Ribo e ()
lockOrSkip key thunk = do
  currentLock <- getOrCreateLock key
  currentState <- atomically $ tryTakeTMVar currentLock
  case currentState of
    Just _ -> finally thunk $ atomically $ tryPutTMVar currentLock ()
    Nothing -> return ()