package field_p384r1

import "core:crypto"
import subtle "core:crypto/_subtle"
import "core:encoding/endian"
import "core:math/bits"

fe_clear :: proc "contextless" (arg1: ^Montgomery_Domain_Field_Element) {
	crypto.zero_explicit(arg1, size_of(Montgomery_Domain_Field_Element))
}

fe_clear_vec :: proc "contextless" (
	arg1: []^Montgomery_Domain_Field_Element,
) {
	for fe in arg1 {
		fe_clear(fe)
	}
}

fe_from_bytes :: proc "contextless" (
	out1: ^Montgomery_Domain_Field_Element,
	arg1: []byte,
	unsafe_assume_canonical := false,
) -> bool {
	ensure_contextless(len(arg1) == 48, "p384r1: invalid fe input buffer")

	// Note: We assume the input is in big-endian.
	tmp := Non_Montgomery_Domain_Field_Element {
		endian.unchecked_get_u64be(arg1[40:]),
		endian.unchecked_get_u64be(arg1[32:]),
		endian.unchecked_get_u64be(arg1[24:]),
		endian.unchecked_get_u64be(arg1[16:]),
		endian.unchecked_get_u64be(arg1[8:]),
		endian.unchecked_get_u64be(arg1[0:]),
	}
	defer crypto.zero_explicit(&tmp, size_of(tmp))

	// Check that tmp is in the the range [0, ELL).
	if !unsafe_assume_canonical {
		_, borrow := bits.sub_u64(ELL[0] - 1, tmp[0], 0)
		_, borrow = bits.sub_u64(ELL[1], tmp[1], borrow)
		_, borrow = bits.sub_u64(ELL[2], tmp[2], borrow)
		_, borrow = bits.sub_u64(ELL[3], tmp[3], borrow)
		_, borrow = bits.sub_u64(ELL[4], tmp[4], borrow)
		_, borrow = bits.sub_u64(ELL[5], tmp[5], borrow)
		if borrow != 0 {
			return false
		}
	}

	fe_to_montgomery(out1, &tmp)

	return true
}

fe_to_bytes :: proc "contextless" (out1: []byte, arg1: ^Montgomery_Domain_Field_Element) {
	ensure_contextless(len(out1) == 48, "p384r1: invalid fe output buffer")

	tmp: Non_Montgomery_Domain_Field_Element = ---
	fe_from_montgomery(&tmp, arg1)

	// Note: Likewise, output in big-endian.
	endian.unchecked_put_u64be(out1[40:], tmp[0])
	endian.unchecked_put_u64be(out1[32:], tmp[1])
	endian.unchecked_put_u64be(out1[24:], tmp[2])
	endian.unchecked_put_u64be(out1[16:], tmp[3])
	endian.unchecked_put_u64be(out1[8:], tmp[4])
	endian.unchecked_put_u64be(out1[0:], tmp[5])

	crypto.zero_explicit(&tmp, size_of(tmp))
}

@(require_results)
fe_equal :: proc "contextless" (arg1, arg2: ^Montgomery_Domain_Field_Element) -> int {
	tmp: Montgomery_Domain_Field_Element = ---
	fe_sub(&tmp, arg1, arg2)

	// This will only underflow if and only if (⟺) arg1 == arg2, and we return the borrow,
	// which will be 1.
	is_eq := subtle.u64_is_zero(fe_non_zero(&tmp))

	fe_clear(&tmp)

	return int(is_eq)
}

@(require_results)
fe_is_odd :: proc "contextless" (arg1: ^Montgomery_Domain_Field_Element) -> int {
	tmp: Non_Montgomery_Domain_Field_Element = ---
	defer crypto.zero_explicit(&tmp, size_of(tmp))

	fe_from_montgomery(&tmp, arg1)
	return int(tmp[0] & 1)
}

fe_pow2k :: proc "contextless" (
	out1: ^Montgomery_Domain_Field_Element,
	arg1: ^Montgomery_Domain_Field_Element,
	arg2: uint,
) {
	// Special case: `arg1^(2 * 0) = 1`, though this should never happen.
	if arg2 == 0 {
		fe_one(out1)
		return
	}

	fe_square(out1, arg1)
	for _ in 1 ..< arg2 {
		fe_square(out1, out1)
	}
}

