// 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

import math.big
import crypto.internal.subtle

// default_integer_tag is the default of ASN.1 INTEGER type.
pub const default_integer_tag = Tag{.universal, false, int(TagType.integer)}

// ASN.1 INTEGER.
//
// The INTEGER type value can be a positive or negative number.
// There are no limits imposed on the magnitude of INTEGER values in the ASN.1 standard.
// Its handles number arbitrary length of number with support of `math.big` module.
// But, for sake of safety, we limit the INTEGER limit to follow allowed length in
// definite form of Length part, ie, 1008 bit, or 126 bytes
// The encoding of an integer number shall be primitive.

// Limit of length of INTEGER type, in bytes
// Known big RSA keys is 4096 bits, ie, 512 bytes
const max_integer_length = 2048

// Integer represent Universal class of arbitrary length type of ASN.1 INTEGER.
// The encoding of an integer value shall be primitive.
// If the contents octets of an integer value encoding consist of more than one octet,
// then the bits of the first octet and bit 8 of the second octet.
// 	a) shall not all be ones; and.
// 	b) shall not all be zero.
// NOTE – These rules ensure that an integer value is always encoded in
// the smallest possible number of octets.
pub struct Integer {
pub:
	// underlying integer value with support from `i64` and `big.Integer`
	value IntValue
}

// hex returns Integer value as a hex string.
pub fn (v Integer) hex() string {
	match v.value {
		i64 {
			val := v.value as i64
			return val.hex_full()
		}
		big.Integer {
			val := v.value as big.Integer
			return val.hex()
		}
	}
}

fn (v Integer) str() string {
	return 'Integer ${v.value.str()}'
}

// IntValue represents arbitrary integer value, currently we support
// through primitive 164 type for integer value below < max_i64, and
// use `big.Integer` for support arbitrary length of integer values.
type IntValue = big.Integer | i64

fn (v IntValue) str() string {
	match v {
		i64 {
			val := v as i64
			return val.str()
		}
		big.Integer {
			val := v as big.Integer
			return val.str()
		}
	}
}

// bytes get the bytes representation from underlying IntValue.
fn (v IntValue) bytes() []u8 {
	match v {
		i64 {
			return i64_to_bytes(v)
		}
		big.Integer {
			// if v == big.zero_int or similar big.Integer values that produces empty bytes,
			// returns v.bytes() directly can lead to undesired behavior thats doesn't aligned with
			// ASN.1 INTEGER requirement. See the discussion on the discord about the issues
			// at https://discord.com/channels/592103645835821068/592294828432424960/1230460279733620777
			// so, we do some hack to get the correct value
			// TODO: find the correct way to tackle this
			if v == big.zero_int {
				return [u8(0x00)]
			}
			// todo: proper check of 0 bytes length
			if v.bit_len() == 0 {
				return [u8(0x00)]
			}
			// otherwise, we use v.bytes() directly
			b, _ := v.bytes()
			return b
		}
	}
}

// from_i64 creates new a ASN.1 Integer from i64 v.
pub fn Integer.from_i64(v i64) Integer {
	return Integer{
		value: IntValue(v)
	}
}

// from_int creates a new ASN.1 Integer from int v.
pub fn Integer.from_int(v int) Integer {
	return Integer{
		value: IntValue(i64(v))
	}
}

// from_bigint creates a new ASN.1 Integer from big.Integer b
pub fn Integer.from_bigint(b big.Integer) Integer {
	return Integer{
		value: IntValue(b)
	}
}

// from_string creates a new ASN.1 Integer from decimal string s.
// If your string value is below max_i64, use from_i64 instead
pub fn Integer.from_string(s string) !Integer {
	v := big.integer_from_string(s)!
	return Integer{
		value: IntValue(v)
	}
}

// from_hex creates a new ASN.1 Integer from hex string in x
// where x is a valid hex string without `0x` prefix.
// If your string value is below max_i64, use from_i64 instead
pub fn Integer.from_hex(x string) !Integer {
	s := big.integer_from_radix(x, 16)!
	return Integer{
		value: IntValue(s)
	}
}

// from_bytes creates a new ASN.1 Integer from bytes array in b.
// Its try to parse bytes as in two's complement form.
fn Integer.from_bytes(b []u8) !Integer {
	return Integer.unpack_from_twoscomplement_bytes(b)!
}

