const std = @import("std");
const testing = std.testing;

const byte_operations = @import("utility").byte_operations;
const word_to_bytes = byte_operations.word_to_bytes_le;
const bytes_to_word = byte_operations.bytes_to_word_le;

const chacha20 = @import("primitive").streamcipher.chacha20;

// ----------------------------------- TEST VECTORS -----------------------------------  //

// https://www.rfc-editor.org/rfc/rfc7539#section-2.2.1
test "ChaCha Quarter Round" {
    var state = [chacha20.BLOCK_WORDS]u32{
        0x879531e0, 0xc5ecf37d, 0x516461b1, 0xc9a62f8a,
        0x44c20ef3, 0x3390af7f, 0xd9fc690b, 0x2a5f714c,
        0x53372767, 0xb00a5631, 0x974c541a, 0x359e9963,
        0x5c971061, 0x3d631689, 0x2098d9d6, 0x91dbd320,
    };
    const reference = [chacha20.BLOCK_WORDS]u32{
        0x879531e0, 0xc5ecf37d, 0xbdb886dc, 0xc9a62f8a,
        0x44c20ef3, 0x3390af7f, 0xd9fc690b, 0xcfacafd2,
        0xe46bea80, 0xb00a5631, 0x974c541a, 0x359e9963,
        0x5c971061, 0xccc07c79, 0x2098d9d6, 0x91dbd320,
    };

    chacha20.quarter_round(&state, 2, 7, 8, 13);
    try testing.expectEqualSlices(u32, reference[0..], state[0..]);
}

// https://www.rfc-editor.org/rfc/rfc7539#section-2.3.2
test "ChaCha20 Block Function" {
    const key = [chacha20.KEY_SIZE]u8{
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    };
    const nonce = [chacha20.Parameters_RFC7539.NONCE_SIZE]u8{
        0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00,
    };

    const reference = [chacha20.BLOCK_SIZE]u8{
        0x10, 0xf1, 0xe7, 0xe4, 0xd1, 0x3b, 0x59, 0x15, 0x50, 0x0f, 0xdd, 0x1f, 0xa3, 0x20, 0x71, 0xc4,
        0xc7, 0xd1, 0xf4, 0xc7, 0x33, 0xc0, 0x68, 0x03, 0x04, 0x22, 0xaa, 0x9a, 0xc3, 0xd4, 0x6c, 0x4e,
        0xd2, 0x82, 0x64, 0x46, 0x07, 0x9f, 0xaa, 0x09, 0x14, 0xc2, 0xd7, 0x05, 0xd9, 0x8b, 0x02, 0xa2,
        0xb5, 0x12, 0x9c, 0xd1, 0xde, 0x16, 0x4e, 0xb9, 0xcb, 0xd0, 0x83, 0xe8, 0xa2, 0x50, 0x3c, 0x4e,
    };

    var chacha = chacha20.chacha20_rfc7539_new(&key, &nonce, &word_to_bytes(1));
    defer chacha20.chacha20_destroy(&chacha);

    chacha20.block_function(&chacha);

    var buffer: [chacha20.BLOCK_SIZE]u8 = undefined;
    chacha20.serialize(chacha20.BLOCK_WORDS, &chacha.state, &buffer);

    try testing.expectEqualSlices(u8, reference[0..], buffer[0..]);
}

// https://www.rfc-editor.org/rfc/rfc7539#section-2.4.2
test "ChaCha20 Cipher" {
    const key = [chacha20.KEY_SIZE]u8{
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    };
    const nonce = [chacha20.Parameters_RFC7539.NONCE_SIZE]u8{
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00,
    };
    const counter = word_to_bytes(1);
    const plaintext = "Ladies and Gentlemen of the class of '99: If I could offer you only one tip for the future, sunscreen would be it.";
    const reference = [_]u8{
        0x6e, 0x2e, 0x35, 0x9a, 0x25, 0x68, 0xf9, 0x80, 0x41, 0xba, 0x07, 0x28, 0xdd, 0x0d, 0x69, 0x81,
        0xe9, 0x7e, 0x7a, 0xec, 0x1d, 0x43, 0x60, 0xc2, 0x0a, 0x27, 0xaf, 0xcc, 0xfd, 0x9f, 0xae, 0x0b,
        0xf9, 0x1b, 0x65, 0xc5, 0x52, 0x47, 0x33, 0xab, 0x8f, 0x59, 0x3d, 0xab, 0xcd, 0x62, 0xb3, 0x57,
        0x16, 0x39, 0xd6, 0x24, 0xe6, 0x51, 0x52, 0xab, 0x8f, 0x53, 0x0c, 0x35, 0x9f, 0x08, 0x61, 0xd8,
        0x07, 0xca, 0x0d, 0xbf, 0x50, 0x0d, 0x6a, 0x61, 0x56, 0xa3, 0x8e, 0x08, 0x8a, 0x22, 0xb6, 0x5e,
        0x52, 0xbc, 0x51, 0x4d, 0x16, 0xcc, 0xf8, 0x06, 0x81, 0x8c, 0xe9, 0x1a, 0xb7, 0x79, 0x37, 0x36,
        0x5a, 0xf9, 0x0b, 0xbf, 0x74, 0xa3, 0x5b, 0xe6, 0xb4, 0x0b, 0x8e, 0xed, 0xf2, 0x78, 0x5e, 0x42,
        0x87, 0x4d,
    };

    var chacha = chacha20.chacha20_rfc7539_new(&key, &nonce, &counter);
    defer chacha20.chacha20_destroy(&chacha);

    var buffer: [plaintext.len]u8 = undefined;
    @memcpy(&buffer, plaintext);

    try chacha20.encrypt_inplace(&chacha, &buffer);
    try testing.expectEqualSlices(u8, reference[0..], buffer[0..]);
}

test "ChaCha20 32-bit counter increment edge cases" {
    const key = [chacha20.KEY_SIZE]u8{
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    };
    const nonce = [chacha20.Parameters_RFC7539.NONCE_SIZE]u8{
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00,
    };
    const counter = word_to_bytes(std.math.maxInt(u32) - 1);

    var chacha = chacha20.chacha20_rfc7539_new(&key, &nonce, &counter);
    defer chacha20.chacha20_destroy(&chacha);

    // Counter: 2^32 - 2 -> 2^32 - 1, OK
    try chacha20.increment_counter(&chacha);

    // Counter: 2^32 - 1 -> 2^32, overflow
    try testing.expectError(chacha20.KeyStreamDepleted, chacha20.increment_counter(&chacha));
}

test "ChaCha20 64-bit counter increment edge cases" {
    const key = [chacha20.KEY_SIZE]u8{
        0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    };
    const nonce = [chacha20.Parameters_Bernstein.NONCE_SIZE]u8{ 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00 };
    const counter = [chacha20.Parameters_Bernstein.COUNTER_SIZE]u8{ 0xff, 0xff, 0xff, 0xff, 0xef, 0xbe, 0xad, 0xde };

    var chacha = chacha20.chacha20_bernstein_new(&key, &nonce, &counter);
    defer chacha20.chacha20_destroy(&chacha);

    // Counter: 0xdeadbeefffffffff -> 0xdeadbef000000000, OK
    try chacha20.increment_counter(&chacha);
    try testing.expectEqualSlices(u32, &.{ 0x00000000, 0xdeadbef0 }, chacha.nonce[0..2]);

    // Counter: 0xffffffffffffffff -> 0x10000000000000000, overflow
    @memset(chacha.nonce[0..2], std.math.maxInt(u32));
    try testing.expectError(chacha20.KeyStreamDepleted, chacha20.increment_counter(&chacha));
}