content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TypeOperators #-}
  
  module Control.Monad.Schedule.OSThreadPool where
  
  -- base
  import Control.Concurrent
  import Control.Monad (forM, replicateM, void)
  import Control.Monad.IO.Class
  import Data.Either (partitionEithers)
  import Data.List.NonEmpty hiding (cycle, zip)
  import Data.Proxy
  import GHC.TypeLits
  import Prelude hiding (take)
  
  -- stm
  import Control.Concurrent.STM
  
  -- monad-schedule
  import Control.Monad.Schedule.Class
  
  newtype OSThreadPool (n :: Nat) a = OSThreadPool {unOSThreadPool :: IO a}
    deriving (Functor, Applicative, Monad, MonadIO)
  
  data WorkerLink a = WorkerLink
    { jobTChan :: TChan (Maybe (IO a))
    , resultTChan :: TChan a
    }
  
  putJob :: WorkerLink a -> OSThreadPool n a -> IO ()
  putJob WorkerLink {..} OSThreadPool {..} =
    atomically $
      writeTChan jobTChan $
        Just unOSThreadPool
  
  makeWorkerLink :: IO (WorkerLink a)
  makeWorkerLink = do
    jobTChan <- atomically newTChan
    resultTChan <- atomically newTChan
    let worker = do
          job <- atomically $ readTChan jobTChan
          case job of
            Nothing -> return ()
            Just action -> do
              result <- action
              atomically $ writeTChan resultTChan result
              worker
    void $ forkOS worker
    return WorkerLink {..}
  
  proxyForActions :: NonEmpty (OSThreadPool n a) -> Proxy n
  proxyForActions _ = Proxy
  
  instance (KnownNat n, (1 <=? n) ~ True) => MonadSchedule (OSThreadPool n) where
    schedule actions = OSThreadPool $ do
      let n = natVal $ proxyForActions actions
      workerLinks <- replicateM (fromInteger n) makeWorkerLink
      backgroundActions <- forM (zip (cycle workerLinks) (toList actions)) $
        \(link, action) -> do
          putJob link action
          return $ resultTChan link
      pollPools backgroundActions
      where
        pollPools :: [TChan a] -> IO (NonEmpty a, [OSThreadPool n a])
        pollPools chans = do
          results <- traverse pollPool chans
          case partitionEithers results of
            (_, []) -> do
              threadDelay 1000
              pollPools chans
            (remainingChans, a : as) ->
              return
                ( a :| as
                , OSThreadPool . atomically . readTChan <$> remainingChans
                )
  
        pollPool :: TChan a -> IO (Either (TChan a) a)
        pollPool chan = maybe (Left chan) Right <$> atomically (tryReadTChan chan)
  