// Copyright (c) blackshirt. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module slhdsa

// PrivateKey represents SLH-DSA keypair.
pub struct PrivateKey {
	key &C.EVP_PKEY
}

// new creates a new SLH-DSA PrivateKey based on the supplied options.
// By default, it will create a `SLH-DSA-SHA2-128s` key with random generator.
// See `enum Kind` for all of availables choice's type (kind) of key.
// Its also support to generate keys based on the seed or private bytes through
// provided options. See available options in `KeyOpts`.
// Usage example:
// ```v
// pk1 := slhdsa.PrivateKey.new()!; println(pk1)
// pk2 := slhdsa.PrivateKey.new(kind: .sha2_128s)!; println(pk2)
// ```
pub fn PrivateKey.new(opt KeyOpts) !PrivateKey {
	match opt.flag {
		// default random generated key, its recommended way
		0 {
			kind := opt.kind.long_name()
			key := C.EVP_PKEY_Q_keygen(0, 0, kind)
			if key == 0 {
				return error('EVP_PKEY_Q_keygen failed')
			}
			return PrivateKey{
				key: key
			}
		}
		// use seed bytes, usually for testing purposes
		1 {
			// if you dont provides seed bytes, it will discarded, and use random one
			return PrivateKey.from_seed(opt.seed, opt.kind)!
		}
		// use private bytes, usually for testing purposes
		2 {
			// if you dont provides private bytes, it will discarded, and use random one
			return PrivateKey.from_bytes(opt.priv, opt.kind)!
		}
		else {
			return error('Unsupported flag')
		}
	}
}

// sign signs the message using this key under desired options in opt.
pub fn (pv PrivateKey) sign(msg []u8, opt SignerOpts) ![]u8 {
	if msg.len == 0 {
		return error('Zero length of messages')
	}
	if opt.context.len > 255 {
		return error('The context string size was over than 255')
	}
	out := slhdsa_do_sign(pv.key, msg, opt)!
	return out
}

// verify verifies signature for the message msg under provided options.
// Its possible because of under the hood, private key is a key pair.
pub fn (pv PrivateKey) verify(sig []u8, msg []u8, opt SignerOpts) !bool {
	if opt.context.len > 255 {
		return error('The context string size was over than 255')
	}
	return slhdsa_do_verify(pv.key, sig, msg, opt)!
}

const default_bioread_bufsize = 4096

// dump_key represents PrivateKey in human readable string.
pub fn (pv PrivateKey) dump_key() !string {
	bo := C.BIO_new(C.BIO_s_mem())
	if bo == 0 {
		C.BIO_free_all(bo)
		return error('BIO_new failed')
	}
	n := C.EVP_PKEY_print_private(bo, pv.key, 2, 0)
	if n <= 0 {
		C.BIO_free_all(bo)
		return error('print private failed')
	}
	size := usize(0)
	buf := []u8{len: default_bioread_bufsize}
	m := C.BIO_read_ex(bo, buf.data, buf.len, &size)

	if m <= 0 {
		unsafe { buf.free() }
		C.BIO_free_all(bo)
		return error('BIO_read_ex failed')
	}
	output := buf[..size].bytestr()

	// Cleans up and return the result
	unsafe { buf.free() }
	C.BIO_free_all(bo)

	return output
}

// public_key gets the public part of this private key as a PublicKey.
pub fn (pv PrivateKey) public_key() !PublicKey {
	pbkey := C.EVP_PKEY_dup(pv.key)
	// we clears out the private bits from the key
	// by setting it into null
	n := C.EVP_PKEY_set_octet_string_param(pbkey, c'priv', 0, 0)
	if n <= 0 {
		C.EVP_PKEY_free(pbkey)
		return error('EVP_PKEY_set_octet_string_param')
	}

	return PublicKey{
		key: pbkey
	}
}

// free releases memory occupied by this key.
pub fn (mut pv PrivateKey) free() {
	C.EVP_PKEY_free(pv.key)
}

// PublicKey represents public key part from the key.
pub struct PublicKey {
	key &C.EVP_PKEY
}

// verify verifies signature sig for messages msg under options provided.
pub fn (pb PublicKey) verify(sig []u8, msg []u8, opt SignerOpts) !bool {
	if opt.context.len > 255 {
		return error('The context string size was over than 255')
	}
	return slhdsa_do_verify(pb.key, sig, msg, opt)!
}

