package sm2m

import (
	"fmt"
	"io"
	"math/big"

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

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

// NewClientSignContext 生成客户端协同签名上下文
// pk没有用,可以传nil, rnd是随机数发生器.
// 注: pk实际是用来验证完成后的签名. 这里由调用者自己验证即可.
func NewClientSignContext(pk *sm2.PublicKey, rnd io.Reader) *ClientSignContext {
	c := &ClientSignContext{
		k1:   new(big.Int),
		rand: rnd,
	}
	return c
}

/*
a）	客户端选择随机数k1 in [1, n - 1]，计算e = SM3(Z || M), P = [k1]G,
其中Z为32字节杂凑值，G为基点，向服务端发送签名请求
同时将out = e, P加密为发送给服务端。
*/
func (c *ClientSignContext) Initial(e []byte) (out []byte, err error) {
	if c.rand == nil {
		c.rand = grand.Reader
	}
	if c.k1 == nil {
		c.k1 = new(big.Int)
	}
	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.Curve256.ScalarMult(sm2.BaseX(), sm2.BaseY(), 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
}

/*
c）	客户端计算s= d1(k1s1+s2) - r mod n,并检查s,r+s是否为0 mod q。
对消息M的签名为（r,s），客户端可验证签名是否正确。
*/
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).SetBytes(in[sm2.ByteSize() : 2*sm2.ByteSize()])
	s2 := new(big.Int).SetBytes(in[2*sm2.ByteSize():])
	s1.Mul(s1, c.k1)
	s1.Add(s1, s2)
	s1.Mod(s1, sm2.OrderN())
	s1.Mul(s1, clientKey.D)
	s1.Sub(s1, sig.R)
	s1.Mod(s1, sm2.OrderN())
	sig.S.Set(s1)
	return sig, nil
}

/*
b）
in = e || P
服务端生成随机数k2, k3 in [1,n-1]，（k = k1 * k2 + k3）
计算(x1,y1)=[k2]P+[k3]G, r=e+x1 mod n ,
检查r是否为0，为0则重新生成k2，重新计算。
计算s1=k2d2 mod n, s2=d2(k3+r) mod n，
输出 r || s1 || s2
*/
func ServerSign(serverKey *sm2.PrivateKey, in []byte, rand io.Reader) ([]byte, error) {

	k2 := make([]byte, sm2.ByteSize())
	if n, err := rand.Read(k2); n != len(k2) || err != nil {
		return nil, err
	}
	k3 := make([]byte, sm2.ByteSize())
	if n, err := rand.Read(k3); n != len(k3) || err != nil {
		return nil, err
	}

	r := new(big.Int).SetBytes(in[:sm2.ByteSize()])
	px := new(big.Int).SetBytes(in[sm2.ByteSize() : 2*sm2.ByteSize()])
	py := new(big.Int).SetBytes(in[2*sm2.ByteSize() : 3*sm2.ByteSize()])
	// x1, _ := sm2.Curve().CombinedMult(px, py, k3, k2)
	x1, _ := ec256.CombinedMult(px, py, k3, k2)

	r.Add(r, x1)
	r.Mod(r, sm2.OrderN())

	if gmath.IsBigInt0(r) {
		return nil, fmt.Errorf("rebuild")
	}

	out := make([]byte, 0, 3*sm2.ByteSize())
	out = append(out, gmath.BigIntToNByte(r, sm2.ByteSize())...)

	s := new(big.Int).SetBytes(k2)
	s.Mul(s, serverKey.D)
	s.Mod(s, sm2.OrderN())
	out = append(out, gmath.BigIntToNByte(s, sm2.ByteSize())...)

	s.SetBytes(k3)
	s.Add(s, r)
	s.Mul(s, serverKey.D)
	s.Mod(s, sm2.OrderN())
	out = append(out, gmath.BigIntToNByte(s, sm2.ByteSize())...)

	return out, nil
}