// Copyright (c) 2022, 2024 blackshirt. All rights reserved.
// Use of this source code is governed by a MIT License
// that can be found in the LICENSE file.
module asn1

// max_definite_length_count is a limit tells how many bytes to represent this length.
// We're going to limit this to 6 bytes following when the length is in long-definite form.
const max_definite_length_count = 6
const max_definite_length_value = max_int

// ASN.1 length handling routines.
//
// The standard of X.690 ITU document defines two length types - definite and indefinite.
// In DER encoding only uses the definite length. The length field must be encoded in the minimum number of octets.
// There are two forms of definite length octets: short (for lengths value between 0 and 127),
// and long definite (for lengths value between 0 and 2^1008 -1).
// In short form contains one octet. Bit 8 has value "0" and bits 7-1 give the length (length value from 0 to 127)
// In long form, its required two or more octets, limited to 127 octets. Bit 8 of first octet has value "1" and bits 7-1 gives
// the number of additional length octets.
// Second and following octets give the length, base 256, most significant digit first.
//
// This module only support definite length, in short or long form. Its required for DER encoding
// the length octets should be in definite length.
//
// Length represent ASN.1 length value
pub type Length = int

// new creates Length from integer  value. Passing negative value (<0) for length
// is not make a sense, so just return error instead if it happen.
pub fn Length.new(v int) !Length {
	if v < 0 {
		return error('Length: supply with positive int')
	}
	if v > max_definite_length_value {
		return error('Length: value provided exceed limit')
	}
	return Length(v)
}

// Length.from_bytes tries to read length from bytes array and return the length
// and the rest of remaining bytes on success, or error on fails.
pub fn Length.from_bytes(bytes []u8) !(Length, []u8) {
	length, next := Length.decode(bytes)!
	if next < bytes.len {
		rest := unsafe { bytes[next..] }
		return length, rest
	}
	if next == bytes.len {
		return length, []u8{}
	}
	return error('Length: too short data')
}

// encode serializes the length v into bytes array stored into buffer buf in .der rule.
pub fn (v Length) encode(mut dst []u8) ! {
	v.encode_with_rule(mut dst, .der)!
}

// encode_with_rule serializes Length v into bytes and append it into `dst`.
// it would use rule of `Encodingrule` to drive how encode operation would be done.
// By default the .der rule is only currently supported.
fn (v Length) encode_with_rule(mut dst []u8, rule EncodingRule) ! {
	// we currently only support .der and (stricter) .ber
	if rule != .der && rule != .ber {
		return error('Length: unsupported rule')
	}
	// TODO: add supports for undefinite form
	// Long form
	if v >= 128 {
		// First, we count how many bytes occupied by this length value.
		// if the count exceed the limit, we return error.
		count := v.bytes_len()
		if count > max_definite_length_count {
			return error('something bad in your length')
		}
		// In definite long form, msb bit of first byte is set into 1, and the remaining bits
		// of first byte tells exact count how many bytes following representing this length value.
		dst << 0x80 | u8(count)
		v.to_bytes(mut dst)
	} else {
		// short form, already tells the length value.
		dst << u8(v)
	}
}

// bytes_len tells how many bytes needed to represent this length
fn (v Length) bytes_len() int {
	mut i := v
	mut num := 1
	for i > 255 {
		num++
		i >>= 8
	}
	return num
}

// to_bytes packs Length v into bytes and apends it into `dst` bytes.
fn (v Length) to_bytes(mut dst []u8) {
	mut n := v.bytes_len()
	for ; n > 0; n-- {
		// pay attention to the brackets
		dst << u8(v >> ((n - 1) * 8))
	}
}

// length_size gets length of length v in .der rule
fn (v Length) length_size() !int {
	n := v.length_size_with_rule(.der)!
	return n
}

// length_size_with_rule calculates the length of bytes needed to store the Length value `v`
// includes one byte marker for long definite form of length value, for value >= 128
fn (v Length) length_size_with_rule(rule EncodingRule) !int {
	// we currently only support .der or (stricter) .ber
	if rule != .der && rule != .ber {
		return error('Length: unsupported rule')
	}
	n := if v < 128 { 1 } else { v.bytes_len() + 1 }
	return n
}

// decode read length from bytes src with default context or return error on fails.
fn Length.decode(src []u8) !(Length, int) {
	ret, next := Length.decode_from_offset(src, 0)!
	return ret, next
}

// decode_from_offset read length from src bytes start from offset pos or return error on fails.
fn Length.decode_from_offset(src []u8, pos int) !(Length, int) {
	ret, next := Length.decode_with_rule(src, pos, .der)!
	return ret, next
}

// decode_with_rule tries to decode and deserializes buffer in src into Length form, start from offset loc in the buffer.
// Its return Length and next offset in the buffer to process on, or return error on fails.
fn Length.decode_with_rule(src []u8, loc int, rule EncodingRule) !(Length, int) {
	if src.len < 1 {
		return error('Length: truncated length')
	}
	// preliminary check
	if rule != .der && rule != .ber {
		return error('Length: unsupported rule')
	}
	// consider b := src[loc] would lead to panic
	if loc >= src.len {
		return error('Length: invalid pos')
	}

	mut pos := loc
	mut b := src[pos]
	pos += 1
	mut length := 0
	// check for the most bit is set or not
	if b & 0x80 == 0 {
		// for lengths between 0 and 127, the one-octet short form can be used.
		// The bit 7 of the length octet is set to 0, and the length is encoded
		// as an unsigned binary value in the octet's rightmost seven bits.
		length = b & 0x7f
	} else {
		if pos > src.len {
			return error('truncated tag or length')
		}
		// Otherwise, its a Long definite form or undefinite form
		num_bytes := b & 0x7f
		if num_bytes == 0 {
			// TODO: add support for undefinite length
			return error('Length: unsupported undefinite length')
		}
		if num_bytes == 0x7f {
			return error('Length: 0x7f is for reserved use')
		}
		// we limit the bytes count for length definite form to `max_definite_length_count`
		if num_bytes > max_definite_length_count {
			return error('Length: count bytes exceed limit')
		}
		for i := 0; i < num_bytes; i++ {
			if pos >= src.len {
				return error('Length: truncated length')
			}
			b = src[pos]
			pos += 1

			// currently, we're only support limited length.
			// The length is in int range
			if length > max_definite_length_value {
				return error('Length: length exceed limit value')
			}
			length <<= 8
			length |= b
			if length == 0 {
				// TODO: leading zeros is allowed in Long form of BER encoding, but
				// not allowed in DER encoding
				if rule == .der {
					return error('Length: leading zeros')
				}
			}
		}
		// do not allow values < 0x80 to be encoded in long form
		if length < 0x80 {
			// TODO: allow in BER
			return error('Length: dont needed in long form')
		}
	}
	ret := Length.new(length)!
	return ret, pos
}