package sm2a

import (
	"fmt"

	"xdx.jelly/xgcl/gmath"
	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/sm/sm2"
)

type ClientSponsor struct {
	mixPoint           *sm2.PublicKey
	KeyBits            uint32
	z                  []byte
	r                  *sm2.PrivateKey // sponsor's temp private key
	clientKey          *sm2.PrivateKey
	publicKey          *sm2.PublicKey
	tempKeyOfSponsor   *sm2.PublicKey
	tempKeyOfResponsor *sm2.PublicKey
}

func NewClientSponsor(id []byte, clientKey *sm2.PrivateKey, publicKey *sm2.PublicKey) *ClientSponsor {
	return &ClientSponsor{
		z:                  sm2.PreComputeWithIdAndPubkey(id, publicKey),
		clientKey:          sm2.NewPrivateKey().Set(clientKey),
		publicKey:          sm2.NewPublicKey().Set(publicKey),
		tempKeyOfSponsor:   sm2.NewPublicKey(),
		tempKeyOfResponsor: sm2.NewPublicKey(),
	}
}

// GenerateAgreementData 发起方客户端计算密钥交换数据
func (c *ClientSponsor) GenerateAgreementData(rnd []byte) (tempKeyOfSponsor *sm2.PublicKey, err error) {
	if len(rnd) < sm2.ByteSize() {
		rnd = make([]byte, sm2.ByteSize())
		if _, err := grand.GenerateRandom(rnd); err != nil {
			return nil, err
		}
	}
	c.r, _ = sm2.GenPrivateKey(rnd[:sm2.ByteSize()])
	tempKeyOfSponsor = sm2.GenPublicKey(c.r)
	c.tempKeyOfSponsor.Set(tempKeyOfSponsor)
	return tempKeyOfSponsor, nil
}

// GenerateKey_1of2 发起方客户端计算密钥第一步
func (c *ClientSponsor) GenerateKey_1of2(pubkeyOfResponsor, tempKeyOfResponsor *sm2.PublicKey) ([]byte, error) {
	if !pubkeyOfResponsor.IsValid() || !tempKeyOfResponsor.IsValid() {
		return nil, fmt.Errorf("Input Public Key are not valid")
	}
	c.mixPoint = mixPoint(pubkeyOfResponsor, tempKeyOfResponsor) // P+[x_R]·R
	return marshalPoint(c.mixPoint), nil
}

// GenerateKey_2of2 发起方客户端计算密钥第二步
func (c *ClientSponsor) GenerateKey_2of2(keyLength int, idOfResponsor []byte, fromServer []byte, pubkeyOfResponsor *sm2.PublicKey, clientKey *sm2.PrivateKey) (key []byte, err error) {
	if len(fromServer) < 2*sm2.ByteSize() {
		return nil, fmt.Errorf("Input Public Key are not valid")
	}

	tempKeyOfServer := unmarshalPoint(fromServer)
	if !tempKeyOfServer.IsValid() || !pubkeyOfResponsor.IsValid() {
		return nil, fmt.Errorf("Input Public Key are not valid")
	}

	// U = [d_c + d_s + \bar{x_1} * r_A]*MixPointB
	//   = [d_s]*MixPointB + [d_c + \bar{x_1} * r_A]*MixPointB
	w := sm2.XBar(c.tempKeyOfSponsor.X)
	w.Mul(w, c.r.D)
	w.Add(w, clientKey.D)
	w.Mod(w, sm2.OrderN())
	x, y := sm2.Curve256.ScalarMult(c.mixPoint.X, c.mixPoint.Y, w.Bytes())
	x, y = sm2.Curve().Add(x, y, tempKeyOfServer.X, tempKeyOfServer.Y) // U = (x,y)

	// printLog("U= %x\n", append(append([]byte{}, x.Bytes()...), y.Bytes()...))

	if gmath.IsBigInt0(x) && gmath.IsBigInt0(y) {
		return nil, fmt.Errorf("U is zero")
	}

	key = make([]byte, keyLength)
	sm2.Kdf(key,
		gmath.BigIntToNByte(x, sm2.ByteSize()),
		gmath.BigIntToNByte(y, sm2.ByteSize()),
		c.z,
		sm2.PreComputeWithIdAndPubkey(idOfResponsor, pubkeyOfResponsor))

	return key, nil

}

type ClientResponsor struct {
	z         []byte
	clientKey *sm2.PrivateKey
	publicKey *sm2.PublicKey

	mixPoint *sm2.PublicKey
}

