package sm2

import (
	"io"
	"math/big"

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/gmath"
	"xdx.jelly/xgcl/grand"
)

var wMask *big.Int
var w int

func initKeyExchange() {
	w = (Curve().Params().N.BitLen()+1)/2 - 1
	wMask = big.NewInt(1)
	wMask.Lsh(wMask, uint(w))
	wMask.Sub(wMask, gmath.BigInt1)
}

var XBar = computeXBar

func computeXBar(x *big.Int) *big.Int {
	xBar := new(big.Int).And(x, wMask)
	xBar.SetBit(xBar, w, 1)
	return xBar
}

// /////////////////////////////////////  包函数的方式，不绑定方法到struct上
type ExchgParameters struct {
	PrivateKey *PrivateKey // ra and Ra = [ra]·G or rb and Rb = [rb]·G
	IsSponsor  bool
}

func (p *ExchgParameters) PublicKey() *PublicKey {
	// return p.privateKey.PublicKey.X, p.privateKey.PublicKey.Y
	return &p.PrivateKey.PublicKey
}

// SharedKey generate shared key.
func (p *ExchgParameters) SharedKey(sk *PrivateKey, peerPk *PublicKey, peerPublicKey *PublicKey) (*big.Int, *big.Int, error) {
	xBar1 := XBar(p.PrivateKey.X)
	xBar2 := XBar(peerPublicKey.X)

	t := new(big.Int).Mul(xBar1, p.PrivateKey.D)
	t.Add(sk.D, t)
	t.Mod(t, Curve().Params().N)
	x, y := Curve().ScalarMult(peerPublicKey.X, peerPublicKey.Y, xBar2.Bytes())
	x, y = Curve().Add(x, y, peerPk.X, peerPk.Y)
	x, y = Curve().ScalarMult(x, y, t.Bytes())

	if gmath.IsBigInt0(x) && gmath.IsBigInt0(y) {
		return nil, nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "U/V is infinity point")
	}
	return x, y, nil
}

// GenerateAgreementData 发起方调用
func GenerateAgreementData(rand io.Reader) (*ExchgParameters, error) {
	r, err := GenerateKey(Curve(), rand)
	if err != nil {
		return nil, err
	}

	return &ExchgParameters{
		PrivateKey: r,
		IsSponsor:  true,
	}, err
}

// GenerateAgreementDataAndKey 响应方调用生成临时公钥及协商密钥
// Za、Zb由PreComputeWithIdAndPubkey计算
func GenerateAgreementDataAndKey(
	id []byte,
	privateKey *PrivateKey,
	peerId []byte,
	peerPublicKey *PublicKey,
	peerParam *PublicKey,
	keyLength int, rand io.Reader) ([]byte, *ExchgParameters, error) {

	r, err := GenerateAgreementData(rand)
	r.IsSponsor = false
	if err != nil {
		return nil, nil, err
	}

	key, err := GenerateSharedKey(r, id, privateKey, peerId, peerPublicKey, peerParam, keyLength)

	return key, r, nil
}

// GenerateSharedKey 发起方调用生成协商密钥。
func GenerateSharedKey(
	r *ExchgParameters,
	id []byte,
	privateKey *PrivateKey,
	peerId []byte,
	peerPublicKey *PublicKey,
	peerParam *PublicKey,
	keyLength int) ([]byte, error) {

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

	if !peerPublicKey.IsValid() {
		return nil, gerrors.WithAnnotating(ErrInvalidPublicKey, "peer's public key is not a valid curve point")
	}
	x, y, err := r.SharedKey(privateKey, peerPublicKey, peerParam)
	if err != nil {
		return nil, gerrors.ChainErrors(ErrKeyExchangeFailed, err)
	}
	key := make([]byte, keyLength)
	if r.IsSponsor {
		Kdf(key, gmath.BigIntToNByte(x, byteSize),
			gmath.BigIntToNByte(y, byteSize),
			PreComputeWithIdAndPubkey(id, &privateKey.PublicKey),
			PreComputeWithIdAndPubkey(peerId, peerPublicKey),
		)
	} else {
		Kdf(key, gmath.BigIntToNByte(x, byteSize),
			gmath.BigIntToNByte(y, byteSize),
			PreComputeWithIdAndPubkey(peerId, peerPublicKey),
			PreComputeWithIdAndPubkey(id, &privateKey.PublicKey),
		)
	}
	return key, nil
}

// Sponsor holds a sponsor's key pairs and temp key pairs, context data.
// Sponsor对象可以重复使用
type Sponsor struct {
	KeyBits            uint32
	z                  []byte
	r                  *PrivateKey // sponsor's temp private key
	privateKey         *PrivateKey
	publicKey          *PublicKey
	tempKeyOfSponsor   *PublicKey
	tempKeyOfResponsor *PublicKey
}

// NewSponsor make a copy of privateKey, so Clear will not change the input privateKey
// 若id=nil, 则使用默认id。
func NewSponsor(id []byte, privateKey *PrivateKey) *Sponsor {
	pk := GenPublicKey(privateKey)
	return &Sponsor{
		z:                  PreComputeWithIdAndPubkey(id, pk),
		privateKey:         NewPrivateKey().Set(privateKey),
		publicKey:          pk,
		tempKeyOfSponsor:   NewPublicKey(),
		tempKeyOfResponsor: NewPublicKey(),
	}
}

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

}