// unpack_from_twoscomplement_bytes parses the bytes in b into the Integer
// value in the big-endian two's complement way. If b[0]&80 != 0, the number
// is negative. If b is empty it would be error.
fn Integer.unpack_from_twoscomplement_bytes(b []u8) !Integer {
	// FIXME: should we return error instead ?
	if b.len == 0 {
		return error('Integer: null bytes')
	}
	if b.len > 7 {
		mut num := big.integer_from_bytes(b)
		// negative number
		if b.len > 0 && b[0] & 0x80 > 0 {
			sub := big.one_int.left_shift(u32(b.len) * 8)
			num -= sub
		}

		return Integer{
			value: IntValue(num)
		}
	}
	// use i64
	val := read_i64(b)!
	res := Integer.from_i64(val)
	return res
}

// bytes return underlying bytes array
fn (v Integer) bytes() []u8 {
	return v.value.bytes()
}

// bytes_len returns underlying bytes length
fn (v Integer) bytes_len() int {
	b := v.value.bytes()
	return b.len
}

// tag returns the tag of Integer type element
pub fn (v Integer) tag() Tag {
	return default_integer_tag
}

// payload returns the payload of Integer type element.
pub fn (v Integer) payload() ![]u8 {
	bytes, _ := v.pack_into_twoscomplement_form()!
	return bytes
}

// pack_into_twoscomplement_form serialize Integer in two's-complement rules.
// 	- 	The integer value contains the encoded integer if it is positive, or its two's complement if it is negative.
// 	- 	If the integer is positive but the high order bit is set to 1, a leading 0x00 is added to the content
// 		to indicate that the number is not negative.
// 	-	If the number is negative after applying two's-complement rules, and the most-significant-bit of the
// 		the high order bit of the bytes results isn't set, pad it with 0xff in order to keep the number negative.
fn (v Integer) pack_into_twoscomplement_form() !([]u8, int) {
	match v.value {
		i64 {
			val := v.value as i64
			mut bytes := i64_to_bytes(val)
			return bytes, bytes.len
		}
		big.Integer {
			match v.value.signum {
				0 {
					return [u8(0x00)], 1
				}
				1 {
					mut b := v.bytes()
					// handle the zero issues
					if b.len == 0 {
						return [u8(0x00)], 1
					}
					// If the integer is positive but the high order bit is set to 1, a leading 0x00 is added
					// to the content to indicate that the number is not negative
					if b[0] & 0x80 > 0 {
						b.prepend(u8(0x00))
					}
					return b, b.len
				}
				-1 {
					// A negative number has to be converted to two's-complement form.
					// by invert the number and then subtract it with big(1), or with other mean
					// Flip all of the bits in the value and then add one to the resulting value.
					// If the most-significant-bit isn't set then we'll need to pad the
					// beginning with 0xff in order to keep the number negative.
					negv := v.value.neg()
					negvminus1 := negv - big.one_int
					mut bytes, _ := negvminus1.bytes()
					for i, _ in bytes {
						bytes[i] ^= 0xff
					}
					if bytes.len == 0 || bytes[0] & 0x80 == 0 {
						bytes.prepend(u8(0xff))
					}
					return bytes, bytes.len
				}
				else {
					return error('should unreachable')
				}
			}
		}
	}
}

// equal do checking if integer n was equal to integer m.
// ISSUE?: There are some issues when compared n == m directly,
// its fails even internally its a same, so we provide and use equality check
pub fn (n Integer) equal(m Integer) bool {
	nbytes := n.bytes()
	mbytes := m.bytes()
	// todo: check sign equality
	// m.tag() == n.tag() by definition, no need to check
	return subtle.constant_time_compare(nbytes, mbytes) == 1
}

// Integer.unpack_and_validate deserializes bytes in b into Integer
// in two's complement way and perform validation on this bytes to
// meet der requirement.
fn Integer.unpack_and_validate(b []u8) !Integer {
	if !valid_bytes(b, true) {
		return error('Integer: check return false')
	}
	ret := Integer.unpack_from_twoscomplement_bytes(b)!
	return ret
}

// as_bigint casts Integer value to big.Integer or error on fails.
pub fn (v Integer) as_bigint() !big.Integer {
	if v.value is big.Integer {
		val := v.value as big.Integer
		return val
	}
	return error('Integer not hold big.Integer type')
}

