// Copyright (c) 2025 blackshirt.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
//
// This file contains benchmarking code for standard AEAD_CHACHA20_POLY1305 encryption
// and decryption compared to AEAD_CHACHA20_POLY1305 with PSIV construct.
//
// This output on my test.
// Standard ChaCha20Poly1305 AEAD and PSIV construct output performance comparison
// ===============================================================================
// Iterations per test: 1000
// ------------------------------------------------------------------------------------------------------
//   Data Size |        Std |       PSIV | Enc (std/psiv) ||        Std |         PSIV | Dec (std/psiv)|
// ------------------------------------------------------------------------------------------------------
//         6 B |    16.00ms |    19.00ms |          0.85x ||    16.00ms |      19.00ms |         0.81x |
//         8 B |    15.00ms |    19.00ms |          0.81x ||    16.00ms |      18.00ms |         0.87x |
//        12 B |    15.00ms |    20.00ms |          0.76x ||    16.00ms |      20.00ms |         0.77x |
//        16 B |    15.00ms |    21.00ms |          0.72x ||    15.00ms |      19.00ms |         0.80x |
//        64 B |    21.00ms |    25.00ms |          0.82x ||    21.00ms |      26.00ms |         0.80x |
//        75 B |    28.00ms |    35.00ms |          0.81x ||    34.00ms |      44.00ms |         0.77x |
//       256 B |    55.00ms |    59.00ms |          0.94x ||    50.00ms |      63.00ms |         0.80x |
//      1028 B |   174.00ms |   242.00ms |          0.72x ||   178.00ms |     227.00ms |         0.78x |
//      2049 B |   361.00ms |   522.00ms |          0.69x ||   399.00ms |     558.00ms |         0.71x |
// ------------------------------------------------------------------------------------------------------
//       Total |   703.00ms |   965.00ms |          0.73x ||   748.00ms |        1.00s |          0.75x|
// ------------------------------------------------------------------------------------------------------
//
// Per-operation averages:
//   Standard ChaCha20Poly1305 encrypt:        78209 ns per hash
//   ChaCha20Poly1305 PSIV encrypt:           107325 ns per hash
//
//   Standard ChaCha20Poly1305 decrypt:        83151 ns per hash
//   ChaCha20Poly1305 PSIV decrypt:           111155 ns per hash
module main

import rand
import time
import x.crypto.chacha20poly1305

const benchmark_iterations = 1000

const test_data_sizes = [6, 8, 12, 16, 64, 75, 256, 1028, 2049]

// standard AEAD_CHACHA20_POLY1305 encryption
fn benchmark_aead_std_encrypt(msg []u8, key []u8, nonce []u8, ad []u8, iterations int) time.Duration {
    start := time.now()
    for _ in 0 .. iterations {
        _ := chacha20poly1305.encrypt(msg, key, nonce, ad) or { panic(err) }
    }
    return time.since(start)
}

// standard AEAD_CHACHA20_POLY1305 decryption
fn benchmark_aead_std_decrypt(data []u8, key []u8, nonce []u8, ad []u8, iterations int) time.Duration {
    start := time.now()
    for _ in 0 .. iterations {
        _ := chacha20poly1305.decrypt(data, key, nonce, ad) or { panic(err) }
    }
    return time.since(start)
}

// psiv AEAD_CHACHA20_POLY1305 encryption
fn benchmark_aead_psiv_encrypt(msg []u8, key []u8, nonce []u8, ad []u8, iterations int) time.Duration {
    start := time.now()
    for _ in 0 .. iterations {
        _ := chacha20poly1305.psiv_encrypt(msg, key, nonce, ad) or { panic(err) }
    }
    return time.since(start)
}

// psiv AEAD_CHACHA20_POLY1305 decryption
fn benchmark_aead_psiv_decrypt(data []u8, key []u8, nonce []u8, ad []u8, iterations int) time.Duration {
    start := time.now()
    for _ in 0 .. iterations {
        _ := chacha20poly1305.psiv_decrypt(data, key, nonce, ad) or { panic(err) }
    }
    return time.since(start)
}

fn format_duration(d time.Duration) string {
    if d.microseconds() < 1000 {
        return '${d.microseconds():6}μs'
    } else if d.milliseconds() < 1000 {
        return '${f64(d.milliseconds()):6.2f}ms'
    } else {
        return '${f64(d.seconds()):6.2f}s'
    }
}

const data_title = 'Data Size'
const aead_std_enc = 'Std'
const aead_psiv_enc = 'PSIV'
const aead_std_dec = 'Std'
const aead_psiv_dec = 'PSIV'
const ratio_std_psiv_enc_title = 'Enc (std/psiv)'
const ratio_std_psiv_dec_title = 'Dec (std/psiv)'