// dump_key dumps this public key as a human readable string.
pub fn (pb PublicKey) dump_key() !string {
	bo := C.BIO_new(C.BIO_s_mem())
	n := C.EVP_PKEY_print_public(bo, pb.key, 2, 0)
	if n <= 0 {
		C.BIO_free_all(bo)
		return error('EVP_PKEY_print_public failed')
	}
	size := usize(0)
	mut m := C.BIO_read_ex(bo, 0, default_bioread_bufsize, &size)
	mut buf := []u8{len: int(size)}
	m = C.BIO_read_ex(bo, buf.data, buf.len, &size)
	if m <= 0 {
		C.BIO_free_all(bo)
		return error('BIO_read_ex failed')
	}

	output := buf[..size].bytestr()
	// Cleans up
	unsafe { buf.free() }
	C.BIO_free_all(bo)

	return output
}

// free releases memory occupied by this public key
pub fn (mut pb PublicKey) free() {
	C.EVP_PKEY_free(pb.key)
}

// the biggest supported size of seed was 32 bytes length, where private key contains 4*32 bytes
const default_privkey_buffer = 128

fn (pv PrivateKey) bytes() ![]u8 {
	priv_len := usize(0)
	priv := []u8{len: default_privkey_buffer}

	n := C.EVP_PKEY_get_octet_string_param(pv.key, c'priv', priv.data, priv.len, &priv_len)
	if n <= 0 {
		return error('EVP_PKEY_get_octet_string_param failed')
	}
	out := priv[..int(priv_len)].clone()
	unsafe { priv.free() }
	return out
}

// the biggest public key size was 64 bytes length, where public key
// was contains 2*n, so we relaxed it to 2*72 bytes
const default_pubkey_buffer = 144

fn (pv PrivateKey) public_bytes() ![]u8 {
	pblen := usize(0)
	buf := []u8{len: default_pubkey_buffer}
	n := C.EVP_PKEY_get_octet_string_param(pv.key, c'pub', buf.data, buf.len, &pblen)
	if n <= 0 {
		return error('EVP_PKEY_get_octet_string_param failed')
	}

	out := buf[..pblen].clone()
	unsafe { buf.free() }

	return out
}

// Signing and verifying Helpers
//
// slhdsa_do_sign performs signing msg with SLH-DSA key.
@[inline]
fn slhdsa_do_sign(key &C.EVP_PKEY, msg []u8, opt SignerOpts) ![]u8 {
	sctx := C.EVP_PKEY_CTX_new_from_pkey(0, key, 0)
	if sctx == 0 {
		C.EVP_PKEY_CTX_free(sctx)
		return error('EVP_PKEY_CTX_new_from_pkey failed')
	}

	// gets algo name from key
	algo_name := key_type_name(key)
	sig_alg := C.EVP_SIGNATURE_fetch(0, algo_name, 0)
	if sig_alg == 0 {
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('EVP_SIGNATURE_fetch failed')
	}

	param_bld := C.OSSL_PARAM_BLD_new()
	if param_bld == 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('OSSL_PARAM_BLD_new failed')
	}

	// writes `context-string` params into context key generator.
	// OSSL_PARAM_octet_string("context-string", (unsigned char *)"A context string", 33),
	cs := C.OSSL_PARAM_BLD_push_octet_string(param_bld, c'context-string', opt.context.str,
		opt.context.len)
	if cs <= 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('OSSL_PARAM_BLD_push_octet_string context-string flag FAILED')
	}

	// write `message-encoding` flag
	me := C.OSSL_PARAM_BLD_push_int(param_bld, c'message-encoding', opt.encoding)
	if me <= 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('OSSL_PARAM_BLD_push_int message-encoding flag FAILED')
	}

	// handle entropy testing
	// `test-entropy flag only handled when `encoding` flag was set into 0 value
	if opt.encoding == 0 {
		te := C.OSSL_PARAM_BLD_push_octet_string(param_bld, c'test-entropy', opt.entropy.data,
			opt.entropy.len)
		if te <= 0 {
			C.OSSL_PARAM_BLD_free(param_bld)
			C.EVP_SIGNATURE_free(sig_alg)
			C.EVP_PKEY_CTX_free(sctx)
			return error('OSSL_PARAM_BLD_push_octet_string test-entropy flag failed')
		}
	}

	// write `deterministic` flag, its getting ignored if "test-entropy" is set.
	dt := C.OSSL_PARAM_BLD_push_int(param_bld, c'deterministic', opt.deterministic)
	if dt <= 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('OSSL_PARAM_BLD_push_int deterministic flag FAILED')
	}
	// build params
	params := C.OSSL_PARAM_BLD_to_param(param_bld)

	// Then init the context with this params
	x := C.EVP_PKEY_sign_message_init(sctx, sig_alg, params)
	if x <= 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.OSSL_PARAM_free(params)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('EVP_PKEY_sign_message_init failed')
	}

	sig_len := usize(C.EVP_PKEY_size(key))
	buf := []u8{len: int(sig_len)}
	// Do signing the msg and updates the sig_len.
	m := C.EVP_PKEY_sign(sctx, buf.data, &sig_len, msg.data, msg.len)
	if m <= 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.OSSL_PARAM_free(params)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('EVP_PKEY_sign failed')
	}

	// return the copy of the sig
	sig := buf[..sig_len].clone()

	// cleans up
	unsafe { buf.free() }
	C.OSSL_PARAM_BLD_free(param_bld)
	C.OSSL_PARAM_free(params)
	C.EVP_SIGNATURE_free(sig_alg)
	C.EVP_PKEY_CTX_free(sctx)

	return sig
}

