{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE NumericUnderscores #-}
import System.Mem
import Control.Concurrent
import Control.Concurrent.MVar
import Control.Concurrent.Thread.Storage
import Control.Monad
import Data.IORef
import Data.List hiding (lookup)
import Test.Hspec
import Prelude hiding (lookup)

main :: IO ()
main = hspec $ do

  describe "cleanup" $ do
    it "works" $ do
      let n = 100_000
      gate <- newEmptyMVar
      doneRef <- newIORef (0 :: Int)
      tsm <- newThreadStorageMapWith (n * 2)
      replicateM_ n $ do
        forkIO $ do
          attach tsm ()
          readMVar gate
          atomicModifyIORef' doneRef (\x -> (x + 1, ()))

      -- Wait for all threads to have attached
      waitForCount tsm n

      -- Release all threads
      putMVar gate ()

      -- Wait for all threads to finish
      spinUntil $ do
        c <- readIORef doneRef
        pure (c >= n)

      -- Give finalizers a chance to run
      waitUntilGC $ do
        items <- storedItems tsm
        pure (null items)

      thingsStillInStorage <- storedItems tsm
      sort thingsStillInStorage `shouldBe` []

    it "doesn't happen while a thread is still alive" $ do
      tsm <- newThreadStorageMapWith 64
      gate <- newEmptyMVar
      resultVar <- newEmptyMVar
      forkIO $ do
        attach tsm ()
        readMVar gate
        putMVar resultVar =<< lookup tsm

      -- Wait until the child has attached
      waitForCount tsm 1

      performGC
      yield

      -- The entry should still be there since the thread is alive
      putMVar gate ()
      result <- readMVar resultVar
      result `shouldBe` Just ()

      -- Now the thread is dead; give finalizers a chance
      replicateM_ 5 $ performGC >> yield

      waitUntilGC $ do
        items <- storedItems tsm
        pure (null items)

      items <- storedItems tsm
      items `shouldBe` []

  describe "detach" $ do
    it "returns previous value and clears the entry" $ do
      tsm <- newThreadStorageMapWith 16
      resultVar <- newEmptyMVar
      forkIO $ do
        attach tsm (42 :: Int)
        prev <- detach tsm
        after <- lookup tsm
        putMVar resultVar (prev, after)

      (prev, after) <- readMVar resultVar
      prev `shouldBe` Just 42
      after `shouldBe` Nothing

    it "returns Nothing when no value is attached" $ do
      tsm <- newThreadStorageMapWith 16
      resultVar <- newEmptyMVar
      forkIO $ putMVar resultVar =<< detach tsm
      result <- readMVar resultVar
      result `shouldBe` (Nothing :: Maybe Int)

    it "makes the value eligible for GC" $ do
      tsm <- newThreadStorageMapWith 16
      gate <- newEmptyMVar
      forkIO $ do
        attach tsm (42 :: Int)
        _ <- detach tsm
        readMVar gate

      spinUntil $ do
        items <- storedItems tsm
        pure (null items)

      putMVar gate ()

  describe "update" $ do
    it "can insert via Nothing -> Just" $ do
      tsm <- newThreadStorageMapWith 16
      resultVar <- newEmptyMVar
      forkIO $ do
        update tsm (\_ -> (Just (99 :: Int), ()))
        val <- lookup tsm
        putMVar resultVar val

      result <- readMVar resultVar
      result `shouldBe` Just 99

    it "can remove via Just -> Nothing" $ do
      tsm <- newThreadStorageMapWith 16
      resultVar <- newEmptyMVar
      forkIO $ do
        attach tsm (7 :: Int)
        removed <- update tsm (\old -> (Nothing, old))
        after <- lookup tsm
        putMVar resultVar (removed, after)

      (removed, after) <- readMVar resultVar
      removed `shouldBe` Just 7
      after `shouldBe` Nothing

  describe "resize" $ do
    it "grows the table when capacity is exceeded" $ do
      tsm <- newThreadStorageMapWith 16
      let n = 200
      gate <- newEmptyMVar
      doneRef <- newIORef (0 :: Int)
      replicateM_ n $ forkIO $ do
        attach tsm ()
        readMVar gate
        atomicModifyIORef' doneRef (\x -> (x + 1, ()))

      waitForCount tsm n

      items <- storedItems tsm
      length items `shouldBe` n

      putMVar gate ()
      spinUntil $ do
        c <- readIORef doneRef
        pure (c >= n)

      waitUntilGC $ do
        remaining <- storedItems tsm
        pure (null remaining)

    it "preserves values across resize" $ do
      tsm <- newThreadStorageMapWith 16
      let n = 100
      resultRefs <- replicateM n newEmptyMVar
      gate <- newEmptyMVar

      forM_ (zip [1 :: Int ..] resultRefs) $ \(i, mv) ->
        forkIO $ do
          attach tsm i
          readMVar gate
          val <- lookup tsm
          putMVar mv val

      waitForCount tsm n

      putMVar gate ()

      results <- mapM readMVar resultRefs
      let expected = fmap Just [1 .. n]
      sort results `shouldBe` sort expected


waitForCount :: ThreadStorageMap a -> Int -> IO ()
waitForCount tsm target = spinUntil $ do
  items <- storedItems tsm
  pure (length items >= target)


-- | Spin-wait without GC.  Suitable for waiting on concurrent threads to
-- make progress (insert, signal, etc.).
spinUntil :: IO Bool -> IO ()
spinUntil check = go (500000 :: Int)
  where
    go 0 = error "spinUntil: timed out"
    go !n = do
      done <- check
      unless done $ do
        yield
        go (n - 1)


-- | Spin-wait with periodic 'performGC'.  Use only when waiting for GC
-- finalizers to fire (e.g. dead-thread cleanup).
waitUntilGC :: IO Bool -> IO ()
waitUntilGC check = go (5000 :: Int)
  where
    go 0 = error "waitUntilGC: timed out"
    go !n = do
      done <- check
      unless done $ do
        yield
        performGC
        go (n - 1)