fn main() {
    println('Standard ChaCha20Poly1305 AEAD and PSIV construct output performance comparison')
    println('===============================================================================')
    println('Iterations per test: ${benchmark_iterations}')

    println('${'-'.repeat(102)}')
    println('${data_title:12} | ${aead_std_enc:10} | ${aead_psiv_enc:10} | ${ratio_std_psiv_enc_title:12} || ${aead_std_dec:10} | ${aead_psiv_dec:12} | ${ratio_std_psiv_dec_title:12}|')
    println('${'-'.repeat(102)}')

    mut total_std_encrypt := time.Duration(0)
    mut total_std_decrypt := time.Duration(0)
    mut total_psiv_encrypt := time.Duration(0)
    mut total_psiv_decrypt := time.Duration(0)

    key := rand.bytes(32)!
    nonce := rand.bytes(12)!
    for size in test_data_sizes {
        ad := rand.bytes(size)!
        test_msg := rand.bytes(size)!

        // Warm up
        out0 := chacha20poly1305.encrypt(test_msg, key, nonce, ad)!
        _ := chacha20poly1305.decrypt(out0, key, nonce, ad)!

        out1 := chacha20poly1305.psiv_encrypt(test_msg, key, nonce, ad)!
        _ := chacha20poly1305.psiv_decrypt(out1, key, nonce, ad)!

        // Benchmark Standard AEAD_CHACHA20_POLY1305 encryption
        std_encrypt_time := benchmark_aead_std_encrypt(test_msg, key, nonce, ad,
            benchmark_iterations)

        // Benchmark Standard AEAD_CHACHA20_POLY1305 decryption
        std_decrypt_time := benchmark_aead_std_decrypt(out0, key, nonce, ad, benchmark_iterations)

        // Benchmark AEAD_CHACHA20_POLY1305 PSIV encryption
        psiv_encrypt_time := benchmark_aead_psiv_encrypt(test_msg, key, nonce, ad,
            benchmark_iterations)

        // Benchmark AEAD_CHACHA20_POLY1305 PSIV decryption
        psiv_decrypt_time := benchmark_aead_psiv_decrypt(out1, key, nonce, ad, benchmark_iterations)

        // Calculate ratio Standard/PSIV encryption
        ratio_std_psiv_encrypt := f64(std_encrypt_time.nanoseconds()) / f64(psiv_encrypt_time.nanoseconds())

        // Calculate ratio Standard/PSIV decryption
        ratio_std_psiv_decrypt := f64(std_decrypt_time.nanoseconds()) / f64(psiv_decrypt_time.nanoseconds())

        stdencrypt_str := format_duration(std_encrypt_time)
        stddecrypt_str := format_duration(std_decrypt_time)
        psivencrypt_str := format_duration(psiv_encrypt_time)
        psivdecrypt_str := format_duration(psiv_decrypt_time)

        ratio_std_psiv_encrypt_str := '${ratio_std_psiv_encrypt:6.2f}x'
        ratio_std_psiv_decrypt_str := '${ratio_std_psiv_decrypt:6.2f}x'

        println('${size:10} B | ${stdencrypt_str:10} | ${psivencrypt_str:10} | ${ratio_std_psiv_encrypt_str:14} || ${stddecrypt_str:10} |  ${psivdecrypt_str:11} |  ${ratio_std_psiv_decrypt_str:12} |')

        total_std_encrypt += std_encrypt_time
        total_std_decrypt += std_decrypt_time
        total_psiv_encrypt += psiv_encrypt_time
        total_psiv_decrypt += psiv_decrypt_time
    }

    println('${'-'.repeat(102)}')

    // Overall performance comparison
    overall_std_psiv_encrypt_ratio := f64(total_std_encrypt.nanoseconds()) / f64(total_psiv_encrypt.nanoseconds())
    overall_std_psiv_decrypt_ratio := f64(total_std_decrypt.nanoseconds()) / f64(total_psiv_decrypt.nanoseconds())

    total_title := 'Total'
    println('${total_title:12} | ${format_duration(total_std_encrypt):10} | ${format_duration(total_psiv_encrypt):10} | ${overall_std_psiv_encrypt_ratio:13.2f}x || ${format_duration(total_std_decrypt):10} | ${format_duration(total_psiv_decrypt):12} |  ${overall_std_psiv_decrypt_ratio:12.2f}x|')
    println('${'-'.repeat(102)}')

    println('')
    println('Per-operation averages:')
    avg_std_encrypt := total_std_encrypt.nanoseconds() / (benchmark_iterations * test_data_sizes.len)
    avg_std_decrypt := total_std_decrypt.nanoseconds() / (benchmark_iterations * test_data_sizes.len)
    avg_psiv_encrypt := total_psiv_encrypt.nanoseconds() / (benchmark_iterations * test_data_sizes.len)
    avg_psiv_decrypt := total_psiv_decrypt.nanoseconds() / (benchmark_iterations * test_data_sizes.len)

    println('  Standard ChaCha20Poly1305 encrypt:\t ${avg_std_encrypt:8} ns per hash')
    println('  ChaCha20Poly1305 PSIV encrypt:\t ${avg_psiv_encrypt:8} ns per hash')
    println('')
    println('  Standard ChaCha20Poly1305 decrypt:\t ${avg_std_decrypt:8} ns per hash')
    println('  ChaCha20Poly1305 PSIV decrypt:\t ${avg_psiv_decrypt:8} ns per hash')
    println('')
}