// Package bn256 implements a particular bilinear group.
//
// Bilinear groups are the basis of many of the new cryptographic protocols that
// have been proposed over the past decade. They consist of a triplet of groups
// (G₁, G₂ and GT) such that there exists a function e(g₁ˣ,g₂ʸ)=gTˣʸ (where gₓ
// is a generator of the respective group). That function is called a pairing
// function.
//
// This package specifically implements the Optimal Ate pairing over a 256-bit
// Barreto-Naehrig curve as described in
// http://cryptojedi.org/papers/dclxvi-20100714.pdf. Its output is compatible
// with the implementation described in that paper.
//
// This package previously claimed to operate at a 128-bit security level.
// However, recent improvements in attacks mean that is no longer true. See
// https://moderncrypto.org/mail-archive/curves/2016/000740.html.
package bn256

import (
	"io"
	"math/big"

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/sm/sm9/errors"
)

var one = big.NewInt(1)

// randomK returns a random integer in [1, N-1].
func randomK(r io.Reader) (k *big.Int, err error) {
	b := make([]byte, numBytes)
	n, err := r.Read(b)
	if err != nil {
		return nil, errors.ErrGenerateRandomFailed
	}

	// it is possible that err != nil but n > 0.
	// In this case, we also consider it succeed.
	if n == 0 {
		return nil, errors.ErrGenerateRandomFailed
	}

	// 0 <= k <= N-2
	k = new(big.Int).SetBytes(b)
	if k.Cmp(nMinusOne) >= 0 {
		k.Sub(k, nMinusOne)
	}
	// 1 <= k <= N-1
	k.Add(k, one)
	return k, nil
}

// G1 is an abstract cyclic group. The zero value is suitable for use as the
// output of an operation, but cannot be used as an input.
type G1 struct {
	p curvePoint
}

// UnmarshalCompressed restore e from x and the LSB of y.
func (e *G1) UnmarshalCompressed(x []byte, yBit0 byte) (*G1, error) {
	ex := &gfP{}
	ey := &gfP{}
	ex.Unmarshal(x)
	montEncode(ex, ex)

	// y^2 = x^3 + B
	gfpMul(ey, ex, ex)
	gfpMul(ey, ey, ex)
	gfpAdd(ey, ey, curveB)
	if legendre(ey) != 1 {
		return e, gerrors.WithAnnotating(errors.ErrInvalidInput, "sqrt failed, input bytes are not a valid compressed point")
	}
	ey.Sqrt(ey)
	var temp gfP
	montDecode(&temp, ey)
	if yBit0 != byte(temp[0]&1) {
		gfpNeg(ey, ey)
	}

	e.p.x = *ex
	e.p.y = *ey
	e.p.z = r
	e.p.t = r
	return e, nil
}

// RandomG1 returns x and g₁ˣ where x is a random, non-zero number read from r.
func RandomG1(r io.Reader) (*big.Int, *G1, error) {
	k, err := randomK(r)
	if err != nil {
		return nil, nil, gerrors.WithMessage(err, "RandomG1 failed")
	}

	return k, new(G1).ScalarBaseMult(k), nil
}

// returns the montgemery domain of x
func (e *G1) X() *big.Int {
	e.p.MakeAffine()
	return e.p.x.toBigInt()
}

// returns the montgemery domain of x
func (e *G1) Y() *big.Int {
	e.p.MakeAffine()
	return e.p.y.toBigInt()
}

// returns the montgemery domain of x
func (e *G1) AffineX() *big.Int {
	e.p.MakeAffine()
	var x gfP
	montDecode(&x, &e.p.x)
	return x.toBigInt()
}

// returns the montgemery domain of x
func (e *G1) AffineY() *big.Int {
	e.p.MakeAffine()
	var y gfP
	montDecode(&y, &e.p.y)
	return y.toBigInt()
}

// String returns a readable string representation of e
func (e *G1) String() string {
	return "G1" + e.p.String()
}

// IsInfinity returns if G1 is the infinity point
func (e *G1) IsInfinity() bool {
	return e.p.IsInfinity()
}

// SetInfinity sets e to the infinity point
func (e *G1) SetInfinity() {
	e.p.SetInfinity()
}

