package sm9

import (
	"encoding/hex"
	"io"
	"math/big"

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

// MastSignPublicKey 签名主公钥
type MastSignPublicKey struct {
	G2
	e *GT // Precomputed. If e != nil, then e = e(pubs, g2).
}

// MastSignPrivateKey 签名主私钥。主公钥使用.Public()方法获得。
type MastSignPrivateKey struct {
	public *MastSignPublicKey
	big.Int
}

// MastEncPublicKey 加密主公钥
type MastEncPublicKey struct {
	G1
	e *GT // if e != nil, then e = e(g1, pube)
}

// MastEncPrivateKey 加密主私钥,主公钥使用.Public()方法获得
type MastEncPrivateKey struct {
	public *MastEncPublicKey
	big.Int
}

// UserSignKey 用户签名私钥
type UserSignKey struct {
	G1
}

// UserEncKey 用户加密私钥
type UserEncKey struct {
	G2
}

// BytesCount return how many bytes an UserSignKey need.
func (k *UserSignKey) BytesCount() int {
	return 2 * 256 / 8
}

// BytesCount return how many bytes an UserEncKey need.
func (k *UserEncKey) BytesCount() int {
	return 4 * 256 / 8
}

// Clear set the memory to zeros.
func (k *MastSignPrivateKey) Clear() {
	gmath.ClearBigInt(&k.Int)
	k.public.G2.SetInfinity()
}

// Clear set the memory to zeros.
func (k *MastEncPrivateKey) Clear() {
	gmath.ClearBigInt(&k.Int)
	k.public.G1.SetInfinity()
}

// Clear set the user's sign key to const g1.
func (k *UserSignKey) Clear() {
	k.G1.SetInfinity()
}

// Clear set the user's encryption key to const g2.
func (k *UserEncKey) Clear() {
	k.G2.SetInfinity()
}

// Bytes 主签名私钥转化为32字节串
func (k *MastSignPrivateKey) Bytes() []byte {
	k.Int.Mod(&k.Int, Order())
	ret := make([]byte, byteSize)
	_ = gmath.FillBytes(&k.Int, ret)
	return ret
}

// fillBytes 主签名私钥转化为32字节串, b 必须为32字节
func (k *MastSignPrivateKey) fillBytes(b []byte) {
	k.Int.Mod(&k.Int, Order())
	_ = gmath.FillBytes(&k.Int, b)
}

// Bytes 主签名公钥转化为128字节串 X0||X1||Y0||Y1
func (k *MastSignPublicKey) Bytes() []byte {
	return k.G2.Marshal()
}

// fillBytes 主签名公钥转化为128字节串 X0||X1||Y0||Y1, b必须为128字节
func (k *MastSignPublicKey) fillBytes(b []byte) {
	k.G2.FillBytes(b)
}

// Bytes 主加密私钥转化为32字节串
func (k *MastEncPrivateKey) Bytes() []byte {
	k.Int.Mod(&k.Int, Order())
	return gmath.BigIntToNByte(&k.Int, byteSize)
}

// fillBytes 主加密私钥转化为32字节串
func (k *MastEncPrivateKey) fillBytes(b []byte) {
	k.Int.Mod(&k.Int, Order())
	gmath.FillBytes(&k.Int, b)
}

// Bytes 主加密公钥转化为64字节串X||Y
func (k *MastEncPublicKey) Bytes() []byte {
	return k.G1.Marshal()
}

// fillBytes 主加密公钥转化为64字节串X||Y, len(b) must be 64
func (k *MastEncPublicKey) fillBytes(b []byte) {
	k.G1.FillBytes(b)
}

// Bytes 用户签名私钥转化为64字节串X||Y
func (k *UserSignKey) Bytes() []byte {
	return k.G1.Marshal()
}

// fillBytes 用户签名私钥转化为64字节串X||Y
func (k *UserSignKey) fillBytes(b []byte) {
	k.G1.FillBytes(b)
}

// Bytes 用户加密私钥转换为128字节串X0||X1||Y0||Y1
func (k *UserEncKey) Bytes() []byte {
	return k.G2.Marshal()
}

// fillBytes 用户加密私钥转换为128字节串X0||X1||Y0||Y1
func (k *UserEncKey) fillBytes(b []byte) {
	k.G2.FillBytes(b)
}

// SetBytes 从字节串转换为主签名私钥
func (k *MastSignPrivateKey) SetBytes(buf []byte) (*MastSignPrivateKey, error) {
	k.Int.SetBytes(buf)
	k.Int.Mod(&k.Int, Order())
	k.computePublic()
	return k, nil
}

// SetBytes 从字节串转换为主签名公钥
func (k *MastSignPublicKey) SetBytes(buf []byte) (*MastSignPublicKey, error) {
	if len(buf) < 4*byteSize {
		return nil, gerrors.WithAnnotating(errors.ErrInvalidInput, "input data too short")
	}
	_, err := k.G2.Unmarshal(buf)
	if err != nil {
		return nil, err
	}
	k.ComputePairing()
	return k, nil
}

// SetBytes 从字节串转换为主加密私钥
func (k *MastEncPrivateKey) SetBytes(buf []byte) (*MastEncPrivateKey, error) {
	k.Int.SetBytes(buf)
	k.Int.Mod(&k.Int, Order())
	k.computePublic()
	return k, nil
}

// SetBytes 从字节串转换为主加密公钥
func (k *MastEncPublicKey) SetBytes(buf []byte) (*MastEncPublicKey, error) {
	if len(buf) < 2*byteSize {
		return nil, gerrors.WithAnnotating(errors.ErrInvalidInput, "input data too short")
	}
	_, err := k.G1.Unmarshal(buf)
	if err != nil {
		return nil, err
	}
	k.ComputePairing()
	return k, nil
}

// SetBytes 从字节串转换为用户签名私钥
func (k *UserSignKey) SetBytes(src []byte) (err error) {
	_, err = k.G1.Unmarshal(src)
	return
}

// SetBytes 从字节串转换为用户加密私钥
func (k *UserEncKey) SetBytes(src []byte) (err error) {
	_, err = k.G2.Unmarshal(src)
	return
}

// String implements the Stringer interface.
func (k *MastSignPrivateKey) String() string {
	return hex.EncodeToString(k.Bytes())
}

// String implements the Stringer interface.
func (k *MastSignPublicKey) String() string {
	return hex.EncodeToString(k.Bytes())
}

// String implements the Stringer interface.
func (k *MastEncPrivateKey) String() string {
	return hex.EncodeToString(k.Bytes())
}

// String implements the Stringer interface.
func (k *MastEncPublicKey) String() string {
	return hex.EncodeToString(k.Bytes())
}

// ComputePairing computes k.e = (g1, Ppubs) if k.e == nil.
// 注意，不要对k主密钥重新赋值，否则这里的预计算值不能更新。
func (k *MastSignPublicKey) ComputePairing() {
	if k.e != nil {
		return
	}
	k.e = Pairing(g1Gen, &k.G2)
}

// ComputePairing computes k.e = (Ppube, g2)
// 注意，不要对k主密钥重新赋值，否则这里的预计算值不能更新。
func (k *MastEncPublicKey) ComputePairing() {
	if k.e != nil {
		return
	}
	k.e = Pairing(&k.G1, g2Gen)
}

// randOrder returns a number in [1, N-1].
//
// It returns an nil error unless the random number generator r reads 0 bytes and got an error.
func randOrder(r io.Reader) (*big.Int, error) {
	b := make([]byte, numBytes)
	ret := new(big.Int)

	_, err := r.Read(b)
	if err != nil {
		return nil, errors.ErrGenerateRandomFailed
	}
	for {
		ret.SetBytes(b)
		if ret.Sign() > 0 && ret.Cmp(Order()) < 0 {
			return ret, nil
		}
		d := sm3.Sum(b)
		b = d[:]
	}
}

// GenerateMastSignPrivateKey 生成签名主私钥
func GenerateMastSignPrivateKey(r io.Reader) (*MastSignPrivateKey, *MastSignPublicKey, error) {
	ks := &MastSignPrivateKey{}
	d, err := randOrder(r)
	if err != nil {
		return nil, nil, err
	}
	ks.Int = *d
	ks.computePublic()
	return ks, ks.public, nil
}

// GenerateMastEncPrivateKey 生成加密主私钥
func GenerateMastEncPrivateKey(r io.Reader) (*MastEncPrivateKey, *MastEncPublicKey, error) {
	ke := &MastEncPrivateKey{}
	d, err := randOrder(r)
	if err != nil {
		return nil, nil, err
	}
	ke.Int = *d
	ke.computePublic()
	return ke, ke.public, nil
}

// GenerateUserSignKey 生成用户私钥。
//
// 注意: 有极小可能返回error。但一旦返回error，则需要更换主签名密钥以及重新生成所有用户密钥。
// 一般来说，这个概率非常……非常低。但一旦返回error，如果KGC可以改变uid的值（比如uid内的时间等参数）
// 则用新uid重新计算。否则若uid完全由用户提交上来，则必须更换主签名密钥以及重新生成所有用户密钥。
// 因为此时该用户可以根据此uid计算出主私钥。
func GenerateUserSignKey(uid []byte, ks *MastSignPrivateKey) (*UserSignKey, error) {
	return ks.GenerateUserSignKey(uid)
}

// GenerateUserEncKey 生成用户加密私钥
//
// 注: 见GenerateUserSignKey说明
func GenerateUserEncKey(uid []byte, ke *MastEncPrivateKey) (*UserEncKey, error) {
	return ke.GenerateUserEncKey(uid)
}

func (k *MastSignPrivateKey) computePublic() {
	if k.public == nil {
		k.public = &MastSignPublicKey{}
	}
	k.public.G2.ScalarBaseMult(&k.Int)
	k.public.ComputePairing()
}

func (k *MastEncPrivateKey) computePublic() {
	if k.public == nil {
		k.public = &MastEncPublicKey{}
	}
	k.public.G1.ScalarBaseMult(&k.Int)
	k.public.ComputePairing()
}

// GenerateUserSignKey 生成用户私钥。注意如果返回错误ErrKGCRebuildKey，则需要重建主密钥。
//
//	k,err:= k.GenerateUserSignKey(uid)
//	if err.Is(ErrKGCRebuildKey){...}
//
// - 也可以更改uid中某些项，比如uid如果是ASN.1编码的identity，则可以通过改变时间等参数来重新计算。
// - 如果uid是用户提交的不能更改，则必须重建主密钥，否则用户可以推算出主密钥。
// - 返回ErrKGCRebuildKey错误的几率非常低，以致可以忽略。
func (k *MastSignPrivateKey) GenerateUserSignKey(uid []byte) (*UserSignKey, error) {
	t1 := H1(uid, []byte{hidSign})
	t1.Add(t1, &k.Int)
	t1.Mod(t1, Order())
	if gmath.IsBigInt0(t1) {
		return nil, errors.ErrKGCRebuildKey
	}

	t1.ModInverse(t1, Order())
	t1.Mul(t1, &k.Int)
	t1.Mod(t1, Order())

	return &UserSignKey{
		*new(G1).ScalarBaseMult(t1),
	}, nil
}

// GenerateUserEncKey 生成用户加密私钥
//
// 注: 见GenerateUserSignKey说明
func (k *MastEncPrivateKey) GenerateUserEncKey(uid []byte) (*UserEncKey, error) {
	t1 := H1(uid, []byte{hidEncryption})
	t1.Add(t1, &k.Int)
	t1.Mod(t1, Order())
	if gmath.IsBigInt0(t1) {
		return nil, errors.ErrKGCRebuildKey
	}

	t1.ModInverse(t1, Order())
	t1.Mul(t1, &k.Int)
	t1.Mod(t1, Order())
	return &UserEncKey{
		*new(G2).ScalarBaseMult(t1),
	}, nil
}

// Public 返回主签名公钥
func (k *MastSignPrivateKey) Public() *MastSignPublicKey {
	if k.public == nil {
		k.computePublic()
	}
	k.public.ComputePairing()
	return k.public
}

// Public 返回主加密公钥
func (k *MastEncPrivateKey) Public() *MastEncPublicKey {
	if k.public == nil {
		k.computePublic()
	}
	k.public.ComputePairing()
	return k.public
}