{-# LANGUAGE OverloadedStrings, FlexibleInstances #-}

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

module Benchmark where

import Types
import Output
import Tunables
import ExpensiveHash
import HTTP.ProofOfWork
import Cost
import Serialization ()
import qualified Data.ByteString.UTF8 as BU8
import qualified Data.Text as T
import qualified Crypto.Argon2 as Argon2
import Data.Time.Clock
import Control.DeepSeq
import Control.Monad
import Data.Maybe

data BenchmarkResult t = BenchmarkResult { expectedBenchmark :: t, actualBenchmark :: t }

instance Show (BenchmarkResult (Cost op)) where
	show br = "   expected: " ++ show (expectedBenchmark br) ++ "s"
		++ "\tactual: " ++ show (actualBenchmark br) ++ "s"

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 = T.pack (show n)
		let baseb = BU8.fromString (show n)
		let ExpensiveHash _ t = expensiveHash tunables
			(Salt (GpgKey (KeyId (base <> "dummy"))))
			(baseb <> "himom")
		t `deepseq` return ()
	end <- getCurrentTime
	let diff = floor (end `diffUTCTime` start) :: Integer
	let maxthreads = Argon2.hashParallelism hashopts
	let actual = CPUCost (Seconds (fromIntegral 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 (/ fromIntegral usedcores) expected
	return $ BenchmarkResult
		{ expectedBenchmark = adjustedexpected
		, actualBenchmark = actual
		}

benchmark :: NFData t => Int -> Cost CreationOp -> (Int -> IO t) -> IO (BenchmarkResult (Cost CreationOp))
benchmark rounds expected a = do
	start <- getCurrentTime
	forM_ [1..rounds] $ \n -> do
		v <- a n
		v `deepseq` return ()
	end <- getCurrentTime
	let diff = floor (end `diffUTCTime` start) :: Integer
	return $ BenchmarkResult
		{ expectedBenchmark = expected
		, actualBenchmark = CPUCost (Seconds (fromIntegral diff)) (Divisibility 1)
		}

benchmarkPoW :: Int  -> Seconds -> IO (BenchmarkResult (Cost CreationOp))
benchmarkPoW rounds seconds = do
	let Just mk = mkProofOfWorkRequirement seconds
	s <- newRequestIDSecret
	rid <- mkRequestID s
	benchmark rounds (CPUCost (seconds * fromIntegral rounds) (Divisibility 1))
		(return . genProofOfWork (mk rid))

benchmarkTunables :: Tunables -> IO ()
benchmarkTunables tunables = do
	say "/proc/cpuinfo:"
	say =<< readFile "/proc/cpuinfo"
	
	say "Benchmarking 1000 rounds of proof of work hash..."
	display =<< benchmarkExpensiveHash 1000 (proofOfWorkHashTunable 0)

	say "Benchmarking 60 rounds of 1 second proofs of work..."
	display =<< benchmarkPoW 60 (Seconds 1)
	
	say "Benchmarking 10 rounds of 8 second proofs of work..."
	display =<< benchmarkPoW 10 (Seconds 8)

	-- 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 = 256 * randomSaltBytes (keyEncryptionKeyTunable tunables)
	say $ "Benchmarking 16/" ++ show normalrounds ++ " rounds of key encryption key hash..."
	r <- benchmarkExpensiveHash' 16
		(keyEncryptionKeyHash $ keyEncryptionKeyTunable tunables)
		(mapCost (/ (fromIntegral normalrounds / 16)) $ randomSaltBytesBruteForceCost $ keyEncryptionKeyTunable tunables)
	display r
	say $ "Estimated time for " ++ show normalrounds ++ " rounds of key encryption key hash..."
	display $ BenchmarkResult
		{ expectedBenchmark = mapCost (* 16) (expectedBenchmark r)
		, actualBenchmark = mapCost (* 16) (actualBenchmark r)
		}

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