// Equal returns if e equals the other point
func (e *G1) Equal(other *G1) bool {
	return e.p.Equal(&other.p)
}

// IsZero return true if e is infinity
func (e *G1) IsZero() bool {
	return e.p.IsInfinity()
}

// IsValid return if e is a valid point of G1
func (e *G1) IsValid() bool {
	return e.p.IsOnCurve()
}

// ScalarBaseMult sets e to [k]g1 where g1 is the generator of the group and returns e.
func (e *G1) ScalarBaseMult(k *big.Int) *G1 {
	e.p.MulBase(k, curverBasePrecompted8)
	return e
}

// ScalarMult sets e to [k]a and then returns e.
func (e *G1) ScalarMult(a *G1, k *big.Int) *G1 {
	e.p.Mul(&a.p, k)
	return e
}

// Add sets e to a+b and then returns e.
func (e *G1) Add(a, b *G1) *G1 {
	e.p.Add(&a.p, &b.p)
	return e
}

// Neg sets e to -a and then returns e.
func (e *G1) Neg(a *G1) *G1 {
	e.p.Neg(&a.p)
	return e
}

// Set sets e to a and then returns e.
func (e *G1) Set(a *G1) *G1 {
	e.p.Set(&a.p)
	return e
}

// FillBytes fills a point in G1 to a byte slice.
//
// The Input slice must be of exactly 64 bytes, i.e., len(b) == 64.
// Also assume the point is not infinity. But if so, then set b to 0s and return.
func (e *G1) FillBytes(b []byte) {
	if e.IsInfinity() {
		for i := 0; i < len(b); i++ {
			b[i] = 0
		}
		return
	}

	e.p.MakeAffine()
	temp := &gfP{}
	montDecode(temp, &e.p.x)
	temp.Marshal(b)
	montDecode(temp, &e.p.y)
	temp.Marshal(b[numBytes:])
}

// Marshal converts a point in G1 to a byte slice of length 64.
// The slice is filled with x || y.
func (e *G1) Marshal() []byte {
	ret := make([]byte, numBytes*2)
	if e.p.IsInfinity() {
		return ret
	}
	e.FillBytes(ret)
	return ret
}

// Unmarshal sets e to the result of converting the output of Marshal back into
// a group element and then returns e.
// The input byte slice m must at least 64 bytes long.
func (e *G1) Unmarshal(m []byte) ([]byte, error) {
	if len(m) < 2*numBytes {
		return m, gerrors.WithAnnotating(errors.ErrInvalidInput, "not enough data to unmarshal to G1")
	}
	e.p.x, e.p.y = gfP{0}, gfP{0}
	e.p.x.Unmarshal(m)
	e.p.y.Unmarshal(m[numBytes:])
	montEncode(&e.p.x, &e.p.x)
	montEncode(&e.p.y, &e.p.y)

	zero := gfP{0}
	if e.p.x == zero && e.p.y == zero {
		// The point at infinity.
		e.p.y = gfPOne
		e.p.z = gfP{0}
		e.p.t = gfP{0}
	} else {
		e.p.z = gfPOne
		e.p.t = gfPOne

		if !e.p.IsOnCurve() {
			return nil, gerrors.WithAnnotating(errors.ErrInvalidPoint, "point is not a valid point on curve")
		}
	}

	return m[2*numBytes:], nil
}

// G2 is an abstract cyclic group. The zero value is suitable for use as the
// output of an operation, but cannot be used as an input.
type G2 struct {
	p twistPoint
}

// RandomG2 returns x and g₂ˣ where x is a random, non-zero number read from r.
func RandomG2(r io.Reader) (*big.Int, *G2, error) {
	k, err := randomK(r)
	if err != nil {
		return nil, nil, gerrors.WithMessage(err, "RandomG2 failed")
	}

	return k, new(G2).ScalarBaseMult(k), nil
}

func (e *G2) X() (*big.Int, *big.Int) {
	e.p.MakeAffine()
	return e.p.x.x.toBigInt(), e.p.x.y.toBigInt()
}

func (e *G2) Y() (*big.Int, *big.Int) {
	e.p.MakeAffine()
	return e.p.y.x.toBigInt(), e.p.y.y.toBigInt()
}

