package bn256

import (
	"encoding/binary"
	"fmt"
	"math/big"

	"golang.org/x/crypto/hkdf"
	"xdx.jelly/xgcl/sm/sm3"
)

// gfP is the finite fileds GF(p), little-endian and Montgomery reprent.
type gfP [4]uint64

// newGFp return the Montgomery encoded of x, i.e., out = x * R mod p
func newGFp(x int64) (out *gfP) {
	if x >= 0 {
		out = &gfP{uint64(x)}
	} else {
		out = &gfP{uint64(-x)}
		gfpNeg(out, out)
	}

	montEncode(out, out)
	return out
}

// bug: the output order
func (e *gfP) toBigInt() *big.Int {
	// r := new(big.Int).SetUint64(e[0])
	// r.Lsh(r, 64).Add(r, new(big.Int).SetUint64(e[1]))
	// r.Lsh(r, 64).Add(r, new(big.Int).SetUint64(e[2]))
	// r.Lsh(r, 64).Add(r, new(big.Int).SetUint64(e[3]))

	// should be:

	r := new(big.Int).SetUint64(e[3])
	r.Lsh(r, 64).Add(r, new(big.Int).SetUint64(e[2]))
	r.Lsh(r, 64).Add(r, new(big.Int).SetUint64(e[1]))
	r.Lsh(r, 64).Add(r, new(big.Int).SetUint64(e[0]))

	return r
}

// hashToBase implements hashing a message to an element of the field.
//
// L = ceil((256+128)/8)=48, ctr = 0, i = 1
//
// Although we do not need it in SM9.
func hashToBase(msg, dst []byte) *gfP { // nolint
	var t [48]byte
	info := []byte{'H', '2', 'C', byte(0), byte(1)}
	// sha256 or sm3? its a question.
	// r := hkdf.New(sha256.New, msg, dst, info)
	r := hkdf.New(sm3.New, msg, dst, info)
	if _, err := r.Read(t[:]); err != nil {
		panic(err)
	}
	var x big.Int
	v := x.SetBytes(t[:]).Mod(&x, P).Bytes()
	v32 := [32]byte{}
	for i := len(v) - 1; i >= 0; i-- {
		v32[len(v)-1-i] = v[i]
	}
	u := &gfP{
		binary.LittleEndian.Uint64(v32[0*8 : 1*8]),
		binary.LittleEndian.Uint64(v32[1*8 : 2*8]),
		binary.LittleEndian.Uint64(v32[2*8 : 3*8]),
		binary.LittleEndian.Uint64(v32[3*8 : 4*8]),
	}
	montEncode(u, u)
	return u
}

// String return the hex of e.
//
// Note e is Montgomery encoded, decode it befor printing if you want print the origin value.
func (e *gfP) String() string {
	return fmt.Sprintf("%16.16X%16.16X%16.16X%16.16X", e[3], e[2], e[1], e[0])
}

func (e *gfP) Equal(other *gfP) bool {
	return *e == *other
}

// Set sets e to f.
func (e *gfP) Set(f *gfP) {
	e[0] = f[0]
	e[1] = f[1]
	e[2] = f[2]
	e[3] = f[3]
}

// exp compute e = f^bits with naive square-mul method.
// Input bits is reprents as little-endian.
func (e *gfP) exp(f *gfP, bits [4]uint64) {
	sum, power := &gfP{}, &gfP{}

	sum[0], sum[1], sum[2], sum[3] = r[0], r[1], r[2], r[3] // sum = 1
	power.Set(f)

	for word := 0; word < 4; word++ {
		for bit := uint(0); bit < 64; bit++ {
			if (bits[word]>>bit)&1 == 1 {
				gfpMul(sum, sum, power)
			}
			gfpMul(power, power, power)
		}
	}
	e.Set(sum)
}

// Invert set e to f^{-1}, by Farmat's little theorem: a^{-1} = a^{p-2} mod p
//
// If input f is 0, then e is 0 after return.
func (e *gfP) Invert(f *gfP) {
	e.exp(f, pMinus2)
}