// as_i64 casts Integer value to i64 value or error on fails.
pub fn (v Integer) as_i64() !i64 {
	if v.value is i64 {
		val := v.value as i64
		return val
	}
	return error('Integer not hold i64 type')
}

// parse tries to read and parse into Integer type or return error on fails.
fn Integer.parse(mut p Parser) !Integer {
	tag := p.read_tag()!
	if !tag.equal(default_integer_tag) {
		return error('Get unexected non Integer tag')
	}
	length := p.read_length()!
	if length < 1 {
		return error('Get length < 1 for Integer length')
	}
	bytes := p.read_bytes(length)!
	ret := Integer.from_bytes(bytes)!

	return ret
}

// decode tries to decode bytes array into Integer type or error on fails
fn Integer.decode(bytes []u8) !(Integer, int) {
	return Integer.decode_with_rule(bytes, 0, .der)!
}

// decode_with_rule tries to decode bytes back into ASN.1 Integer.
// Its accepts `loc` params, the location (offset) within bytes where the unpack start from.
// If not sure set to 0 to drive unpacking and rule of `Encodingrule`, currently only support`.der`.
fn Integer.decode_with_rule(bytes []u8, loc int, rule EncodingRule) !(Integer, int) {
	if bytes.len < 3 {
		return error('Integer: bad bytes length')
	}
	tag, length_pos := Tag.decode_with_rule(bytes, loc, rule)!
	if !tag.equal(default_integer_tag) {
		return error('Get unexpected Integer tag')
	}
	length, content_pos := Length.decode_with_rule(bytes, length_pos, rule)!
	payload := if length == 0 {
		[]u8{}
	} else {
		if content_pos + length > bytes.len {
			return error('Not enought bytes to read on')
		}
		unsafe { bytes[content_pos..content_pos + length] }
	}

	// buf := trim_bytes(payload)!
	next := content_pos + length
	result := Integer.from_bytes(payload)!

	return result, next
}

// Utility function
//
fn is_highest_bit_set(src []u8) bool {
	if src.len > 0 {
		return src[0] & 0x80 == 0
	}
	return false
}

fn trim_bytes(src []u8) ![]u8 {
	if src.len == 0 {
		return error('bad src')
	}
	// TODO: removes prepended bytes when its meet criteria
	// positive value but its prepended with 0x00
	if src.len > 1 && src[0] == 0x00 && src[1] & 0x80 > 0 {
		bytes := src[1..]
		return bytes
	}
	// TODO: how to do with multiples 0xff
	if src.len > 1 && src[0] == 0xff && src[1] & 0x80 == 0 {
		bytes := src[1..]
		return bytes
	}
	return src
}

// length_i64 gets bytes length needed to reperesent this i64 value
fn length_i64(val i64) int {
	mut i := val
	mut n := 1

	for i > 127 {
		n++
		i >>= 8
	}

	for i < -128 {
		n++
		i >>= 8
	}

	return n
}

// i64_to_bytes transforms i64 value into bytes representation
fn i64_to_bytes(i i64) []u8 {
	mut n := length_i64(i)
	mut dst := []u8{len: n}
	for j := 0; j < n; j++ {
		dst[j] = u8(i >> u32((n - 1 - j) * 8))
	}
	return dst
}

// read_i64 read src as signed i64
fn read_i64(src []u8) !i64 {
	if !valid_bytes(src, true) {
		return error('i64 check return false')
	}
	mut ret := i64(0)

	if src.len > 8 {
		return error('too large integer')
	}
	for i := 0; i < src.len; i++ {
		ret <<= 8
		ret |= i64(src[i])
	}

	ret <<= 64 - u8(src.len) * 8
	ret >>= 64 - u8(src.len) * 8

	// try to serialize back, and check its matching original one
	// and gives a warning when its not match.

	dst := i64_to_bytes(ret)
	if dst != src {
		eprintln('maybe integer bytes not in shortest form')
	}

	return ret
}

// i32 handling
//
// read_i32 read from bytes
fn read_i32(src []u8) !int {
	if !valid_bytes(src, true) {
		return error('i32 check return false')
	}

	ret := read_i64(src)!
	if ret != i64(int(ret)) {
		return error('integer too large')
	}

	return int(ret)
}

fn length_i32(v i32) int {
	return length_i64(i64(v))
}

fn i32_to_bytes(v i32) []u8 {
	return i64_to_bytes(i64(v))
}