module Jobs.Handler.Intervals.Utils
  ( mkIntervals, mkIntervalsCached
  , getCurrentInterval
  , currentIntervalCached
  ) where

import Import hiding (init, maximumBy, cached)

import Control.Monad.Random.Lazy (evalRand, mkStdGen)
import System.Random.Shuffle (shuffleM)
import Data.List ((!!), unfoldr, maximumBy, init, genericLength)
import qualified Data.ByteString as ByteString
import Data.Bits (Bits(shiftR))

import qualified Data.Map.Strict as Map
  

mkIntervals :: forall a. Integer -> (ByteString -> Maybe a) -> Natural -> [(Maybe a, Maybe a)]
mkIntervals bytes fromBS numIterations = zip (Nothing : init intervals') intervals'
  where
    bits = bytes * 8
    base :: Integer
    base = 2 ^ bits

    -- | Exclusive upper bounds
    intervals
      | numIterations <= 0 = pure base
      | otherwise = go protoIntervals ^.. folded . _1
      where
        go [] = []
        go ints
          | maximumOf (folded . _1) ints == Just base = ints
          | otherwise = go $ lts ++ over _1 succ (over _2 (subtract $ toInteger numIterations) closest) : map (over _1 succ) gts
          where
            closest = maximumBy (comparing $ view _2) ints
            (lts, geqs) = partition (((>) `on` view _1) closest) ints
            gts = filter (((<) `on` view _1) closest) geqs
        -- | Exclusive upper bounds
        protoIntervals :: [(Integer, Integer)]
        protoIntervals = [ over _1 (i *) $ base `divMod` toInteger numIterations
                         | i <- [1 .. toInteger numIterations]
                         ]

    intervals' = map (fromBS' <=< assertM' (> 0)) intervals

    fromBS' :: Integer -> Maybe a
    fromBS' = fromBS . pad . ByteString.pack . reverse . unfoldr step
      where step i
              | i <= 0 || i >= base = Nothing
              | otherwise = Just (fromIntegral i, i `shiftR` 8)
            pad bs
              | toInteger (ByteString.length bs) >= bytes = bs
              | otherwise = pad $ ByteString.cons 0 bs

getCurrentInterval :: forall a. Natural -> Natural -> [a] -> a
getCurrentInterval epoch iteration intervals = permIntervals !! fromIntegral (toInteger iteration `mod` genericLength permIntervals)
  where permIntervals = shuffleM intervals `evalRand` mkStdGen (hash epoch)


mkIntervalsCached :: forall m a. (NFData a, MonadIO m)
                  => TVar (Map Natural [(Maybe a, Maybe a)])
                  -> Integer -> (ByteString -> Maybe a) -> Natural -> m [(Maybe a, Maybe a)]
mkIntervalsCached cacheTVar bytes fromBS numIterations = atomically $ do
    cached <- readTVar cacheTVar
    case Map.lookup numIterations cached of
      Just c -> return c
      Nothing -> do
        modifyTVar' cacheTVar $ force . Map.insert numIterations intervals'
        return intervals'
  where intervals' = mkIntervals bytes fromBS numIterations

currentIntervalCached :: forall m a. (NFData a, MonadIO m)
                      => TVar (Map Natural [(Maybe a, Maybe a)])
                      -> Integer -> (ByteString -> Maybe a)
                      -> Natural -> Natural -> Natural -> m (Maybe a, Maybe a)
currentIntervalCached cacheTVar bytes fromBS numIterations epoch iteration
  = getCurrentInterval epoch iteration <$> mkIntervalsCached cacheTVar bytes fromBS numIterations