// NewResponsor return a responsor's instance
func NewClientResponsor(id []byte, clientKey *sm2.PrivateKey, publicKey *sm2.PublicKey) *ClientResponsor {
	return &ClientResponsor{
		z:         sm2.PreComputeWithIdAndPubkey(id, publicKey),
		clientKey: sm2.NewPrivateKey().Set(clientKey),
		publicKey: sm2.NewPublicKey().Set(publicKey),
	}
}

// GenerateAgreementDataAndKey_1of2 响应方客户端生成数据和密钥第一步
func (cr *ClientResponsor) GenerateAgreementDataAndKey_1of2(pubkeyOfSponsor, tempKeyOfSponsor *sm2.PublicKey) (toServer []byte, err error) {
	if !pubkeyOfSponsor.IsValid() || !tempKeyOfSponsor.IsValid() {
		return nil, fmt.Errorf("Input Public Key are not valid")
	}
	cr.mixPoint = mixPoint(pubkeyOfSponsor, tempKeyOfSponsor)
	return marshalPoint(cr.mixPoint), nil
}

// GenerateAgreementDataAndKey_1of2 响应方客户端生成数据和密钥第二步
func (cr *ClientResponsor) GenerateAgreementDataAndKey_2of2(keyLength int, idOfSponsor []byte, pubkeyOfSponsor *sm2.PublicKey, fromServer []byte, rnd []byte) (key []byte, tempKeyOfResponsor *sm2.PublicKey, err error) {
	if len(fromServer) < 2*sm2.ByteSize() {
		return nil, nil, fmt.Errorf("Input Public Key are not valid")
	}

	tempKeyOfServer := unmarshalPoint(fromServer)
	if !pubkeyOfSponsor.IsValid() || !tempKeyOfServer.IsValid() {
		return nil, nil, fmt.Errorf("Input Public Key are not valid")
	}

	if len(rnd) < sm2.ByteSize() {
		rnd = make([]byte, sm2.ByteSize())
		if _, err := grand.GenerateRandom(rnd); err != nil {
			return nil, nil, err
		}
	}

	rb, _ := sm2.GenPrivateKey(rnd[:sm2.ByteSize()])
	tempKeyOfResponsor = sm2.GenPublicKey(rb)

	w := sm2.XBar(tempKeyOfResponsor.X)
	w.Mul(w, rb.D)
	w.Add(w, cr.clientKey.D)
	w.Mod(w, sm2.OrderN())
	x, y := sm2.Curve256.ScalarMult(cr.mixPoint.X, cr.mixPoint.Y, w.Bytes())
	x, y = sm2.Curve().Add(x, y, tempKeyOfServer.X, tempKeyOfServer.Y) // V = (x,y)

	// printLog("V=", append(append([]byte{}, x.Bytes()...), y.Bytes()...))

	if gmath.IsBigInt0(x) && gmath.IsBigInt0(y) {
		return nil, nil, fmt.Errorf("U is zero")
	}

	key = make([]byte, keyLength)
	sm2.Kdf(key,
		gmath.BigIntToNByte(x, sm2.ByteSize()),
		gmath.BigIntToNByte(y, sm2.ByteSize()),
		sm2.PreComputeWithIdAndPubkey(idOfSponsor, pubkeyOfSponsor),
		cr.z,
	)

	return key, tempKeyOfResponsor, nil
}

// mixPoint return p+[x_r]·r
func mixPoint(p, r *sm2.PublicKey) *sm2.PublicKey {
	xBar := sm2.XBar(r.X)
	x, y := sm2.Curve256.ScalarMult(r.X, r.Y, xBar.Bytes())
	x, y = sm2.Curve().Add(x, y, p.X, p.Y)
	return &sm2.PublicKey{X: x, Y: y}
}

func marshalPoint(p *sm2.PublicKey) []byte {
	out := make([]byte, 0, 2*sm2.ByteSize())
	out = append(out, gmath.BigIntToNByte(p.X, sm2.ByteSize())...)
	out = append(out, gmath.BigIntToNByte(p.Y, sm2.ByteSize())...)
	return out
}

func unmarshalPoint(b []byte) *sm2.PublicKey {
	ret := sm2.NewPublicKey()
	ret.X.SetBytes(b[:sm2.ByteSize()])
	ret.Y.SetBytes(b[sm2.ByteSize() : 2*sm2.ByteSize()])
	return ret
}