func (e *G2) IsInfinity() bool {
	return e.p.IsInfinity()
}

func (e *G2) SetInfinity() {
	e.p.SetInfinity()
}
func (e *G2) Equal(other *G2) bool {
	return e.p.Equal(&other.p)
}

// FromX restore e from x and the LSB of y. yBit0 can only be 0 or 1.
func (e *G2) UnmarshalCompressed(x0, x1 []byte, yBit0 byte) (*G2, error) {
	if len(x0) < numBytes || len(x1) < numBytes {
		return nil, gerrors.WithAnnotating(errors.ErrInvalidInput, "point is not a valid point on curve")
	}

	if yBit0&byte(0xfe) != 0 {
		return nil, gerrors.WithAnnotatingf(errors.ErrInvalidInput, "yBit0 can only be 0 or 1, but it's %d", yBit0)
	}

	ex := &gfP2{}
	ey := &gfP2{}
	ex0 := &ex.x
	ex1 := &ex.y

	ex0.Unmarshal(x0)
	montEncode(ex0, ex0)
	ex1.Unmarshal(x1)
	montEncode(ex1, ex1)

	ey.Mul(ex, ex)
	ey.Mul(ey, ex)
	ey.Add(ey, twistB) // ey = x^3 + B

	if !ey.Sqrt(ey) {
		return e, gerrors.WithAnnotatingf(errors.ErrInvalidPoint, "sqrt failed, input bytes are not a valid compressed point")
	}
	tmp := &gfP{}
	montDecode(tmp, &ey.y)
	if yBit0 != byte(tmp[0]&1) {
		ey.Neg(ey)
	}

	e.p.x = *ex
	e.p.y = *ey
	e.p.t = gfP2{y: r}
	e.p.z = gfP2{y: r}
	return e, nil
}

func (e *G2) String() string {
	return "G2" + e.p.String()
}

// ScalarBaseMult sets e to [k]g where g is the generator of the group and then
// returns out.
func (e *G2) ScalarBaseMult(k *big.Int) *G2 {
	e.p.MulBase(k, twistBasePrecomputed8)
	return e
}

// ScalarMult sets e to [k]a and then returns e.
func (e *G2) ScalarMult(a *G2, k *big.Int) *G2 {
	e.p.Mul(&a.p, k)
	return e
}

// Add sets e to a+b and then returns e.
func (e *G2) Add(a, b *G2) *G2 {
	e.p.Add(&a.p, &b.p)
	return e
}

// Neg sets e to -a and then returns e.
func (e *G2) Neg(a *G2) *G2 {
	e.p.Neg(&a.p)
	return e
}

// Set sets e to a and then returns e.
func (e *G2) Set(a *G2) *G2 {
	e.p.Set(&a.p)
	return e
}

// FillBytes fills a point in G2 to a byte slice.
//
// The Input slice must be of exactly 128 bytes, i.e., len(b) == 128.
// Also assume the point is not infinity. But if so, then set b to 0s and return.
func (e *G2) FillBytes(b []byte) {
	if e.IsInfinity() {
		for i := 0; i < len(b); i++ {
			b[i] = 0
		}
		return
	}
	e.p.MakeAffine()
	temp := &gfP{}
	montDecode(temp, &e.p.x.x)
	temp.Marshal(b)
	montDecode(temp, &e.p.x.y)
	temp.Marshal(b[numBytes:])
	montDecode(temp, &e.p.y.x)
	temp.Marshal(b[2*numBytes:])
	montDecode(temp, &e.p.y.y)
	temp.Marshal(b[3*numBytes:])
}

// Marshal converts e into a byte slice.
func (e *G2) Marshal() []byte {
	ret := make([]byte, numBytes*4)
	if e.p.IsInfinity() {
		return ret
	}
	e.FillBytes(ret)
	return ret
}

