package sm9

import (
	_subtle "crypto/subtle"
	"encoding/binary"
	"errors"
	"io"
	"math/bits"
	"sync"

	"slices"

	"github.com/emmansun/gmsm/internal/bigmod"
	"github.com/emmansun/gmsm/internal/randutil"
	"github.com/emmansun/gmsm/internal/sm9/bn256"
	"github.com/emmansun/gmsm/internal/subtle"
)

// SignMasterPrivateKey is a signature master private key, generated by KGC
type SignMasterPrivateKey struct {
	*SignMasterPublicKey        // master public key
	privateKey           []byte // master private key
}

// SignMasterPublicKey is a signature master public key, generated by KGC
type SignMasterPublicKey struct {
	MasterPublicKey *bn256.G2 // master public key
	pairOnce        sync.Once
	basePoint       *bn256.GT // the result of Pair(Gen1, pub.MasterPublicKey)
	tableGenOnce    sync.Once
	table           *[32 * 2]bn256.GTFieldTable // precomputed basePoint^n
}

// SignPrivateKey is a signature private key, generated by KGC
type SignPrivateKey struct {
	PrivateKey           *bn256.G1 // user private key
	*SignMasterPublicKey           // master public key
}

// EncryptMasterPrivateKey is an encryption master private key, generated by KGC
type EncryptMasterPrivateKey struct {
	*EncryptMasterPublicKey        // master public key
	privateKey              []byte // master private key
}

// EncryptMasterPublicKey is an encryption master public key, generated by KGC
type EncryptMasterPublicKey struct {
	MasterPublicKey *bn256.G1 // public key
	pairOnce        sync.Once
	basePoint       *bn256.GT // the result of Pair(pub.MasterPublicKey, Gen2)
	tableGenOnce    sync.Once
	table           *[32 * 2]bn256.GTFieldTable // precomputed basePoint^n
}

// EncryptPrivateKey is an encryption private key, generated by KGC
type EncryptPrivateKey struct {
	PrivateKey              *bn256.G2 // user private key
	*EncryptMasterPublicKey           // master public key
}

// GenerateSignMasterKey generates a signature master key pair for DSA usage.
func GenerateSignMasterKey(rand io.Reader) (*SignMasterPrivateKey, error) {
	key := make([]byte, len(bn256.OrderMinus1Bytes))
	randutil.MaybeReadByte(rand)

	for {
		if _, err := io.ReadFull(rand, key); err != nil {
			return nil, err
		}
		// In tests, rand will return all zeros and NewPrivateKey will reject
		// the zero key as it generates the identity as a public key. This also
		// makes this function consistent with crypto/elliptic.GenerateKey.
		key[1] ^= 0x42

		k, err := NewSignMasterPrivateKey(key)
		if err == errInvalidPrivateKey {
			continue
		}
		return k, err
	}
}

// NewSignMasterPrivateKey creates a new SignMasterPrivateKey from the given byte slice.
// The provided key must have the same length as bn256.OrderMinus1Bytes and must be less than bn256.OrderMinus1Bytes.
// If the key is invalid, an error is returned.
//
// Parameters:
//   - key: A byte slice representing the master private key.
//
// Returns:
//   - *SignMasterPrivateKey: A pointer to the newly created SignMasterPrivateKey.
//   - error: An error if the key is invalid or if there is an issue during key generation.
func NewSignMasterPrivateKey(key []byte) (*SignMasterPrivateKey, error) {
	if len(key) > len(bn256.OrderMinus1Bytes) {
		return nil, errInvalidPrivateKey
	}
	key = bn256.NormalizeScalar(key)
	if subtle.ConstantTimeAllZero(key) == 1 || !isLess(key, bn256.OrderMinus1Bytes) {
		return nil, errInvalidPrivateKey
	}
	p, err := new(bn256.G2).ScalarBaseMult(key)
	if err != nil {
		return nil, err
	}
	priv := new(SignMasterPrivateKey)
	priv.privateKey = slices.Clone(key)
	priv.SignMasterPublicKey = new(SignMasterPublicKey)
	priv.MasterPublicKey = p
	return priv, nil
}

