module Main
  ( main )
where

import Control.Concurrent (threadDelay)
import Control.Concurrent.Async qualified as Async
import Control.Concurrent.STM (atomically)
import Control.Concurrent.STM.TChan qualified as TChan
import Control.Concurrent.STM.TVar qualified as TVar
import Control.Concurrent.Throttle (throttle)
import Data.Time.Clock.POSIX (getPOSIXTime)
import Test.Tasty
import Test.Tasty.HUnit

main :: IO ()
main = defaultMain tests

tests :: TestTree
tests = testGroup "throttle tests" unitTests



unitTests :: [TestTree]
unitTests =
  [ testCase "can call a function once" $ do
      let val = 100

      withTestContext $ \(tvar, delay, tchan, _a) -> do
        start <- unixTime
        atomically $ TChan.writeTChan tchan (1, val)
        
        threadDelay (delay * 3)
        
        actual <- TVar.readTVarIO tvar
        length actual @?= 1
        let [(val', n)] = actual
        val @?= val'
        assertBool "not enough delay" $ n - start >= delay

  , testCase "doesn't throttle independent messages" $ do
      withTestContext $ \(tvar, delay, tchan, _a) -> do
        atomically $ TChan.writeTChan tchan (1, 100)

        -- this one has a different id so should be called as well
        atomically $ TChan.writeTChan tchan (2, 101)
        
        threadDelay (delay * 3)
        
        actual <- TVar.readTVarIO tvar
        length actual @?= 2

  , testCase "can handle simple throttling" $ do
      withTestContext $ \(tvar, delay, tchan, _a) -> do
        -- this one should be discarded
        atomically $ TChan.writeTChan tchan (1, 100)

        -- this one should be called (the same id)
        atomically $ TChan.writeTChan tchan (1, 101)
        
        threadDelay (delay * 3)
        
        actual <- TVar.readTVarIO tvar
        length actual @?= 1
        let [(val', _)] = actual
        val' @?= 101

  , testCase "can handle throttling for more massive messages" $ do
      withTestContext $ \(tvar, delay, tchan, _a) -> do
        mapM_ (\val ->
          atomically $ TChan.writeTChan tchan (1, val)) [100..200]
        
        threadDelay (delay * 3)

        -- Only last one should be called
        actual <- TVar.readTVarIO tvar
        length actual @?= 1
        let [(val', _)] = actual
        val' @?= 200
  ]


withTestContext :: ((TVar.TVar [(Int, Int)], Int, TChan.TChan (Int, Int), Async.Async ()) -> IO b)
                -> IO b
withTestContext cb = do
  tvar <- TVar.newTVarIO []
  
  let delay = 200
  tchan <- TChan.newTChanIO :: IO (TChan.TChan (Int, Int))
  Async.withAsync (throttle delay tchan (action tvar)) $ \a ->
    cb (tvar, delay, tchan, a)


action :: TVar.TVar [(Int, Int)] -> Int -> IO ()
action tvar v = do
  now <- unixTime
  atomically $ TVar.modifyTVar tvar (\l -> l ++ [(v, now)])


unixTime :: IO Int
unixTime = (round . (* 1000000)) <$> getPOSIXTime