// slhdsa_do_verify performs verifying of provided signature in sign for the msg message.
@[inline]
fn slhdsa_do_verify(key &C.EVP_PKEY, sign []u8, msg []u8, opt SignerOpts) !bool {
	sctx := C.EVP_PKEY_CTX_new_from_pkey(0, key, 0)
	if sctx == 0 {
		C.EVP_PKEY_CTX_free(sctx)
		return error('EVP_PKEY_CTX_new_from_pkey failed')
	}

	// get the algo signature name from the key
	algo_name := key_type_name(key)
	sig_alg := C.EVP_SIGNATURE_fetch(0, algo_name, 0)
	if sig_alg == 0 {
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('EVP_SIGNATURE_fetch failed')
	}

	param_bld := C.OSSL_PARAM_BLD_new()
	if param_bld == 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('OSSL_PARAM_BLD_new failed')
	}

	// handle context option
	if opt.context.len > 0 {
		o := C.OSSL_PARAM_BLD_push_octet_string(param_bld, c'context-string', opt.context.str,
			opt.context.len)
		if o <= 0 {
			C.OSSL_PARAM_BLD_free(param_bld)
			C.EVP_SIGNATURE_free(sig_alg)
			C.EVP_PKEY_CTX_free(sctx)
			return error('OSSL_PARAM_BLD_push failed')
		}
	}
	// for testing
	if opt.encoding == 0 {
		oo := C.OSSL_PARAM_BLD_push_int(param_bld, c'message-encoding', opt.encoding)
		if oo <= 0 {
			C.OSSL_PARAM_BLD_free(param_bld)
			C.EVP_SIGNATURE_free(sig_alg)
			C.EVP_PKEY_CTX_free(sctx)
			return error('OSSL_PARAM_BLD_push FAILED')
		}
	}
	// build params
	params := C.OSSL_PARAM_BLD_to_param(param_bld)
	n := C.EVP_PKEY_verify_message_init(sctx, sig_alg, params)
	if n <= 0 {
		C.OSSL_PARAM_BLD_free(param_bld)
		C.OSSL_PARAM_free(params)
		C.EVP_SIGNATURE_free(sig_alg)
		C.EVP_PKEY_CTX_free(sctx)
		return error('OSSL_PARAM_BLD_push FAILED')
	}

	// do verifying of the signature.
	m := C.EVP_PKEY_verify(sctx, sign.data, sign.len, msg.data, msg.len)

	// Cleans up
	C.EVP_SIGNATURE_free(sig_alg)
	C.EVP_PKEY_CTX_free(sctx)
	C.OSSL_PARAM_BLD_free(param_bld)
	C.OSSL_PARAM_free(params)

	return m == 1
}