// Equal compares the receiver SignMasterPrivateKey with another SignMasterPrivateKey x.
// It returns true if both the MasterPublicKey and privateKey fields are equal, using
// constant time comparison for the privateKey to prevent timing attacks.
func (master *SignMasterPrivateKey) Equal(x *SignMasterPrivateKey) bool {
	return master.SignMasterPublicKey.Equal(x.SignMasterPublicKey) && _subtle.ConstantTimeCompare(master.privateKey, x.privateKey) == 1
}

// Bytes returns the byte representation of the SignMasterPrivateKey.
// It creates a new byte slice and appends the private key bytes to it,
// ensuring that the original private key data is not modified.
func (master *SignMasterPrivateKey) Bytes() []byte {
	return slices.Clone(master.privateKey)
}

// GenerateUserKey generate an user dsa key.
func (master *SignMasterPrivateKey) GenerateUserKey(uid []byte, hid byte) (*SignPrivateKey, error) {
	var id []byte
	id = append(append(id, uid...), hid)

	t1Nat := hashH1(id)

	d, err := bigmod.NewNat().SetBytes(master.privateKey, orderNat)
	if err != nil {
		return nil, err
	}

	t1Nat.Add(d, orderNat)
	if t1Nat.IsZero() == 1 {
		return nil, errors.New("sm9: need to re-generate signature master private key")
	}

	t1Nat = bigmod.NewNat().Exp(t1Nat, bn256.OrderMinus2Bytes, orderNat)
	t1Nat.Mul(d, orderNat)

	priv := new(SignPrivateKey)
	priv.SignMasterPublicKey = master.SignMasterPublicKey
	g1, err := new(bn256.G1).ScalarBaseMult(t1Nat.Bytes(orderNat))
	if err != nil {
		return nil, err
	}
	priv.PrivateKey = g1

	return priv, nil
}

// Public returns the public key corresponding to the private key.
func (master *SignMasterPrivateKey) Public() *SignMasterPublicKey {
	return master.SignMasterPublicKey
}

// Equal compares the receiver SignMasterPublicKey with another SignMasterPublicKey
// and returns true if they are equal, otherwise it returns false.
func (pub *SignMasterPublicKey) Equal(x *SignMasterPublicKey) bool {
	pubBytes := pub.MasterPublicKey.MarshalUncompressed()
	xBytes := x.MasterPublicKey.MarshalUncompressed()
	return _subtle.ConstantTimeCompare(pubBytes, xBytes) == 1
}

// Bytes returns the byte representation of the SignMasterPublicKey
// by marshaling the MasterPublicKey in an uncompressed format.
func (pub *SignMasterPublicKey) Bytes() []byte {
	return pub.MasterPublicKey.MarshalUncompressed()
}

// pair generate the basepoint once
func (pub *SignMasterPublicKey) pair() *bn256.GT {
	pub.pairOnce.Do(func() {
		pub.basePoint = bn256.Pair(bn256.Gen1, pub.MasterPublicKey)
	})
	return pub.basePoint
}

func (pub *SignMasterPublicKey) generatorTable() *[32 * 2]bn256.GTFieldTable {
	pub.tableGenOnce.Do(func() {
		pub.table = bn256.GenerateGTFieldTable(pub.pair())
	})
	return pub.table
}

// ScalarBaseMult compute basepoint^r with precomputed table
// The base point = pair(Gen1, <master public key>)
func (pub *SignMasterPublicKey) ScalarBaseMult(scalar []byte) (*bn256.GT, error) {
	tables := pub.generatorTable()
	return bn256.ScalarBaseMultGT(tables, scalar)
}

// GenerateUserPublicKey generate a signature public key for given user.
func (pub *SignMasterPublicKey) GenerateUserPublicKey(uid []byte, hid byte) *bn256.G2 {
	var buffer []byte
	buffer = append(append(buffer, uid...), hid)
	h1 := hashH1(buffer)
	p, err := new(bn256.G2).ScalarBaseMult(h1.Bytes(orderNat))
	if err != nil {
		panic(err)
	}
	p.Add(p, pub.MasterPublicKey)
	return p
}