// half set e to f/2. Use shift to void mul 1/2 mod p
func (e *gfP) half(f *gfP) {
	sign := f[0] & 1
	if sign == 1 {
		gfpNeg(e, f)
	} else {
		if e != f {
			*e = *f
		}
	}
	// e = e >> 1
	e[0] = (e[0] >> 1) | (e[1] << 63)
	e[1] = (e[1] >> 1) | (e[2] << 63)
	e[2] = (e[2] >> 1) | (e[3] << 63)
	e[3] >>= 1
	if sign == 1 {
		gfpNeg(e, e)
	}
}

// Sqrt set e to be the square root of f if f is a square.
// If f is not a square, the result is undefined.
func (e *gfP) Sqrt(f *gfP) {
	// Since P = 8u+5, then:
	// if f^{2u+1} = 1, e = f^(u+1)
	// or
	// f^{2u+1} = -1, e = f^(u+1) * sqrt(-1).
	if false {
		// If we do not care side channel attack, we just square f^(u+1) and compare with f.
		// if f is not a square, then f is undefined
		root := &gfP{}
		ff := &gfP{}
		root.exp(f, pPlus3Over8) // sum = f^{u+1}
		gfpMul(ff, root, root)   // power = f^{2u+2}
		if *ff == *f {
			e.Set(root)
		} else {
			gfpMul(e, root, sqrtRootOfMinus1ModP)
		}
	} else {
		// constant time, but extra 2 mul.
		f2u1 := &gfP{}
		fu1 := &gfP{}
		fu1s1 := &gfP{}

		fu1s1.exp(f, pMinus5Over8)               // fu = f^u
		gfpMul(fu1, fu1s1, f)                    // fu1 = f^{u+1}
		gfpMul(f2u1, fu1s1, fu1)                 // g=f^{2u+1}
		gfpMul(fu1s1, fu1, sqrtRootOfMinus1ModP) // fu = f^{u+1} * sqrt(-1)

		// g must be -1, 0 or 1 if f is a square.
		switch {
		case *f2u1 == gfP(r):
			e.Set(fu1)
		case *f2u1 == gfP(nr):
			e.Set(fu1s1)
		case *f2u1 == gfP{}:
			e[0] = 0
			e[1] = 0
			e[2] = 0
			e[3] = 0
		default:
			// f is not a square, e unchange

		}
	}

}

// Marshal marshal e to bytes
func (e *gfP) Marshal(out []byte) {
	for w := uint(0); w < 4; w++ {
		for b := uint(0); b < 8; b++ {
			out[8*w+b] = byte(e[3-w] >> (56 - 8*b))
		}
	}
}

// Unmarshal restore e from bytes
func (e *gfP) Unmarshal(in []byte) {
	for w := uint(0); w < 4; w++ {
		e[3-w] = 0
		for b := uint(0); b < 8; b++ {
			e[3-w] += uint64(in[8*w+b]) << (56 - 8*b)
		}
	}
}

// montEncode set c to a's Montgomery reprent, c = a*R mod p
func montEncode(c, a *gfP) {
	gfpMul(c, a, r2)
}

// montDecode a Montgomery reprent a, c = a*R^-1 mod p
func montDecode(c, a *gfP) {
	gfpMul(c, a, &gfP{1})
}

// sign0 returns the sign of e - (p-1)/2
func sign0(e *gfP) int { // nolint
	x := &gfP{}
	montDecode(x, e)
	for w := 3; w >= 0; w-- {
		if x[w] > pMinus1Over2[w] {
			return 1
		} else if x[w] < pMinus1Over2[w] {
			return -1
		}
	}
	return 1
}

// legendre return the legendre symbol of e
func legendre(e *gfP) int {
	f := &gfP{}
	// Since P = 4k+3, then e^(2k+1) is the Legendre symbol of e.
	f.exp(e, pMinus1Over2)

	montDecode(f, f)

	if *f != [4]uint64{} {
		return 2*int(f[0]&1) - 1
	}

	return 0
}