fe_inv :: proc "contextless" (out1, arg1: ^Montgomery_Domain_Field_Element) {
	// Inversion computation is derived from the addition chain:
	//
	//	_10     = 2*1
	//	_11     = 1 + _10
	//	_110    = 2*_11
	//	_111    = 1 + _110
	//	_111000 = _111 << 3
	//	_111111 = _111 + _111000
	//	x12     = _111111 << 6 + _111111
	//	x24     = x12 << 12 + x12
	//	x30     = x24 << 6 + _111111
	//	x31     = 2*x30 + 1
	//	x32     = 2*x31 + 1
	//	x63     = x32 << 31 + x31
	//	x126    = x63 << 63 + x63
	//	x252    = x126 << 126 + x126
	//	x255    = x252 << 3 + _111
	//	return    ((x255 << 33 + x32) << 94 + x30) << 2
	//
	// Operations: 383 squares 14 multiplies
	//
	// Generated by github.com/mmcloughlin/addchain v0.4.0.

	// Note: Need to stash `arg1` (`xx`) in the case that `out1`/`arg1` alias,
	// as `arg1` is used after `out1` has been altered.
	t0, t1, t2, t3, xx: Montgomery_Domain_Field_Element = ---, ---, ---, ---, arg1^

	// Step 1: z = x^0x2
	fe_square(out1, arg1)

	// Step 2: z = x^0x3
	fe_mul(out1, &xx, out1)

	// Step 3: z = x^0x6
	fe_square(out1, out1)

	// Step 4: t1 = x^0x7
	fe_mul(&t1, &xx, out1)

	// Step 7: z = x^0x38
	fe_pow2k(out1, &t1, 3)

	// Step 8: z = x^0x3f
	fe_mul(out1, &t1, out1)

	// Step 14: t0 = x^0xfc0
	fe_pow2k(&t0, out1, 6)

	// Step 15: t0 = x^0xfff
	fe_mul(&t0, out1, &t0)

	// Step 27: t2 = x^0xfff000
	fe_pow2k(&t2, &t0, 12)

	// Step 28: t0 = x^0xffffff
	fe_mul(&t0, &t0, &t2)

	// Step 34: t0 = x^0x3fffffc0
	fe_pow2k(&t0, &t0, 6)

	// Step 35: z = x^0x3fffffff
	fe_mul(out1, out1, &t0)

	// Step 36: t0 = x^0x7ffffffe
	fe_square(&t0, out1)

	// Step 37: t2 = x^0x7fffffff
	fe_mul(&t2, &xx, &t0)

	// Step 38: t0 = x^0xfffffffe
	fe_square(&t0, &t2)

	// Step 39: t0 = x^0xffffffff
	fe_mul(&t0, &xx, &t0)

	// Step 70: t3 = x^0x7fffffff80000000
	fe_pow2k(&t3, &t0, 31)

	// Step 71: t2 = x^0x7fffffffffffffff
	fe_mul(&t2, &t2, &t3)

	// Step 134: t3 = x^0x3fffffffffffffff8000000000000000
	fe_pow2k(&t3, &t2, 63)

	// Step 135: t2 = x^0x3fffffffffffffffffffffffffffffff
	fe_mul(&t2, &t2, &t3)

	// Step 261: t3 = x^0xfffffffffffffffffffffffffffffffc0000000000000000000000000000000
	fe_pow2k(&t3, &t2, 126)

	// Step 262: t2 = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff
	fe_mul(&t2, &t2, &t3)

	// Step 265: t2 = x^0x7ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff8
	fe_pow2k(&t2, &t2, 3)

	// Step 266: t1 = x^0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff
	fe_mul(&t1, &t1, &t2)

	// Step 299: t1 = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffe00000000
	fe_pow2k(&t1, &t1, 33)

	// Step 300: t0 = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff
	fe_mul(&t0, &t0, &t1)

	// Step 394: t0 = x^0x3fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffbfffffffc00000000000000000000000
	fe_pow2k(&t0, &t0, 94)

	// Step 395: z = x^0x3fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffbfffffffc0000000000000003fffffff
	fe_mul(out1, out1, &t0)

	// Step 397: z = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff0000000000000000fffffffc
	fe_pow2k(out1, out1, 2)

	fe_mul(out1, out1, &xx)

	fe_clear_vec([]^Montgomery_Domain_Field_Element{&t0, &t1, &t2, &t3, &xx})
}