// Unmarshal sets e to the result of converting the output of Marshal back into
// a group element and then returns e.
func (e *G2) Unmarshal(m []byte) ([]byte, error) {
	if len(m) < 4*numBytes {
		return m, gerrors.WithAnnotatingf(errors.ErrInvalidInput, "not enough data to unmarshal to G1")
	}

	e.p.x.x.Unmarshal(m)
	e.p.x.y.Unmarshal(m[numBytes:])
	e.p.y.x.Unmarshal(m[2*numBytes:])
	e.p.y.y.Unmarshal(m[3*numBytes:])
	montEncode(&e.p.x.x, &e.p.x.x)
	montEncode(&e.p.x.y, &e.p.x.y)
	montEncode(&e.p.y.x, &e.p.y.x)
	montEncode(&e.p.y.y, &e.p.y.y)

	if e.p.x.IsZero() && e.p.y.IsZero() {
		// This is the point at infinity.
		e.p.y.SetOne()
		e.p.z.SetZero()
		e.p.t.SetZero()
	} else {
		e.p.z.SetOne()
		e.p.t.SetOne()

		if !e.p.IsOnCurve() {
			return m, gerrors.WithAnnotatingf(errors.ErrInvalidPoint, "unmarshaled point is not a valid point on curve")
		}
	}

	return m[4*numBytes:], nil
}

// GT is an abstract cyclic group. The zero value is suitable for use as the
// output of an operation, but cannot be used as an input.
type GT struct {
	p gfP12
}

func (e *GT) Equal(other *GT) bool {
	return e.p.Equal(&other.p)
}

// Order12412 change e to 1-2-4-12 field extension represent
func (e *GT) Order12412() {
	e.p.x.y.x, e.p.x.y.y, e.p.y.y.x, e.p.y.y.y, e.p.x.z.x, e.p.x.z.y, e.p.y.x.x, e.p.y.x.y =
		e.p.y.y.x, e.p.y.y.y, e.p.x.y.x, e.p.x.y.y, e.p.y.x.x, e.p.y.x.y, e.p.x.z.x, e.p.x.z.y
}

// RandomGT returns x and e(g₁, g₂)ˣ where x is a random, non-zero number read
// from r.
func RandomGT(r io.Reader) (*big.Int, *GT, error) {
	k, err := randomK(r)
	if err != nil {
		return nil, nil, gerrors.WithMessage(err, "RandomGT failed")
	}

	return k, new(GT).ScalarBaseMult(k), nil
}

// Pair calculates an Optimal Ate pairing.
func Pair(g1 *G1, g2 *G2) *GT {
	var e GT
	optimalAte(&e.p, &g2.p, &g1.p)
	return &e
}

// Pair calculates an Optimal Ate pairing.
func PairLol(e *GT, g1 *G1, g2 *G2) {
	optimalAte(&e.p, &g2.p, &g1.p)
}

// Miller applies Miller's algorithm, which is a bilinear function from the
// source groups to F_p^12. Miller(g1, g2).Finalize() is equivalent to Pair(g1,
// g2).
func Miller(g1 *G1, g2 *G2) *GT {
	var e GT
	miller(&e.p, &g2.p, &g1.p)
	return &e
}

func (e *GT) String() string {
	p := *e
	p.Order12412()
	return "GT" + gfP12Decode(&p.p).String()
}

// ScalarBaseMult sets e to g*k where g is the generator of the group and then
// returns out.
func (e *GT) ScalarBaseMult(k *big.Int) *GT {
	if useLattice {
		e.p.latticeExp(gfP12Gen, k)
		return e
	} else {
		return e.ScalarMultSimple(&GT{*gfP12Gen}, k)
	}
}

// ScalarMult sets e to a*k and then returns e. (If e is not guaranteed to be an element of the group because it is the
// output of Miller(), use ScalarMultSimple.)
func (e *GT) ScalarMult(a *GT, k *big.Int) *GT {
	if useLattice {
		e.p.latticeExp(&a.p, k)
		return e
	} else {
		return e.ScalarMultSimple(a, k)
	}
}

// ScalarMultSimple sets e to a*k and then returns e.
func (e *GT) ScalarMultSimple(a *GT, k *big.Int) *GT {
	e.p.Exp(&a.p, k)
	return e
}

func (e *GT) Mul(a, b *GT) *GT {
	e.p.Mul(&a.p, &b.p)
	return e
}

// Add sets e to a+b and then returns e.
func (e *GT) Add(a, b *GT) *GT {
	e.p.Mul(&a.p, &b.p)
	return e
}

// Neg sets e to -a and then returns e.
func (e *GT) Neg(a *GT) *GT {
	e.p.Conjugate(&a.p)
	return e
}

