package concentration

import (
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"math/big"

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

var (
	ErrClientGenerateAgreementData = fmt.Errorf("client GenerateAgreementData error")
)

type ClientKeyGenContext struct {
	ks        *sm2.PrivateKey
	ppub      *sm2.PublicKey
	clientPub *sm2.PublicKey
}

func NewClientKeyGenContext() *ClientKeyGenContext {
	return &ClientKeyGenContext{}
}

// GenerateTempKGC output data, and send to server
// in: serverData, server临时公钥U
// out: []byte, ppub_x || ppub_y || W_x || W_y || t || check
func (c *ClientKeyGenContext) GenerateAgreementData(serverData []byte, clientID []byte) ([]byte, *sm2.PublicKey, error) {
	var data []byte
	var err error

	serverTempPubkey := sm2.NewPublicKey()
	if err = serverTempPubkey.UnmarshalBinary(serverData); err != nil {
		return []byte{}, nil, ErrClientGenerateAgreementData
	}

	if c.ks, err = sm2.GenPrivateKey(grand.GetRandom(sm2.ByteSize())); err != nil {
		return []byte{}, nil, ErrClientGenerateAgreementData
	}
	c.ppub = sm2.GenPublicKey(c.ks)
	data = make([]byte, 0, 128+32+4)
	data = append(data, gmath.BigIntToNByte(c.ppub.X, 32)...)
	data = append(data, gmath.BigIntToNByte(c.ppub.Y, 32)...)

	w, err := sm2.GenPrivateKey(grand.GetRandom(sm2.ByteSize()))
	if err != nil {
		return []byte{}, nil, ErrClientGenerateAgreementData
	}
	W := sm2.GenPublicKey(w)
	data = append(data, gmath.BigIntToNByte(W.X, 32)...)
	data = append(data, gmath.BigIntToNByte(W.Y, 32)...)

	// FIXME: W.X.Bytes() may not 32-bytes
	digest := sm3.Sum(W.Bytes(), sm2.PreComputeWithIdAndPubkey(clientID, c.ppub))
	// digest := sm3.Sum(W.X.Bytes(), W.Y.Bytes(), sm2.PreComputeWithIdAndPubkey(clientID, c.ppub))
	t := new(big.Int)
	t.SetBytes(digest[:])
	t.Mul(t, c.ks.D)
	t.Add(t, w.D)
	t.Mod(t, sm2.OrderN())
	data = append(data, gmath.BigIntToNByte(t, sm2.ByteSize())...)

	checkSum := crc32.ChecksumIEEE(data)
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, checkSum)
	data = append(data, buf...)

	x, y := sm2.Curve256.ScalarBaseMult(t.Bytes())
	x, y = sm2.Curve().Add(x, y, serverTempPubkey.X, serverTempPubkey.Y)
	pk := &sm2.PublicKey{
		X: x,
		Y: y,
	}
	c.clientPub = sm2.NewPublicKey()
	c.clientPub.X.Set(x)
	c.clientPub.Y.Set(y)
	return data, pk, nil
}

type ClientSignContext struct {
	k *sm2.PrivateKey
	r *big.Int
}

func NewClientSignContext() *ClientSignContext {
	return &ClientSignContext{}
}

// GenerateSignData 生成签名中间数据和r
func (c *ClientSignContext) GenerateSignData(serverData []byte, e []byte) (dataToServer []byte, r []byte, err error) {
	c.k = sm2.NewPrivateKey()
	c.k.Random(grand.Reader)
	c.r = new(big.Int)
	c.r.SetBytes(e)
	tempPK := sm2.NewPublicKey()
	tempPK.X.SetBytes(serverData[:32])
	tempPK.Y.SetBytes(serverData[32:])
	x, _ := sm2.Curve256.ScalarMult(tempPK.X, tempPK.Y, c.k.Bytes())
	c.r.Add(c.r, x)
	c.r.Mod(c.r, sm2.OrderN())

	return gmath.BigIntToNByte(c.k.D, 32), gmath.BigIntToNByte(c.r, 32), nil
}