// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Generate cross-platform test vectors for REX encryption compatibility.
//!
//! This script generates deterministic test vectors to verify that
//! the TypeScript and Rust implementations of:
//! - RexValue Borsh serialization
//! - REX encryption constants
//!
//! are compatible across platforms.

use std::fs;

use rialo_tee_secret_sharing::{SECRET_SHARING_HPKE_INFO, USER_SECRET_AAD};
use rialo_types::RexValue;
use serde::Serialize;

// ========== Test Vector Structures ==========

#[derive(Debug, Serialize)]
struct EncryptionTestVectors {
    version: String,
    constants: ConstantsVectors,
    rex_value_borsh: Vec<RexValueVector>,
}

#[derive(Debug, Serialize)]
struct ConstantsVectors {
    user_secret_aad: ConstantVector,
    secret_sharing_hpke_info: ConstantVector,
}

#[derive(Debug, Serialize)]
struct ConstantVector {
    id: String,
    description: String,
    value_string: String,
    value_hex: String,
    byte_length: usize,
}

#[derive(Debug, Serialize)]
struct RexValueVector {
    id: String,
    description: String,
    variant: String,
    /// For Plain: UTF-8 string if applicable, null if binary
    input_string: Option<String>,
    /// Hex-encoded raw bytes
    input_hex: String,
    /// Expected Borsh serialization (hex)
    expected_borsh_hex: String,
}

// ========== Helper Functions ==========

fn bytes_to_hex(bytes: &[u8]) -> String {
    bytes.iter().map(|b| format!("{:02x}", b)).collect()
}

// ========== Vector Generators ==========

fn generate_constants_vectors() -> ConstantsVectors {
    ConstantsVectors {
        user_secret_aad: ConstantVector {
            id: "user_secret_aad".to_string(),
            description: "AAD prefix for user secret encryption (13 bytes)".to_string(),
            value_string: String::from_utf8_lossy(USER_SECRET_AAD).to_string(),
            value_hex: bytes_to_hex(USER_SECRET_AAD),
            byte_length: USER_SECRET_AAD.len(),
        },
        secret_sharing_hpke_info: ConstantVector {
            id: "secret_sharing_hpke_info".to_string(),
            description: "HPKE info string for secret sharing (32 bytes)".to_string(),
            value_string: String::from_utf8_lossy(SECRET_SHARING_HPKE_INFO).to_string(),
            value_hex: bytes_to_hex(SECRET_SHARING_HPKE_INFO),
            byte_length: SECRET_SHARING_HPKE_INFO.len(),
        },
    }
}