@(require_results)
fe_sqrt :: proc "contextless" (out1, arg1: ^Montgomery_Domain_Field_Element) -> int {
	// Square root candidate can be derived via exponentiation by `(p + 1) / 4`
	// From sage: 9850501549098619803069760025035903451269934817616361666987073351061430442874217582261816522064734500465401743278080
	//
	// Inversion computation is derived from the addition chain:
	//
	//	_10     = 2*1
	//	_11     = 1 + _10
	//	_110    = 2*_11
	//	_111    = 1 + _110
	//	_111000 = _111 << 3
	//	_111111 = _111 + _111000
	//	x12     = _111111 << 6 + _111111
	//	x24     = x12 << 12 + x12
	//	x30     = x24 << 6 + _111111
	//	x31     = 2*x30 + 1
	//	x32     = 2*x31 + 1
	//	x63     = x32 << 31 + x31
	//	x126    = x63 << 63 + x63
	//	x252    = x126 << 126 + x126
	//	x255    = x252 << 3 + _111
	//	return    ((x255 << 33 + x32) << 94 + x30) << 2
	//
	// Operations: 383 squares 14 multiplies
	//
	// Generated by github.com/mmcloughlin/addchain v0.4.0.

	// Likewise this tramples over arg1, so stash another copy.
	t0, t1, t2, t3, xx: Montgomery_Domain_Field_Element =  ---, ---, ---, ---, arg1^

	// Step 1: z = x^0x2
	fe_square(out1, arg1)

	// Step 2: z = x^0x3
	fe_mul(out1, &xx, out1)

	// Step 3: z = x^0x6
	fe_square(out1, out1)

	// Step 4: t1 = x^0x7
	fe_mul(&t1, &xx, out1)

	// Step 7: z = x^0x38
	fe_pow2k(out1, &t1, 3)

	// Step 8: z = x^0x3f
	fe_mul(out1, &t1, out1)

	// Step 14: t0 = x^0xfc0
	fe_pow2k(&t0, out1, 6)

	// Step 15: t0 = x^0xfff
	fe_mul(&t0, out1, &t0)

	// Step 27: t2 = x^0xfff000
	fe_pow2k(&t2, &t0, 12)

	// Step 28: t0 = x^0xffffff
	fe_mul(&t0, &t0, &t2)

	// Step 34: t0 = x^0x3fffffc0
	fe_pow2k(&t0, &t0, 6)

	// Step 35: z = x^0x3fffffff
	fe_mul(out1, out1, &t0)

	// Step 36: t0 = x^0x7ffffffe
	fe_square(&t0, out1)

	// Step 37: t2 = x^0x7fffffff
	fe_mul(&t2, &xx, &t0)

	// Step 38: t0 = x^0xfffffffe
	fe_square(&t0, &t2)

	// Step 39: t0 = x^0xffffffff
	fe_mul(&t0, &xx, &t0)

	// Step 70: t3 = x^0x7fffffff80000000
	fe_pow2k(&t3, &t0, 31)

	// Step 71: t2 = x^0x7fffffffffffffff
	fe_mul(&t2, &t2, &t3)

	// Step 134: t3 = x^0x3fffffffffffffff8000000000000000
	fe_pow2k(&t3, &t2, 63)

	// Step 135: t2 = x^0x3fffffffffffffffffffffffffffffff
	fe_mul(&t2, &t2, &t3)

	// Step 261: t3 = x^0xfffffffffffffffffffffffffffffffc0000000000000000000000000000000
	fe_pow2k(&t3, &t2, 126)

	// Step 262: t2 = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff
	fe_mul(&t2, &t2, &t3)

	// Step 265: t2 = x^0x7ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff8
	fe_pow2k(&t2, &t2, 3)

	// Step 266: t1 = x^0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff
	fe_mul(&t1, &t1, &t2)

	// Step 299: t1 = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffe00000000
	fe_pow2k(&t1, &t1, 33)

	// Step 300: t0 = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff
	fe_mul(&t0, &t0, &t1)

	// Step 394: t0 = x^0x3fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffbfffffffc00000000000000000000000
	fe_pow2k(&t0, &t0, 94)

	// Step 395: z = x^0x3fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffbfffffffc0000000000000003fffffff
	fe_mul(out1, out1, &t0)

	// Step 397: z = x^0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffeffffffff0000000000000000fffffffc
	fe_pow2k(out1, out1, 2)

	// Ensure that our candidate is actually the square root.
	check, zero: Montgomery_Domain_Field_Element
	fe_square(&check, out1)

	is_valid := fe_equal(&check, &xx)
	fe_cond_select(out1, &zero, out1, is_valid)

	fe_clear_vec([]^Montgomery_Domain_Field_Element{&t0, &t1, &t2, &t3, &xx, &check})

	return is_valid
}