// Equal compares the SignPrivateKey receiver with another SignPrivateKey x
// and returns true if they are equal, otherwise it returns false.
func (priv *SignPrivateKey) Equal(x *SignPrivateKey) bool {
	return priv.PrivateKey.Equal(x.PrivateKey)
}

// Bytes returns the byte representation of the SignPrivateKey.
// It marshals the private key in an uncompressed format.
func (priv *SignPrivateKey) Bytes() []byte {
	return priv.PrivateKey.MarshalUncompressed()
}

// MasterPublic returns the master public key corresponding to the private key.
func (priv *SignPrivateKey) MasterPublic() *SignMasterPublicKey {
	return priv.SignMasterPublicKey
}

// SetMasterPublicKey bind the signature master public key to it.
func (priv *SignPrivateKey) SetMasterPublicKey(pub *SignMasterPublicKey) {
	if priv.SignMasterPublicKey == nil || priv.SignMasterPublicKey.MasterPublicKey == nil {
		priv.SignMasterPublicKey = pub
	}
}

func unmarshalG2(bytes []byte) (*bn256.G2, error) {
	g2 := new(bn256.G2)
	switch bytes[0] {
	case 4:
		_, err := g2.Unmarshal(bytes[1:])
		if err != nil {
			return nil, err
		}
	case 2, 3:
		_, err := g2.UnmarshalCompressed(bytes)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("sm9: invalid point identity byte")
	}
	return g2, nil
}

// UnmarshalRaw unmarsal raw bytes data to signature master public key
func (pub *SignMasterPublicKey) UnmarshalRaw(bytes []byte) error {
	g2, err := unmarshalG2(bytes)
	if err != nil {
		return err
	}
	pub.MasterPublicKey = g2
	return nil
}

func unmarshalG1(bytes []byte) (*bn256.G1, error) {
	g := new(bn256.G1)
	switch bytes[0] {
	case 4:
		_, err := g.Unmarshal(bytes[1:])
		if err != nil {
			return nil, err
		}
	case 2, 3:
		_, err := g.UnmarshalCompressed(bytes)
		if err != nil {
			return nil, err
		}
	default:
		return nil, errors.New("sm9: invalid point encoding")
	}
	return g, nil
}

// UnmarshalRaw unmarsal raw bytes data to the signature private key
// Note, priv's SignMasterPublicKey should be handled separately.
func (priv *SignPrivateKey) UnmarshalRaw(bytes []byte) error {
	g, err := unmarshalG1(bytes)
	if err != nil {
		return err
	}
	priv.PrivateKey = g
	return nil
}

// GenerateEncryptMasterKey generates an encryption master key pair.
func GenerateEncryptMasterKey(rand io.Reader) (*EncryptMasterPrivateKey, error) {
	key := make([]byte, len(bn256.OrderMinus1Bytes))
	randutil.MaybeReadByte(rand)

	for {
		if _, err := io.ReadFull(rand, key); err != nil {
			return nil, err
		}
		// In tests, rand will return all zeros and NewPrivateKey will reject
		// the zero key as it generates the identity as a public key. This also
		// makes this function consistent with crypto/elliptic.GenerateKey.
		key[1] ^= 0x42

		k, err := NewEncryptMasterPrivateKey(key)
		if err == errInvalidPrivateKey {
			continue
		}
		return k, err
	}
}