fn generate_rex_value_vectors() -> Vec<RexValueVector> {
    let mut vectors = Vec::new();

    // Plain: Empty
    let plain_empty = RexValue::plain(vec![]);
    let borsh_bytes = borsh::to_vec(&plain_empty).unwrap();

    vectors.push(RexValueVector {
        id: "plain_empty".to_string(),
        description: "RexValue::Plain with empty data".to_string(),
        variant: "Plain".to_string(),
        input_string: Some("".to_string()),
        input_hex: "".to_string(),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Plain: Simple ASCII string
    let plain_hello = RexValue::plain_string("hello");
    let borsh_bytes = borsh::to_vec(&plain_hello).unwrap();

    vectors.push(RexValueVector {
        id: "plain_hello".to_string(),
        description: "RexValue::Plain with 'hello' string".to_string(),
        variant: "Plain".to_string(),
        input_string: Some("hello".to_string()),
        input_hex: bytes_to_hex(b"hello"),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Plain: Longer string
    let plain_sentence = RexValue::plain_string("Hello, World!");
    let borsh_bytes = borsh::to_vec(&plain_sentence).unwrap();

    vectors.push(RexValueVector {
        id: "plain_hello_world".to_string(),
        description: "RexValue::Plain with 'Hello, World!' string".to_string(),
        variant: "Plain".to_string(),
        input_string: Some("Hello, World!".to_string()),
        input_hex: bytes_to_hex(b"Hello, World!"),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Plain: Unicode string
    let plain_unicode = RexValue::plain_string("Hello 世界 🌍");
    let borsh_bytes = borsh::to_vec(&plain_unicode).unwrap();

    vectors.push(RexValueVector {
        id: "plain_unicode".to_string(),
        description: "RexValue::Plain with Unicode characters".to_string(),
        variant: "Plain".to_string(),
        input_string: Some("Hello 世界 🌍".to_string()),
        input_hex: bytes_to_hex("Hello 世界 🌍".as_bytes()),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Plain: Binary data (URL example)
    let plain_url = RexValue::plain_string("https://api.example.com/data");
    let borsh_bytes = borsh::to_vec(&plain_url).unwrap();

    vectors.push(RexValueVector {
        id: "plain_url".to_string(),
        description: "RexValue::Plain with URL string".to_string(),
        variant: "Plain".to_string(),
        input_string: Some("https://api.example.com/data".to_string()),
        input_hex: bytes_to_hex(b"https://api.example.com/data"),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Encrypted: Empty ciphertext (edge case)
    let encrypted_empty = RexValue::encrypted(vec![]);
    let borsh_bytes = borsh::to_vec(&encrypted_empty).unwrap();

    vectors.push(RexValueVector {
        id: "encrypted_empty".to_string(),
        description: "RexValue::Encrypted with empty ciphertext".to_string(),
        variant: "Encrypted".to_string(),
        input_string: None,
        input_hex: "".to_string(),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Encrypted: Realistic ciphertext size (48 bytes = 32 enc + 16 tag, no plaintext)
    let mock_ciphertext = vec![0xab; 48];
    let encrypted_48 = RexValue::encrypted(mock_ciphertext.clone());
    let borsh_bytes = borsh::to_vec(&encrypted_48).unwrap();

    vectors.push(RexValueVector {
        id: "encrypted_48_bytes".to_string(),
        description: "RexValue::Encrypted with 48-byte mock ciphertext (min overhead)".to_string(),
        variant: "Encrypted".to_string(),
        input_string: None,
        input_hex: bytes_to_hex(&mock_ciphertext),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Encrypted: Larger ciphertext (100 bytes)
    let mock_ciphertext_100: Vec<u8> = (0..100).map(|i| (i % 256) as u8).collect();
    let encrypted_100 = RexValue::encrypted(mock_ciphertext_100.clone());
    let borsh_bytes = borsh::to_vec(&encrypted_100).unwrap();

    vectors.push(RexValueVector {
        id: "encrypted_100_bytes".to_string(),
        description: "RexValue::Encrypted with 100-byte ciphertext".to_string(),
        variant: "Encrypted".to_string(),
        input_string: None,
        input_hex: bytes_to_hex(&mock_ciphertext_100),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    // Plain: Binary data (non-UTF8)
    let plain_binary = RexValue::plain(vec![0x00, 0x01, 0xff, 0xfe, 0x80]);
    let borsh_bytes = borsh::to_vec(&plain_binary).unwrap();

    vectors.push(RexValueVector {
        id: "plain_binary".to_string(),
        description: "RexValue::Plain with non-UTF8 binary data".to_string(),
        variant: "Plain".to_string(),
        input_string: None,
        input_hex: bytes_to_hex(&[0x00, 0x01, 0xff, 0xfe, 0x80]),
        expected_borsh_hex: bytes_to_hex(&borsh_bytes),
    });

    vectors
}

// ========== Public API ==========

pub fn generate() {
    eprintln!("Generating cross-platform encryption test vectors");

    let test_vectors = EncryptionTestVectors {
        version: "1.0.0".to_string(),
        constants: generate_constants_vectors(),
        rex_value_borsh: generate_rex_value_vectors(),
    };

    // Serialize to JSON with pretty printing
    let json =
        serde_json::to_string_pretty(&test_vectors).expect("Failed to serialize test vectors");

    // Write to file
    let manifest_dir = env!("CARGO_MANIFEST_DIR");
    let output_path = format!("{manifest_dir}/test-vectors/encryption-vectors.json");
    fs::write(&output_path, json).expect("Failed to write test vectors");

    eprintln!("Encryption test vectors generated successfully");
    eprintln!("Output: {output_path}");
    eprintln!(
        "Generated {} RexValue vectors",
        test_vectors.rex_value_borsh.len()
    );
    eprintln!(
        "Constants verified: USER_SECRET_AAD ({} bytes), SECRET_SHARING_HPKE_INFO ({} bytes)",
        test_vectors.constants.user_secret_aad.byte_length,
        test_vectors.constants.secret_sharing_hpke_info.byte_length
    );
}