// Responsor hold a responsor's key pairs and id
// Responsor对象可以重复使用
type Responsor struct {
	z          []byte
	privateKey *PrivateKey
	publicKey  *PublicKey
}

// NewResponsor return a responsor's instance
func NewResponsor(id []byte, privateKey *PrivateKey) *Responsor {
	pk := GenPublicKey(privateKey)
	return &Responsor{
		z:          PreComputeWithIdAndPubkey(id, pk),
		privateKey: NewPrivateKey().Set(privateKey),
		publicKey:  pk,
	}
}

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

// A is sponsor and B is responsor

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

// GenerateAgreementData 发起方首先调用，za先计算好
// Za、Zb由PreComputeWithIdAndPubkey计算
// 输出tempPubkey需转化为[]byte
func (s *Sponsor) GenerateAgreementData(rnd []byte) (tempKeyOfSponsor *PublicKey, err error) {
	if len(rnd) < byteSize {
		rnd = make([]byte, byteSize)
		if _, err := grand.GenerateRandom(rnd); err != nil {
			return nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "generate random numbers failed")
		}
	}
	// return error never happen
	s.r, _ = GenPrivateKey(rnd[:byteSize])
	tempKeyOfSponsor = GenPublicKey(s.r)
	s.tempKeyOfSponsor.Set(tempKeyOfSponsor)
	return tempKeyOfSponsor, nil
}

// GenerateKey 发起方调用生成协商密钥。
func (s *Sponsor) GenerateKey(idOfResponsor []byte, pubkeyOfResponsor, tempKeyOfResponsor *PublicKey, keyLength int) (outKey []byte, err error) {
	if !tempKeyOfResponsor.IsValid() {
		return nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "responsor's temp param is not a valid curve point")
	}

	if !pubkeyOfResponsor.IsValid() {
		return nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "responsor's public key is not a valid curve point")
	}

	xBar1 := computeXBar(s.tempKeyOfSponsor.X)
	xBar2 := computeXBar(tempKeyOfResponsor.X)
	ta := new(big.Int).Mul(xBar1, s.r.D)
	ta.Mod(ta, sm2Curve.Params().N)
	// fmt.Println("x1ra= ", ta.Text(16))
	ta.Add(s.privateKey.D, ta)
	ta.Mod(ta, sm2Curve.Params().N)

	ux, uy := sm2Curve.ScalarMult(tempKeyOfResponsor.X, tempKeyOfResponsor.Y, xBar2.Bytes())
	ux, uy = sm2Curve.Add(ux, uy, pubkeyOfResponsor.X, pubkeyOfResponsor.Y)
	// fmt.Println("x2Rb= ", ux.Text(16), uy.Text(16))
	ux, uy = sm2Curve.ScalarMult(ux, uy, ta.Bytes())
	// fmt.Println("发起方U= ", ux.Text(16), uy.Text(16))
	// U = O?
	if gmath.IsBigInt0(ux) && gmath.IsBigInt0(uy) {
		return nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "U/V is infinity point")
	}

	outKey = make([]byte, keyLength)
	zb := PreComputeWithIdAndPubkey(idOfResponsor, pubkeyOfResponsor)
	Kdf(outKey, gmath.BigIntToNByte(ux, byteSize), gmath.BigIntToNByte(uy, byteSize), s.z, zb)
	return outKey, nil
}

// 响应方

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

	if len(rnd) < byteSize {
		rnd = make([]byte, byteSize)
		if _, err := grand.GenerateRandom(rnd); err != nil {
			return nil, nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "generate random numbers failed")
		}
	}

	rb, _ := GenPrivateKey(rnd[:byteSize])
	tempKeyOfResponsor = GenPublicKey(rb)

	if !tempKeyOfSponsor.IsValid() {
		return nil, nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "responsor's temp param is not a valid curve point")
	}
	if !pubkeyOfSponsor.IsValid() {
		return nil, nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "responsor's public key is not a valid curve point")
	}

	xBar1 := computeXBar(tempKeyOfSponsor.X)
	xBar2 := computeXBar(tempKeyOfResponsor.X)
	tb := new(big.Int).Mul(xBar2, rb.D)
	tb.Mod(tb, sm2Curve.Params().N)
	tb.Add(rs.privateKey.D, tb)
	tb.Mod(tb, sm2Curve.Params().N)

	vx, vy := sm2Curve.ScalarMult(tempKeyOfSponsor.X, tempKeyOfSponsor.Y, xBar1.Bytes())
	vx, vy = sm2Curve.Add(vx, vy, pubkeyOfSponsor.X, pubkeyOfSponsor.Y)
	vx, vy = sm2Curve.ScalarMult(vx, vy, tb.Bytes())
	// U = O?
	if gmath.IsBigInt0(vx) && gmath.IsBigInt0(vy) {
		return nil, nil, gerrors.WithAnnotating(ErrKeyExchangeFailed, "U/V is infinity point")
	}
	za := PreComputeWithIdAndPubkey(idOfSponsor, pubkeyOfSponsor)
	outKey = make([]byte, keyLength)
	Kdf(outKey, gmath.BigIntToNByte(vx, byteSize), gmath.BigIntToNByte(vy, byteSize), za, rs.z)

	return outKey, tempKeyOfResponsor, nil

}