package sm9

import (
	"io"
	"math/big"

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

// KeyPackage 输出密钥封装结构
type KeyPackage struct {
	G1
}

// KeyEncapsule output KeyPackage and key with kenlen.
//
// rnd可以为:
//   - []byte, 并且len(rnd)= 32
//   - io.Reader, 随机数发生器， 如crypto/rand.Reader, 也可以是包装的硬件随机数发生器
//   - nil, 则会使用默认的软件随机数发生器
func KeyEncapsulate(id []byte, keylen int, pube *MastEncPublicKey, rnd interface{}) (*KeyPackage, []byte, error) {
	if rnd == nil {
		rnd = grand.Reader
	}

	if b, ok := rnd.([]byte); ok {
		return keyEncapsulate(id, keylen, pube, b)
	} else if reader, ok := rnd.(io.Reader); ok {
		b := make([]byte, numBytes)
		if _, err := reader.Read(b); err != nil {
			return nil, nil, gerrors.ChainErrors(errors.ErrKeyEncapsuleFailed, err)
		}
		return keyEncapsulate(id, keylen, pube, b)
	} else {
		panic("rnd can only be of nil, []byte or io.Reader")
	}
}

func keyEncapsulate(id []byte, keylen int, pube *MastEncPublicKey, rnd []byte) (*KeyPackage, []byte, error) {
	if len(rnd) != numBytes {
		panic("input rnd invalid")
	}
	C := hashToG1(id, &pube.G1, hidKeyEncapsule)
	r := new(big.Int).SetBytes(rnd)
	C = C.ScalarMult(C, r)

	key := make([]byte, keylen)
	// genKey(key, id, C, &pube.G1, g2Gen,r)
	w := &GT{}
	if pube.e != nil {
		w.ScalarMult(pube.e, r)
	} else {
		pairing(w, &pube.G1, g2Gen)
		w.ScalarMult(w, r)
	}
	genKey(key, id, C, w)
	return &KeyPackage{*C}, key, nil
}

// KeyUnencapsule unencapsule keypackage
func KeyDecapsulate(id []byte, pack *KeyPackage, keylen int, de *UserEncKey) ([]byte, error) {
	if !pack.G1.IsValid() {
		return nil, gerrors.WithAnnotating(errors.ErrKeyUnencapsuleFailed, "the KeyPackage's C1 is not a valid point on curve")
	}
	key := make([]byte, keylen)
	//genKey(key, id, pack.X, pack.Y, pack.X, pack.Y, de.X0, de.X1, de.Y0, de.Y1, nil)
	// genKey(key, id, &pack.G1, &pack.G1, &de.G2, nil)
	w := &GT{}
	pairing(w, &pack.G1, &de.G2)
	genKey(key, id, &pack.G1, w)
	return key, nil
}
