/// Copyright (c) 2018 xdx. All rights reserved.
///
/// \file: keyexchange.go
///
/// \brief: SM9密钥交换协议
///
/// \author: xdx
///
/// 注意：密钥交换协议中要求双方的主公钥一致。

package sm9

import (
	"io"
	"math/big"

	"crypto/rand"

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

// Sponsor holds a sponsor's key pairs and temp key pairs, context data.
// Sponsor对象可以重复使用
type Sponsor struct {
	id []byte
	r  big.Int // update: 20200624, change []byte to *big.Int
	de UserEncKey
	//
	idOfResponsor    []byte
	tempKeyOfSponsor MastEncPublicKey
	pube             MastEncPublicKey //responsor的主公钥
	state            int              // Sponsor的状态。
}

// NewSponsor make a copy of privateKey, so Clear will not change the input privateKey
// input:
//
//	id		sponsor's id
//	pube 	responsor's master encryption publick key
//	de		sponsor's private key
func NewSponsor(id []byte, de *UserEncKey) *Sponsor {
	ida := make([]byte, len(id))
	copy(ida, id)
	return &Sponsor{
		id: ida,
		de: *de,
	}
}

// Clear 无论成功与否，最后调用Clear
func (s *Sponsor) Clear() {
	// s.de.Clear()
}

// Responsor hold a responsor's key pairs and id
// Responsor对象可以重复使用
type Responsor struct {
	id []byte
	de UserEncKey
}

// NewResponsor return a responsor's instance
func NewResponsor(id []byte, de *UserEncKey) *Responsor {
	return &Responsor{
		id: dup(id),
		de: *de,
	}
}

// Clear 无论成功与否，最后调用Clear
func (rs *Responsor) Clear() {
	// rs.de.Clear()
}

// // A is sponsor and B is responsor

// //////////////////////
// // 发起方函数

// GenerateAgreementData 发起方首先调用
// 输出tempKeyOfSponsor发送给响应方。调用MashalBinary转换为字节
func (s *Sponsor) GenerateAgreementData(
	idOfResponsor []byte,
	pube *MastEncPublicKey,
	rnd []byte) (tempKeyOfSponsor *MastEncPublicKey, err error) {

	s.idOfResponsor = make([]byte, len(idOfResponsor))
	copy(s.idOfResponsor, idOfResponsor)

	if len(rnd) < byteSize {
		rnd = make([]byte, byteSize)
		if _, err := grand.GenerateRandom(rnd); err != nil {
			return nil, err
		}
		// return nil, gerrors.ERR_SM_INVALID_INPUT
	}
	if !pube.G1.IsValid() {
		return nil, gerrors.WithAnnotating(errors.ErrKeyExchangeFailed, "input MastEncPublicKey is not a valid point on curve")
	}
	s.pube.G1.Set(&pube.G1)

	s.r.SetBytes(rnd)
	tempKeyOfSponsor = &MastEncPublicKey{}
	tempKeyOfSponsor.G1.ScalarMult(hashToG1(idOfResponsor, &pube.G1, hidKeyExchange), &s.r)
	s.tempKeyOfSponsor.G1.Set(&tempKeyOfSponsor.G1)
	s.state = 1
	return tempKeyOfSponsor, nil
}

// GenerateKey 发起方调用生成协商密钥。
func (s *Sponsor) GenerateKey(tempKeyOfResponsor *MastEncPublicKey, keyLength int) (outKey []byte, err error) {
	// TODO: 改为状态模式？
	if s.state != 1 {
		return nil, gerrors.WithAnnotating(errors.ErrKeyExchangeFailed, "call GenerateAgreementData first")
	}

	if !tempKeyOfResponsor.G1.IsValid() {
		return nil, gerrors.WithAnnotating(errors.ErrKeyExchangeFailed, "input MastEncPublicKey is not a valid point on curve")
	}
	g1 := &GT{}
	g2 := &GT{}
	pairing(g1, &s.pube.G1, g2Gen)
	g1.ScalarMult(g1, &s.r)
	pairing(g2, &tempKeyOfResponsor.G1, &s.de.G2)
	g3 := (&GT{}).ScalarMult(g2, &s.r)

	buf := make([]byte, 0, byteSize*4+byteSize*12*3)
	buf = append(buf, s.tempKeyOfSponsor.G1.Marshal()...)
	buf = append(buf, tempKeyOfResponsor.G1.Marshal()...)
	buf = append(buf, g1.Marshal()...)
	buf = append(buf, g2.Marshal()...)
	buf = append(buf, g3.Marshal()...)

	outKey = make([]byte, keyLength)
	_ = Kdf(outKey, s.id, s.idOfResponsor, buf)
	s.state = 0
	return outKey, nil
}

// 响应方

// GenerateAgreementDataAndKey 响应方调用生成临时公钥及协商密钥
// Za、Zb由PreComputeWithIdAndPubkey计算
func (rs *Responsor) GenerateAgreementDataAndKey(
	idOfSponsor []byte,
	pubeOfSponsor, tempKeyOfSponsor *MastEncPublicKey,
	keyLength int,
	rnd []byte) (outKey []byte, tempKeyOfResponsor *MastEncPublicKey, err error) {

	if len(rnd) < byteSize {
		rnd = make([]byte, byteSize)
		if _, err := grand.GenerateRandom(rnd); err != nil {
			return nil, nil, err
		}
		// return nil, nil, gerrors.ERR_SM_INVALID_INPUT
	}
	r := new(big.Int).SetBytes(rnd)
	tempKeyOfResponsor = &MastEncPublicKey{}
	tempKeyOfResponsor.G1.ScalarMult(hashToG1(idOfSponsor, &pubeOfSponsor.G1, hidKeyExchange), r)

	if !pubeOfSponsor.G1.IsValid() {
		return nil, nil, errors.ErrInvalidMastPublicKey
	}

	if !tempKeyOfSponsor.G1.IsValid() {
		return nil, nil, gerrors.WithAnnotating(errors.ErrInvalidPublicKey, "sponsor's temp param is not a valid curve point")
	}

	g1 := &GT{}
	g2 := &GT{}
	pairing(g1, &tempKeyOfSponsor.G1, &rs.de.G2)
	pairing(g2, &pubeOfSponsor.G1, g2Gen)
	g2.ScalarMult(g2, r)
	g3 := (&GT{}).ScalarMult(g1, r)
	buf := make([]byte, 0, byteSize*4+byteSize*12*3)
	buf = append(buf, tempKeyOfSponsor.G1.Marshal()...)
	buf = append(buf, tempKeyOfResponsor.G1.Marshal()...)
	buf = append(buf, g1.Marshal()...)
	buf = append(buf, g2.Marshal()...)
	buf = append(buf, g3.Marshal()...)

	outKey = make([]byte, keyLength)
	_ = Kdf(outKey, idOfSponsor, rs.id, buf)

	return outKey, tempKeyOfResponsor, nil

}

// functional

func GenerateAgreementData(idb []byte, pube *MastEncPublicKey, rnd io.Reader) (ra *big.Int, Ra *G1, err error) {
	if !pube.G1.IsValid() {
		return nil, nil, errors.ErrInvalidMastPublicKey
	}
	ra, err = rand.Int(rnd, N)
	if err != nil {
		return nil, nil, err
	}
	for ra.Sign() == 0 {
		ra, err = rand.Int(rnd, N)
		if err != nil {
			return nil, nil, err
		}
	}

	Ra = &G1{}
	Ra.ScalarMult(hashToG1(idb, &pube.G1, hidKeyExchange), ra)
	return
}

// GenerateKey 发起方调用生成协商密钥。
func GenerateKey(ida []byte, idb []byte, Ra, Rb *G1, ra *big.Int, de *UserEncKey, pube *MastEncPublicKey, keyLength int) (outKey []byte, err error) {
	if !Rb.IsValid() {
		return nil, gerrors.WithAnnotating(errors.ErrInvalidPublicKey, "responsor's temp param(Rb) is not a valid curve point")
	}
	g1 := &GT{}
	g2 := &GT{}
	pairing(g1, &pube.G1, g2Gen)
	g1.ScalarMult(g1, ra)
	pairing(g2, Rb, &de.G2)
	g3 := (&GT{}).ScalarMult(g2, ra)

	buf := make([]byte, 0, byteSize*4+byteSize*12*3)
	buf = append(buf, Ra.Marshal()...)
	buf = append(buf, Rb.Marshal()...)
	buf = append(buf, g1.Marshal()...)
	buf = append(buf, g2.Marshal()...)
	buf = append(buf, g3.Marshal()...)

	outKey = make([]byte, keyLength)
	_ = Kdf(outKey, ida, idb, buf)
	return outKey, nil
}

// 响应方

// GenerateAgreementDataAndKey 响应方调用生成临时公钥及协商密钥
// Za、Zb由PreComputeWithIdAndPubkey计算
func GenerateAgreementDataAndKey(ida []byte, idb []byte, Ra *G1, de *UserEncKey, pube *MastEncPublicKey, keyLength int, rnd io.Reader) (outKey []byte, Rb *G1, err error) {

	rb, err := rand.Int(rnd, N)
	if err != nil {
		return nil, nil, gerrors.WithAnnotating(errors.ErrKeyExchangeFailed, "generate random number failed")
	}
	for rb.Sign() == 0 {
		rb, err = rand.Int(rnd, N)
		if err != nil {
			return nil, nil, gerrors.WithAnnotating(errors.ErrKeyExchangeFailed, "generate random number failed")
		}
	}
	Rb = &G1{}
	Rb.ScalarMult(hashToG1(ida, &pube.G1, hidKeyExchange), rb)

	if !pube.G1.IsValid() {
		return nil, nil, gerrors.WithAnnotating(errors.ErrInvalidPublicKey, "sponsor's temp param(Ra) is not a valid curve point")
	}
	if !Ra.IsValid() {
		return nil, nil, gerrors.WithAnnotating(errors.ErrInvalidPublicKey, "sponsor's temp param(Ra) is not a valid curve point")
	}

	g1 := &GT{}
	g2 := &GT{}
	pairing(g1, Ra, &de.G2)
	pairing(g2, &pube.G1, g2Gen)
	g2.ScalarMult(g2, rb)
	g3 := (&GT{}).ScalarMult(g1, rb)
	buf := make([]byte, 0, byteSize*4+byteSize*12*3)
	buf = append(buf, Ra.Marshal()...)
	buf = append(buf, Rb.Marshal()...)
	buf = append(buf, g1.Marshal()...)
	buf = append(buf, g2.Marshal()...)
	buf = append(buf, g3.Marshal()...)

	outKey = make([]byte, keyLength)
	_ = Kdf(outKey, ida, idb, buf)
	return outKey, Rb, nil
}