// NewEncryptMasterPrivateKey creates a new EncryptMasterPrivateKey from the given key bytes.
// The key must have a length equal to bn256.OrderMinus1Bytes. If the key is all zeros or not
// less than bn256.OrderMinus1Bytes, an error is returned.
//
// Parameters:
//   - key: A byte slice representing the master private key.
//
// Returns:
//   - *EncryptMasterPrivateKey: A pointer to the newly created EncryptMasterPrivateKey.
//   - error: An error if the key is invalid or if there is an issue during key generation.
func NewEncryptMasterPrivateKey(key []byte) (*EncryptMasterPrivateKey, error) {
	if len(key) > len(bn256.OrderMinus1Bytes) {
		return nil, errInvalidPrivateKey
	}
	key = bn256.NormalizeScalar(key)
	if subtle.ConstantTimeAllZero(key) == 1 || !isLess(key, bn256.OrderMinus1Bytes) {
		return nil, errInvalidPrivateKey
	}
	p, err := new(bn256.G1).ScalarBaseMult(key)
	if err != nil {
		return nil, err
	}
	priv := new(EncryptMasterPrivateKey)
	priv.privateKey = slices.Clone(key)
	priv.EncryptMasterPublicKey = new(EncryptMasterPublicKey)
	priv.MasterPublicKey = p
	return priv, nil
}

// Bytes returns a copy of the private key bytes from the EncryptMasterPrivateKey.
// This method ensures that the original private key data is not modified by
// returning a new slice containing the same data.
func (master *EncryptMasterPrivateKey) Bytes() []byte {
	return slices.Clone(master.privateKey)
}

// Equal compares the receiver EncryptMasterPrivateKey with another EncryptMasterPrivateKey x.
// It returns true if both the public keys and private keys are equal, otherwise it returns false.
// The comparison of private keys is done in constant time to prevent timing attacks.
func (master *EncryptMasterPrivateKey) Equal(x *EncryptMasterPrivateKey) bool {
	return master.EncryptMasterPublicKey.Equal(x.EncryptMasterPublicKey) && _subtle.ConstantTimeCompare(master.privateKey, x.privateKey) == 1
}

// GenerateUserKey generate an encryption private key for the given user.
func (master *EncryptMasterPrivateKey) GenerateUserKey(uid []byte, hid byte) (*EncryptPrivateKey, error) {
	var id []byte
	id = append(append(id, uid...), hid)

	t1Nat := hashH1(id)

	d, err := bigmod.NewNat().SetBytes(master.privateKey, orderNat)
	if err != nil {
		return nil, err
	}

	t1Nat.Add(d, orderNat)
	if t1Nat.IsZero() == 1 {
		return nil, errors.New("sm9: need to re-generate encryption master private key")
	}

	t1Nat = bigmod.NewNat().Exp(t1Nat, bn256.OrderMinus2Bytes, orderNat)
	t1Nat.Mul(d, orderNat)

	priv := new(EncryptPrivateKey)
	priv.EncryptMasterPublicKey = master.EncryptMasterPublicKey
	p, err := new(bn256.G2).ScalarBaseMult(t1Nat.Bytes(orderNat))
	if err != nil {
		panic(err)
	}
	priv.PrivateKey = p

	return priv, nil
}

// Public returns the public key corresponding to the private key.
func (master *EncryptMasterPrivateKey) Public() *EncryptMasterPublicKey {
	return master.EncryptMasterPublicKey
}

// Equal compares the receiver EncryptMasterPublicKey with another EncryptMasterPublicKey
// and returns true if they are equal, otherwise false.
func (pub *EncryptMasterPublicKey) Equal(x *EncryptMasterPublicKey) bool {
	pubBytes := pub.MasterPublicKey.MarshalUncompressed()
	xBytes := x.MasterPublicKey.MarshalUncompressed()
	return _subtle.ConstantTimeCompare(pubBytes, xBytes) == 1
}

func (pub *EncryptMasterPublicKey) Bytes() []byte {
	return pub.MasterPublicKey.MarshalUncompressed()
}

// pair generate the basepoint once
func (pub *EncryptMasterPublicKey) pair() *bn256.GT {
	pub.pairOnce.Do(func() {
		pub.basePoint = bn256.Pair(pub.MasterPublicKey, bn256.Gen2)
	})
	return pub.basePoint
}

