{-# LANGUAGE OverloadedStrings #-}

{- Copyright 2016 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

module Benchmark where

import Types
import Tunables
import ExpensiveHash
import Cost
import Serialization ()
import qualified Data.ByteString.UTF8 as BU8
import qualified Crypto.Argon2 as Argon2
import Data.Time.Clock
import Control.DeepSeq
import Control.Monad
import Data.Monoid
import Data.Maybe

benchmarkExpensiveHash :: Int -> ExpensiveHashTunable -> IO (BenchmarkResult (Cost CreationOp))
benchmarkExpensiveHash rounds tunables =
	benchmarkExpensiveHash' rounds tunables (getexpected tunables)
  where
	getexpected (UseArgon2 cost _) = mapCost (* fromIntegral rounds) cost

benchmarkExpensiveHash' :: Int -> ExpensiveHashTunable -> Cost op -> IO (BenchmarkResult (Cost op))
benchmarkExpensiveHash' rounds tunables@(UseArgon2 _ hashopts) expected = do
	numcores <- fromIntegral . fromMaybe (error "Unknown number of physical cores.") 
		<$> getNumCores
	start <- getCurrentTime
	forM_ [1..rounds] $ \n -> do
		-- Must vary the data being hashed to avoid laziness
		-- caching hash results.
		let base = BU8.fromString (show n)
		let ExpensiveHash _ t = expensiveHash tunables
			(Salt (GpgKey (KeyId (base <> "dummy"))))
			(base <> "himom")
		t `deepseq` return ()
	end <- getCurrentTime
	let diff = floor $ end `diffUTCTime` start
	let maxthreads = Argon2.hashParallelism hashopts
	let actual = CPUCost (Seconds diff) (Divisibility $ fromIntegral maxthreads)
	-- The expected cost is for a single core, so adjust it
	-- based on the number of cores, up to a maximum of the number
	-- of threads that the hash is configred to use.
	let usedcores = min maxthreads numcores
	let adjustedexpected = mapCost (`div` fromIntegral usedcores) expected
	return $ BenchmarkResult
		{ expectedBenchmark = adjustedexpected
		, actualBenchmark = actual
		}

benchmarkTunables :: Tunables -> IO ()
benchmarkTunables tunables = do
	putStrLn "/proc/cpuinfo:"
	putStrLn =<< readFile "/proc/cpuinfo"
	
	putStrLn $ "Benchmarking 16 rounds of proof of work hash..."
	print =<< benchmarkExpensiveHash 16 (proofOfWorkHashTunable 0)

	-- Rather than run all 256 rounds of this hash, which would
	-- probably take on the order of 1 hour, run only 16, and scale
	-- the expected cost accordingly.
	let normalrounds = fromIntegral $ 
		256 * randomSaltBytes (keyEncryptionKeyTunable tunables)
	putStrLn $ "Benchmarking 16/" ++ show normalrounds ++ " rounds of key encryption key hash..."
	r <- benchmarkExpensiveHash' 16
		(keyEncryptionKeyHash $ keyEncryptionKeyTunable tunables)
		(mapCost (`div` (normalrounds `div` 16)) $ randomSaltBytesBruteForceCost $ keyEncryptionKeyTunable tunables)
	print r
	putStrLn $ "Estimated time for " ++ show normalrounds ++ " rounds of key encryption key hash..."
	print $ BenchmarkResult
		{ expectedBenchmark = mapCost (* 16) (expectedBenchmark r)
		, actualBenchmark = mapCost (* 16) (actualBenchmark r)
		}

	putStrLn "Benchmarking 1 round of name generation hash..."
	print =<< benchmarkExpensiveHash 1
		(nameGenerationHash $ nameGenerationTunable tunables)