package sm2a

import (
	"io"
	"math/big"

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

type ClientSignContext struct {
	pk   *sm2.PublicKey
	k1   *big.Int
	rand io.Reader
}

func NewClientSignContext(pk *sm2.PublicKey, r io.Reader) *ClientSignContext {
	c := &ClientSignContext{
		pk:   sm2.NewPublicKey(),
		k1:   new(big.Int),
		rand: r,
	}
	c.pk.Set(pk)
	return c
}

// Initial 客户端签名第一步，初始化，输入预处理后的待签名数据e,输出协同计算中间数据给服务端
// 输出 = e || [k_1](P+G)
func (c *ClientSignContext) Initial(e []byte) ([]byte, error) {
	buf := make([]byte, sm2.ByteSize())
	if n, err := c.rand.Read(buf); n != len(buf) || err != nil {
		return nil, err
	}
	c.k1.SetBytes(buf)
	c.k1.Mod(c.k1, sm2.OrderN())

	x, y := sm2.Curve().Add(c.pk.X, c.pk.Y, sm2.BaseX(), sm2.BaseY())
	x, y = sm2.Curve256.ScalarMult(x, y, c.k1.Bytes())

	out := make([]byte, sm3.Size+2*sm2.ByteSize())
	pos := copy(out, e)
	pos += copy(out[pos:], gmath.BigIntToNByte(x, sm2.ByteSize()))
	copy(out[pos:], gmath.BigIntToNByte(y, sm2.ByteSize()))
	return out, nil

}

// Final 客户端根据服务端协同计算结果，计算签名值
// in = r || s1
func (c *ClientSignContext) Final(clientKey *sm2.PrivateKey, in []byte) (*sm2.Signature, error) {
	sig := sm2.NewSignature()
	sig.R.SetBytes(in[:sm2.ByteSize()])
	s1 := new(big.Int)
	s1.SetBytes(in[sm2.ByteSize():])
	sig.S.Mul(sig.R, clientKey.D)
	sig.S.Add(sig.S, c.k1)
	sig.S.Add(sig.S, s1)
	sig.S.Sub(sig.S, sig.R)

	sig.S.Mod(sig.S, sm2.OrderN())
	sig.R.Mod(sig.R, sm2.OrderN())
	return sig, nil
}