// Set sets e to a and then returns e.
func (e *GT) Set(a *GT) *GT {
	e.p.Set(&a.p)
	return e
}

// Set sets e to a and then returns e.
func (e *GT) SetOne() *GT {
	e.p.SetOne()
	return e
}

// Set sets e to a and then returns e.
func (e *GT) Invert(a *GT) *GT {
	e.p.Invert(&a.p)
	return e
}

// Finalize is a linear function from F_p^12 to GT.
// func (e *GT) Finalize() *GT {
// 	finalExponentiation(&e.p, &e.p)
// 	return e
// }

// Marshal converts e into a byte slice.
func (e *GT) Marshal() []byte {
	p := *e
	p.Order12412()

	ret := make([]byte, numBytes*12)
	temp := &gfP{}

	montDecode(temp, &p.p.x.x.x)
	temp.Marshal(ret)
	montDecode(temp, &p.p.x.x.y)
	temp.Marshal(ret[numBytes:])
	montDecode(temp, &p.p.x.y.x)
	temp.Marshal(ret[2*numBytes:])
	montDecode(temp, &p.p.x.y.y)
	temp.Marshal(ret[3*numBytes:])
	montDecode(temp, &p.p.x.z.x)
	temp.Marshal(ret[4*numBytes:])
	montDecode(temp, &p.p.x.z.y)
	temp.Marshal(ret[5*numBytes:])
	montDecode(temp, &p.p.y.x.x)
	temp.Marshal(ret[6*numBytes:])
	montDecode(temp, &p.p.y.x.y)
	temp.Marshal(ret[7*numBytes:])
	montDecode(temp, &p.p.y.y.x)
	temp.Marshal(ret[8*numBytes:])
	montDecode(temp, &p.p.y.y.y)
	temp.Marshal(ret[9*numBytes:])
	montDecode(temp, &p.p.y.z.x)
	temp.Marshal(ret[10*numBytes:])
	montDecode(temp, &p.p.y.z.y)
	temp.Marshal(ret[11*numBytes:])

	return ret
}

// Unmarshal sets e to the result of converting the output of Marshal back into
// a group element and then returns e.
func (e *GT) Unmarshal(m []byte) ([]byte, error) {
	if len(m) < 12*numBytes {
		return m, gerrors.WithAnnotating(errors.ErrInvalidInput, "not enough data to unmarshal to GT")
	}

	e.p.x.x.x.Unmarshal(m)
	e.p.x.x.y.Unmarshal(m[numBytes:])
	e.p.x.y.x.Unmarshal(m[2*numBytes:])
	e.p.x.y.y.Unmarshal(m[3*numBytes:])
	e.p.x.z.x.Unmarshal(m[4*numBytes:])
	e.p.x.z.y.Unmarshal(m[5*numBytes:])
	e.p.y.x.x.Unmarshal(m[6*numBytes:])
	e.p.y.x.y.Unmarshal(m[7*numBytes:])
	e.p.y.y.x.Unmarshal(m[8*numBytes:])
	e.p.y.y.y.Unmarshal(m[9*numBytes:])
	e.p.y.z.x.Unmarshal(m[10*numBytes:])
	e.p.y.z.y.Unmarshal(m[11*numBytes:])
	montEncode(&e.p.x.x.x, &e.p.x.x.x)
	montEncode(&e.p.x.x.y, &e.p.x.x.y)
	montEncode(&e.p.x.y.x, &e.p.x.y.x)
	montEncode(&e.p.x.y.y, &e.p.x.y.y)
	montEncode(&e.p.x.z.x, &e.p.x.z.x)
	montEncode(&e.p.x.z.y, &e.p.x.z.y)
	montEncode(&e.p.y.x.x, &e.p.y.x.x)
	montEncode(&e.p.y.x.y, &e.p.y.x.y)
	montEncode(&e.p.y.y.x, &e.p.y.y.x)
	montEncode(&e.p.y.y.y, &e.p.y.y.y)
	montEncode(&e.p.y.z.x, &e.p.y.z.x)
	montEncode(&e.p.y.z.y, &e.p.y.z.y)

	e.Order12412()
	return m[12*numBytes:], nil
}
