// Copyright (c) 2023 Kim Shrier. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
// Package blake3 implements the Blake3 cryptographic hash
// as described in:
// https://github.com/BLAKE3-team/BLAKE3-specs/blob/master/blake3.pdf
// Version 20211102173700

module blake3

import encoding.hex
import json
import os

struct TestObject {
    comment        string
    key            string
    context_string string
    cases          []TestVectors
}

struct TestVectors {
    input_len  u64
    hash       string
    keyed_hash string
    derive_key string
}

const object_string = os.read_file(os.real_path(os.join_path(os.dir(@FILE), 'testdata',
    'test_vectors.json'))) or { panic(err) }

const test_object = json.decode(TestObject, object_string) or { panic(err) }

const data_segment = []u8{len: 251, cap: 251, init: u8(index)}

fn test_run_test_vectors() {
    mut data := []u8{}
    for i in 0 .. 408 {
        data << data_segment
    }

    for case in test_object.cases {
        extended_length := u64(case.hash.len)

        // test Blake3 hash

        mut hash_d := Digest.new_hash() or {
            assert false, 'Digest.new_hash error: ${err}'
            return
        }

        hash_d.write(data[..int(case.input_len)]) or {
            assert false, 'hash_d.write error: ${err}'
            return
        }

        hash_bytes := hex.decode(case.hash) or {
            assert false, 'invalid hex string: ${err}'
            return
        }

        // test various sizes and see if we can ask for different
        // sizes without recomputing the hash

        for i in [u64(16), 20, 24, 28, 32, 48, 64, 72, 96] {
            if hash_bytes.len > i {
                assert hash_d.checksum(i) == hash_bytes[..int(i)], 'hash failed output length ${i}'
            }
        }

        assert hash_d.checksum(u64(hash_bytes.len)) == hash_bytes, 'hash failed output length ${extended_length}'

        // test Blake3 keyed hash

        mut keyed_hash_d := Digest.new_keyed_hash(test_object.key.bytes()) or {
            assert false, 'Digest.new_keyed_hash error: ${err}'
            return
        }

        keyed_hash_d.write(data[..int(case.input_len)]) or {
            assert false, 'keyed_hash_d.write error: ${err}'
            return
        }

        keyed_hash_bytes := hex.decode(case.keyed_hash) or {
            assert false, 'invalid hex string: ${err}'
            return
        }

        // test various sizes and see if we can ask for different
        // sizes without recomputing the hash

        for i in [u64(16), 20, 24, 28, 32, 48, 64, 72, 96] {
            if keyed_hash_bytes.len > i {
                assert keyed_hash_d.checksum(i) == keyed_hash_bytes[..int(i)], 'keyed hash failed output length ${i}'
            }
        }

        assert keyed_hash_d.checksum(u64(keyed_hash_bytes.len)) == keyed_hash_bytes, 'keyed hash failed output length ${extended_length}'
        // test Blake3 derive key hash

        mut derive_key_hash_d := Digest.new_derive_key_hash(test_object.context_string.bytes()) or {
            assert false, 'Digest.new_derive_key_hash error: ${err}'
            return
        }

        derive_key_hash_d.write(data[..int(case.input_len)]) or {
            assert false, 'derive_key_hash_d.write error: ${err}'
            return
        }

        derive_key_hash_bytes := hex.decode(case.derive_key) or {
            assert false, 'invalid hex string: ${err}'
            return
        }

        // test various sizes and see if we can ask for different
        // sizes without recomputing the hash

        for i in [u64(16), 20, 24, 28, 32, 48, 64, 72, 96] {
            if derive_key_hash_bytes.len > i {
                assert derive_key_hash_d.checksum(i) == derive_key_hash_bytes[..int(i)], 'derive key hash failed output length ${i}'
            }
        }

        assert derive_key_hash_d.checksum(u64(derive_key_hash_bytes.len)) == derive_key_hash_bytes, 'derive key hash failed output length ${extended_length}'
    }
}