fe_zero :: proc "contextless" (out1: ^Montgomery_Domain_Field_Element) {
	out1[0] = 0
	out1[1] = 0
	out1[2] = 0
	out1[3] = 0
	out1[4] = 0
	out1[5] = 0
}

fe_set :: proc "contextless" (out1, arg1: ^Montgomery_Domain_Field_Element) {
	x1 := arg1[0]
	x2 := arg1[1]
	x3 := arg1[2]
	x4 := arg1[3]
	x5 := arg1[4]
	x6 := arg1[5]
	out1[0] = x1
	out1[1] = x2
	out1[2] = x3
	out1[3] = x4
	out1[4] = x5
	out1[5] = x6
}

@(optimization_mode = "none")
fe_cond_swap :: #force_no_inline proc "contextless" (out1, out2: ^Montgomery_Domain_Field_Element, arg1: int) {
	mask := (u64(arg1) * 0xffffffffffffffff)
	x := (out1[0] ~ out2[0]) & mask
	x1, y1 := out1[0] ~ x, out2[0] ~ x
	x = (out1[1] ~ out2[1]) & mask
	x2, y2 := out1[1] ~ x, out2[1] ~ x
	x = (out1[2] ~ out2[2]) & mask
	x3, y3 := out1[2] ~ x, out2[2] ~ x
	x = (out1[3] ~ out2[3]) & mask
	x4, y4 := out1[3] ~ x, out2[3] ~ x
	x = (out1[4] ~ out2[4]) & mask
	x5, y5 := out1[4] ~ x, out2[4] ~ x
	x = (out1[5] ~ out2[5]) & mask
	x6, y6 := out1[5] ~ x, out2[5] ~ x
	out1[0], out2[0] = x1, y1
	out1[1], out2[1] = x2, y2
	out1[2], out2[2] = x3, y3
	out1[3], out2[3] = x4, y4
	out1[4], out2[4] = x5, y5
	out1[5], out2[5] = x6, y6
}

@(optimization_mode = "none")
fe_cond_select :: #force_no_inline proc "contextless" (
	out1, arg1, arg2: ^Montgomery_Domain_Field_Element,
	arg3: int,
) {
	mask := (u64(arg3) * 0xffffffffffffffff)
	x1 := ((mask & arg2[0]) | ((~mask) & arg1[0]))
	x2 := ((mask & arg2[1]) | ((~mask) & arg1[1]))
	x3 := ((mask & arg2[2]) | ((~mask) & arg1[2]))
	x4 := ((mask & arg2[3]) | ((~mask) & arg1[3]))
	x5 := ((mask & arg2[4]) | ((~mask) & arg1[4]))
	x6 := ((mask & arg2[5]) | ((~mask) & arg1[5]))
	out1[0] = x1
	out1[1] = x2
	out1[2] = x3
	out1[3] = x4
	out1[4] = x5
	out1[5] = x6
}

fe_cond_negate :: proc "contextless" (out1, arg1: ^Montgomery_Domain_Field_Element, ctrl: int) {
	tmp1: Montgomery_Domain_Field_Element = ---
	fe_opp(&tmp1, arg1)
	fe_cond_select(out1, arg1, &tmp1, ctrl)

	fe_clear(&tmp1)
}
