use criterion::{black_box, criterion_group, criterion_main, Criterion};
use kindi::pattern::PatternDb;

/// Minimal but realistic keyword database for benchmarking.
const KEYWORD_DB: &str = r#"
[keyword_list_version]
    1

[{"evidence_type": "algorithm/asymmetric/RSA", "language": "all"}]
    "\bRSA_"
    "RSAPrivateKey"
    "RSAPublicKey"

[{"evidence_type": "algorithm/hash/SHA256", "language": "all"}]
    "SHA256"
    "SHA-256"
    "\bsha256\b"

[{"evidence_type": "algorithm/symmetric/AES", "language": "all"}]
    "\bAES_encrypt\b"
    "\bAES_decrypt\b"
    "\bAES_set_encrypt_key\b"
    "\bAES_cbc_encrypt\b"

[{"evidence_type": "algorithm/hash/MD5", "language": "all"}]
    "\bMD5_Init\b"
    "\bMD5_Update\b"
    "\bMD5_Final\b"
    "MD5_DIGEST_LENGTH"

[{"evidence_type": "protocol/TLS", "language": "all"}]
    "SSL_CTX_new"
    "SSL_connect"
    "TLS_method"
    "SSL_library_init"

[{"evidence_type": "algorithm/symmetric/DES", "language": "all"}]
    "\bDES_"
    "_DES\b"
    "DES_ecb_encrypt"
    "DES_set_key"
"#;

/// Synthetic source file with realistic crypto patterns for benchmarking.
const SAMPLE_FILE: &str = r#"
#include <openssl/sha.h>
#include <openssl/aes.h>
#include <openssl/ssl.h>
#include <openssl/md5.h>
#include <openssl/des.h>
#include <openssl/rsa.h>

// Initialize the SSL library
void init_crypto() {
    SSL_library_init();
    SSL_CTX *ctx = SSL_CTX_new(TLS_method());
    if (!ctx) {
        fprintf(stderr, "Failed to create SSL context\n");
        return;
    }

    // Generate RSA key pair
    RSA *rsa = RSA_new();
    RSAPublicKey *pub_key = NULL;
    RSAPrivateKey *priv_key = NULL;

    // Hash some data with SHA256
    SHA256_CTX sha_ctx;
    unsigned char hash[SHA256_DIGEST_LENGTH];
    SHA256_Init(&sha_ctx);
    SHA256_Update(&sha_ctx, data, len);
    SHA256_Final(hash, &sha_ctx);

    // Also compute MD5 for legacy compatibility
    MD5_CTX md5_ctx;
    unsigned char md5_hash[MD5_DIGEST_LENGTH];
    MD5_Init(&md5_ctx);
    MD5_Update(&md5_ctx, data, len);
    MD5_Final(md5_hash, &md5_ctx);

    // AES encryption
    AES_KEY aes_key;
    AES_set_encrypt_key(key, 256, &aes_key);
    AES_cbc_encrypt(in, out, len, &aes_key, iv, AES_ENCRYPT);
    AES_encrypt(block, encrypted, &aes_key);
    AES_decrypt(encrypted, decrypted, &aes_key);

    // Legacy DES (should be flagged)
    DES_cblock des_key;
    DES_key_schedule schedule;
    DES_set_key(&des_key, &schedule);
    DES_ecb_encrypt(&input, &output, &schedule, DES_ENCRYPT);
}

// Connection handler
int handle_connection(int sock) {
    SSL *ssl = SSL_new(ctx);
    SSL_set_fd(ssl, sock);
    if (SSL_connect(ssl) <= 0) {
        return -1;
    }

    // Verify peer certificate
    X509 *cert = SSL_get_peer_certificate(ssl);
    RSA *peer_rsa = EVP_PKEY_get1_RSA(EVP_PKEY_get0(X509_get_pubkey(cert)));

    // Generate session key using SHA-256
    unsigned char session_key[32];
    SHA256(random_bytes, sizeof(random_bytes), session_key);

    return 0;
}
"#;

fn bench_keyword_search(c: &mut Criterion) {
    let db = PatternDb::parse(KEYWORD_DB, false).unwrap();

    c.bench_function("keyword_search", |b| {
        b.iter(|| {
            black_box(db.search(SAMPLE_FILE, "all"));
        });
    });
}

fn bench_keyword_load(c: &mut Criterion) {
    c.bench_function("keyword_db_load", |b| {
        b.iter(|| {
            black_box(PatternDb::parse(KEYWORD_DB, false).unwrap());
        });
    });
}

fn bench_has_match(c: &mut Criterion) {
    let db = PatternDb::parse(KEYWORD_DB, false).unwrap();

    c.bench_function("has_match_quick_check", |b| {
        b.iter(|| {
            black_box(db.has_match(SAMPLE_FILE.as_bytes(), "all"));
        });
    });
}

fn bench_case_insensitive(c: &mut Criterion) {
    let db = PatternDb::parse(KEYWORD_DB, true).unwrap();

    c.bench_function("keyword_search_case_insensitive", |b| {
        b.iter(|| {
            black_box(db.search(SAMPLE_FILE, "all"));
        });
    });
}

fn bench_whole_word(c: &mut Criterion) {
    let db = PatternDb::parse_with_options(KEYWORD_DB, false, true, &[]).unwrap();

    c.bench_function("keyword_search_whole_word", |b| {
        b.iter(|| {
            black_box(db.search(SAMPLE_FILE, "all"));
        });
    });
}

criterion_group!(
    benches,
    bench_keyword_search,
    bench_keyword_load,
    bench_has_match,
    bench_case_insensitive,
    bench_whole_word,
);
criterion_main!(benches);
