module Main where

import Codec.FEC (FECParams (paramK, paramN), decode, encode, fec, initialize)
import Control.Monad (replicateM)
import Criterion.Main (Benchmark, bench, bgroup, defaultMain, env, nf)
import Data.Bifunctor (bimap)
import qualified Data.ByteString as B
import Data.List (unfoldr)
import System.Random (genByteString, mkStdGen)

main :: IO ()
main =
    defaultMain
        -- Run against some somewhat arbitrarily chosen configurations.  Notably,
        -- though, 94/100 matches the numbers recorded in the readme.
        [ env (setupFEC 2 3) makeFECBenchmarks
        , env (setupFEC 16 31) makeFECBenchmarks
        , env (setupFEC 94 100) makeFECBenchmarks
        ]
  where
    setupFEC :: Int -> Int -> IO FECParams
    setupFEC k n = do
        initialize
        pure (fec k n)

    makeFECBenchmarks = fecGroup [10 ^ 6]

    fecGroup sizes params =
        bgroup
            (show (paramK params) <> "/" <> show (paramN params))
            ( []
                ++ (decodePrimaryBenchmark params <$> sizes)
                ++ (decodeSecondaryBenchmark params <$> sizes)
                ++ (encodeBenchmark params <$> sizes)
            )

    encodeBenchmark params size =
        env (setupBlocks (paramK params) size) $
            benchmarkEncode params
    decodePrimaryBenchmark params size =
        env (setupBlocks (paramK params) size) $
            benchmarkPrimaryDecode params
    decodeSecondaryBenchmark params size =
        env (setupBlocks (paramK params) size) $
            benchmarkSecondaryDecode params

    setupBlocks :: Int -> Int -> IO [B.ByteString]
    setupBlocks k blockSize = pure $ makeBlocks k blockSize

    benchmarkEncode params blocks =
        bench ("encode blockSize=" <> showWithUnit (B.length $ head blocks)) $
            -- We choose normal form here because the typical thing to do with the
            -- result is serialize use all of the bytes (eg, to write them to a
            -- file or send them over the network) so they will certainly all be
            -- used.
            nf (uncurry encode) (params, blocks)

    benchmarkPrimaryDecode params blocks =
        bench ("decode [0..] blockSize=" <> showWithUnit (B.length $ head blocks)) $
            -- normal form here for the same reason as in benchmarkEncode.
            -- assign block numbers to use only primary blocks
            nf (uncurry decode) (params, (zip [0 ..] blocks))

    benchmarkSecondaryDecode params blocks =
        bench ("decode [" <> show n <> "..] blockSize=" <> showWithUnit (B.length $ head blocks)) $
            -- normal form here for the same reason as in benchmarkEncode.
            -- assign block numbers to use as many non-primary blocks as
            -- possible
            nf (uncurry decode) (params, (zip [n ..] blocks))
      where
        n = paramN params - paramK params

makeBlocks :: Int -> Int -> [B.ByteString]
makeBlocks k size = take k . go $ mkStdGen 42
  where
    go = uncurry ($) . bimap (:) go . genByteString size

data BytesUnit = B | KB | MB deriving (Eq, Ord, Enum, Show, Bounded)

bestUnit :: Int -> BytesUnit
bestUnit n
    | n < 1000 = minBound
    | maxBound == nextUnit = nextUnit
    | otherwise = succ nextUnit
  where
    nextUnit = bestUnit . (`div` 1000) $ n

showWithUnit :: Int -> String
showWithUnit n = show (scale n) <> show u
  where
    scale n = n `div` (10 ^ (3 * fromEnum u))

    u = bestUnit n