{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TypeFamilies #-}

module ResourceQueue (
    ResourceQueue,
    ResourceQueue.take,
    resource,
    request,
    fulfill,
    flush,
    takeRequested,
) where

import Control.Monad.State

import Data.Bifunctor
import Data.Map (Map)
import qualified Data.Map as Map

import Numeric.Natural

import Refined.Unsafe

import Deck.Category
import Util

type ResourceQueue a =
    State (Map CategoryName PositiveNumber, Map CategoryName PositiveNumber) a

flush ::
    ResourceQueue
        ([(CategoryName, PositiveNumber)], [(CategoryName, PositiveNumber)])
flush = do
    (available, requested) <- get

    put (Map.empty, Map.empty)
    pure (Map.toAscList available, Map.toAscList requested)

takeRequested :: ResourceQueue [(CategoryName, PositiveNumber)]
takeRequested = do
    (available, requested) <- get

    put (available, Map.empty)
    pure (Map.toAscList requested)

take :: PositiveNumber -> ResourceQueue [(CategoryName, PositiveNumber)]
take count = do
    (available, requested) <- get

    let (res, remainder) = takeFromQueue count (Map.toAscList available)
    put (Map.fromAscList remainder, requested)
    pure res

resource :: CategoryName -> PositiveNumber -> ResourceQueue ()
resource categoryName count = do
    (available, requested) <- get

    let updateAvailable =
            Map.alter
                ( \case
                    Nothing -> pure count
                    Just x -> pure $ uncheckedAdd x count
                )
                categoryName
                available

    put (updateAvailable, requested)

request :: CategoryName -> PositiveNumber -> ResourceQueue ()
request categoryName count = do
    (available, requested) <- get

    let updateRequested =
            Map.alter
                ( \case
                    Nothing -> pure count
                    Just x -> pure $ uncheckedAdd x count
                )
                categoryName
                requested

    put (available, updateRequested)

fulfill :: PositiveNumber -> ResourceQueue [(CategoryName, PositiveNumber)]
fulfill count = do
    (available, requested) <- get

    let updateRequested = Map.fromAscList remainder
        (res, remainder) = takeFromQueue count $ Map.toAscList requested

    put (available, updateRequested)
    pure res

takeFromQueue ::
    PositiveNumber
    -> [(CategoryName, PositiveNumber)]
    -> ([(CategoryName, PositiveNumber)], [(CategoryName, PositiveNumber)])
takeFromQueue = go mempty . toNumber
  where
    go ::
        [(CategoryName, PositiveNumber)]
        -> Natural
        -> [(CategoryName, PositiveNumber)]
        -> ([(CategoryName, PositiveNumber)], [(CategoryName, PositiveNumber)])
    go acc 0 xs = (acc, xs)
    go acc n (item@(_, count) : xs) = case compare n (toNumber count) of
        LT ->
            go
                (second (const (PositiveNumber $ unsafeRefine n)) item : acc)
                0
                (second (PositiveNumber . unsafeRefine . subtract n . toNumber) item : xs)
        EQ -> go (item : acc) 0 xs
        GT -> go (item : acc) (n - toNumber count) xs
    go acc _ xs = go acc 0 xs