func (pub *EncryptMasterPublicKey) generatorTable() *[32 * 2]bn256.GTFieldTable {
	pub.tableGenOnce.Do(func() {
		pub.table = bn256.GenerateGTFieldTable(pub.pair())
	})
	return pub.table
}

// ScalarBaseMult compute basepoint^r with precomputed table.
// The base point = pair(<master public key>, Gen2)
func (pub *EncryptMasterPublicKey) ScalarBaseMult(scalar []byte) (*bn256.GT, error) {
	tables := pub.generatorTable()
	return bn256.ScalarBaseMultGT(tables, scalar)
}

// GenerateUserPublicKey generate an encryption public key for the given user.
func (pub *EncryptMasterPublicKey) GenerateUserPublicKey(uid []byte, hid byte) *bn256.G1 {
	var buffer []byte
	buffer = append(append(buffer, uid...), hid)
	h1 := hashH1(buffer)
	p, err := new(bn256.G1).ScalarBaseMult(h1.Bytes(orderNat))
	if err != nil {
		panic(err)
	}
	p.Add(p, pub.MasterPublicKey)
	return p
}

// Equal compares the calling EncryptPrivateKey with another EncryptPrivateKey
// and returns true if they are equal, otherwise it returns false.
func (priv *EncryptPrivateKey) Equal(x *EncryptPrivateKey) bool {
	return priv.PrivateKey.Equal(x.PrivateKey)
}

// Bytes returns the byte representation of the EncryptPrivateKey.
// It marshals the private key in an uncompressed format.
func (priv *EncryptPrivateKey) Bytes() []byte {
	return priv.PrivateKey.MarshalUncompressed()
}

// MasterPublic returns the master public key corresponding to priv.
func (priv *EncryptPrivateKey) MasterPublic() *EncryptMasterPublicKey {
	return priv.EncryptMasterPublicKey
}

// SetMasterPublicKey bind the encryption master public key to it.
func (priv *EncryptPrivateKey) SetMasterPublicKey(pub *EncryptMasterPublicKey) {
	if priv.EncryptMasterPublicKey == nil || priv.EncryptMasterPublicKey.MasterPublicKey == nil {
		priv.EncryptMasterPublicKey = pub
	}
}

// UnmarshalRaw unmarsal raw bytes data to the encryption master public key
func (pub *EncryptMasterPublicKey) UnmarshalRaw(bytes []byte) error {
	g, err := unmarshalG1(bytes)
	if err != nil {
		return err
	}
	pub.MasterPublicKey = g
	return nil
}

// UnmarshalRaw unmarsal raw bytes data to the encryption private key
// Note, priv's EncryptMasterPublicKey should be handled separately.
func (priv *EncryptPrivateKey) UnmarshalRaw(bytes []byte) error {
	g, err := unmarshalG2(bytes)
	if err != nil {
		return err
	}
	priv.PrivateKey = g
	return nil
}

// isLess returns whether a < b, where a and b are big-endian buffers of the
// same length and shorter than 72 bytes.
func isLess(a, b []byte) bool {
	if len(a) != len(b) {
		panic("sm9: internal error: mismatched isLess inputs")
	}

	// Copy the values into a fixed-size preallocated little-endian buffer.
	// 72 bytes is enough for every scalar in this package, and having a fixed
	// size lets us avoid heap allocations.
	if len(a) > 72 {
		panic("sm9: internal error: isLess input too large")
	}
	bufA, bufB := make([]byte, 72), make([]byte, 72)
	for i := range a {
		bufA[i], bufB[i] = a[len(a)-i-1], b[len(b)-i-1]
	}

	// Perform a subtraction with borrow.
	var borrow uint64
	for i := 0; i < len(bufA); i += 8 {
		limbA, limbB := binary.LittleEndian.Uint64(bufA[i:]), binary.LittleEndian.Uint64(bufB[i:])
		_, borrow = bits.Sub64(limbA, limbB, borrow)
	}

	// If there is a borrow at the end of the operation, then a < b.
	return borrow == 1
}

var errInvalidPrivateKey = errors.New("sm9: invalid private key")