-- mark.hs: hOpenPGP benchmark suite
-- Copyright © 2014-2026  Clint Adams
-- This software is released under the terms of the Expat license.
-- (See the LICENSE file).
{-# LANGUAGE FlexibleContexts #-}

import Criterion.Main
import Data.Binary (get)
import qualified Data.ByteString.Lazy as BL
import qualified Data.Conduit as DC
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import Data.Conduit.Serialization.Binary (conduitGet)
import Data.Either (rights)
import qualified Data.IxSet.Typed as IxSet
import Data.Maybe (catMaybes)

import Codec.Encryption.OpenPGP.Policy
    ( defaultVerificationPolicy
    )
import Codec.Encryption.OpenPGP.Serialize
    ( conduitParsePktsWithWireRep
    , parsePkts
    , parsePktsEither
    , parsePktsWithWireRep
    )
import Codec.Encryption.OpenPGP.Signatures
    ( verifyAgainstKeyring
    , verifyAgainstKeys
    , verifySigWith
    , verifyTKWith
    , verifyUnknownTKWith
    )
import Codec.Encryption.OpenPGP.Types
    ( someTKToPublicViewTK
    , wireRepRef
    )
import Data.Conduit.OpenPGP.Keyring
    ( conduitToSomeTKsEither
    , conduitToTKsEither
    )

main :: IO ()
main =
    defaultMain
        [ bgroup
            "keyring"
            [ bench "load keys" $ whnfIO (loadKeys "tests/data/pubring.gpg")
            , bench "load keyring" $
                whnfIO (loadKeyring "tests/data/pubring.gpg")
            , bench "self-verify keys" $
                whnfIO (selfVerifyKeys "tests/data/pubring.gpg")
            , bench "self-verify keyring" $
                whnfIO (selfVerifyKeyring "tests/data/pubring.gpg")
            ]
        , env (BL.readFile "tests/data/pubring.gpg") $ \pubringPayload ->
            let pubringRef = wireRepRef pubringPayload
             in bgroup
                    "packet-parse"
                    [ bench "parsePkts/count" $ nf (length . parsePkts) pubringPayload
                    , bench "parsePktsEither/count" $
                        nf
                            (either (const 0) length . parsePktsEither)
                            pubringPayload
                    , bench "parsePktsWithWireRep/count" $
                        nf
                            (length . parsePktsWithWireRep pubringRef)
                            pubringPayload
                    , bench "conduitParsePktsWithWireRep/count" $
                        whnfIO
                            ( fmap
                                length
                                ( DC.runConduitRes $
                                    CB.sourceLbs pubringPayload
                                        DC..| conduitParsePktsWithWireRep Nothing
                                        DC..| CL.consume
                                )
                            )
                    ]
        ]
  where
    loadKeys fp =
        fmap
            (catMaybes . rights)
            ( DC.runConduitRes $
                CB.sourceFile fp
                    DC..| conduitGet get
                    DC..| conduitToTKsEither
                    DC..| CL.consume
            )
    loadKeyring fp =
        fmap
            (sinkFromSomeTKs . rights)
            ( DC.runConduitRes $
                CB.sourceFile fp
                    DC..| conduitGet get
                    DC..| conduitToSomeTKsEither
                    DC..| CL.consume
            )
    selfVerifyKeys fp =
        fmap
            ( \ks ->
                mapM
                    ( verifyUnknownTKWith
                        (verifySigWith defaultVerificationPolicy (verifyAgainstKeys ks))
                        Nothing
                    )
                    ks
            )
            (loadKeys fp)
    selfVerifyKeyring fp =
        fmap
            ( \kr ->
                mapM
                    ( verifyTKWith
                        (verifySigWith defaultVerificationPolicy (verifyAgainstKeyring kr))
                        Nothing
                    )
                    (IxSet.toList kr)
            )
            (loadKeyring fp)
    sinkFromSomeTKs =
        IxSet.fromList
            . map someTKToPublicViewTK
            . catMaybes
