{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}


module Main where

import qualified Data.ByteString               as B
import qualified Data.ByteString.Lazy          as BL
import qualified Proto3.Wire.Decode as De
import qualified Proto3.Wire.Encode as En
import Proto3.Wire

import Control.Applicative (liftA2, liftA3)
import Control.Monad (forM)
import Data.Maybe
import Data.Word
import Data.IORef

import Criterion (bench)
import qualified Criterion as C
import Criterion.Main (defaultMain)

data Tree a = Leaf | Branch a (Tree a) (Tree a)
  deriving (Eq, Functor)

data Rose a = Bud | Rose a [Rose a]

instance Foldable Tree where
  foldr _ z Leaf = z
  foldr f z (Branch a t1 t2) = foldr f (f a (foldr f z t2)) t1

  sum Leaf = 0
  sum (Branch a t1 t2) =
    let !a1 = sum t1
        !a2 = sum t2
    in a + a1 + a2

instance Foldable Rose where
  foldMap f Bud = mempty
  foldMap f (Rose x rs) = f x <> ((foldMap.foldMap) f rs)

intTreeParser :: De.Parser De.RawMessage (Tree Word64)
intTreeParser = liftA3 combine
    (De.at (De.repeated De.fixed64) (FieldNumber 0))
    (De.at (De.one (De.embedded' intTreeParser) Leaf) (FieldNumber 1))
    (De.at (De.one (De.embedded' intTreeParser) Leaf) (FieldNumber 2))
  where
    combine xs y z = Branch (sum xs) y z

intRoseParser :: De.Parser De.RawMessage (Rose Word64)
intRoseParser = liftA2 (Rose @Word64)
  (De.at (De.one De.fixed64 0) (FieldNumber 0))
  (De.at (De.repeated (De.embedded' intRoseParser)) (FieldNumber 1))

detRandom :: [Word64]
detRandom = concat . replicate 10 $
  [ 227, 133, 16, 164, 43,
    159, 207, 87, 180, 236,
    245, 128, 249, 170, 216,
    181, 164, 162, 239, 249,
    76, 237, 197, 246, 209,
    231, 124, 154, 55, 64,
    4, 114, 79, 199, 252,
    163, 116, 237, 209, 138,
    240, 148, 212, 224, 88,
    131, 122, 114, 158, 97,
    186, 3, 223, 230, 223,
    207, 93, 168, 48, 130,
    77, 122, 30, 222, 221,
    224, 243, 19, 175, 61,
    112, 246, 201, 57, 185,
    19, 128, 129, 138, 209,
    4, 153, 196, 238, 72,
    254, 157, 233, 81, 30,
    106, 249, 57, 214, 104,
    171, 146, 175, 185, 192,
    159, 207, 87, 180, 236,
    227, 133, 16, 164, 43,
    245, 128, 249, 170, 216,
    181, 164, 162, 239, 249,
    76, 237, 197, 246, 209,
    231, 124, 154, 55, 64,
    4, 114, 79, 199, 252,
    163, 116, 237, 209, 138,
    240, 148, 212, 224, 88,
    131, 122, 114, 158, 97,
    186, 3, 223, 230, 223,
    207, 93, 168, 48, 130,
    77, 122, 30, 222, 221,
    224, 243, 19, 175, 61,
    112, 246, 201, 57, 185,
    19, 128, 129, 138, 209,
    4, 153, 196, 238, 72,
    254, 157, 233, 81, 30,
    106, 249, 57, 214, 104,
    171, 146, 175, 185, 192,
    159, 207, 87, 180, 236,
    227, 133, 16, 164, 43,
    245, 128, 249, 170, 216,
    181, 164, 162, 239, 249,
    76, 237, 197, 246, 209,
    231, 124, 154, 55, 64,
    4, 114, 79, 199, 252,
    163, 116, 237, 209, 138,
    240, 148, 212, 224, 88,
    131, 122, 114, 158, 97,
    186, 3, 223, 230, 223,
    207, 93, 168, 48, 130,
    77, 122, 30, 222, 221,
    224, 243, 19, 175, 61,
    112, 246, 201, 57, 185,
    19, 128, 129, 138, 209,
    4, 153, 196, 238, 72,
    254, 157, 233, 81, 30,
    106, 249, 57, 214, 104,
    171, 146, 175, 185, 192
    ]

pullInt :: IORef [Word64] -> IO Word64
pullInt xs = do
  xs' <- readIORef xs
  case xs' of
    [] -> pure (-1)
    x : xs' -> do
      writeIORef xs xs'
      pure x

mkTree0 :: IO Word64 -> IO En.MessageBuilder
mkTree0 ints = do
  shouldFork <- (\(i :: Word64) -> (i `mod` 8) < 6) <$> ints
  if shouldFork
    then do
      i <- En.fixed64 (FieldNumber 0) <$> ints
      left <- En.embedded (FieldNumber 1) <$> mkTree0 ints
      right <- En.embedded (FieldNumber 2) <$> mkTree0 ints
      pure (i <> left <> right)
    else pure mempty

mkRose0 :: IO Word64 -> IO En.MessageBuilder
mkRose0 ints = do
  next <- fromIntegral <$> ints
  if next == -1 then pure mempty else do
    let nBranches = next `mod` 9
    if nBranches == 0 then pure mempty else do
      loc <- (\i -> (i `mod` nBranches)) . fromIntegral <$> ints
      i <- En.fixed64 (FieldNumber 0) <$> ints
      rs1 <- forM (replicate loc ()) $ \() ->
        En.embedded (FieldNumber 1) <$> mkTree0 ints
      rs2 <- forM (replicate (nBranches - loc) ()) $ \() ->
        En.embedded (FieldNumber 1) <$> mkTree0 ints
      pure (mconcat rs1 <> i <> mconcat rs2)

mkTree :: IO B.ByteString
mkTree = BL.toStrict . En.toLazyByteString <$> (mkTree0 . pullInt =<< newIORef detRandom)

mkRose :: IO B.ByteString
mkRose = BL.toStrict . En.toLazyByteString <$> (mkRose0 . pullInt =<< newIORef detRandom)

decode :: Foldable f => De.Parser De.RawMessage (f Word64) -> B.ByteString -> IO (Maybe Word64)
decode p = pure . fmap sum . toMaybe . De.parse p
  where
    toMaybe (Left _) = Nothing
    toMaybe (Right x) = Just x

unwrap :: (Functor m, Foldable f) => m (f a) -> m a
unwrap = fmap (foldr1 const)

main :: IO ()
main =
  defaultMain
    [ bench "Parse int tree" $ C.perRunEnv mkTree (unwrap . decode intTreeParser)
    , bench "Parse int rose tree" $ C.perRunEnv mkRose (unwrap . decode intRoseParser)]