package sm2m

import (
	"crypto/rand"
	"errors"
	"io"
	"math/big"

	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/he/paillier"
	"xdx.jelly/xgcl/sm/sm2"
	"xdx.jelly/xgcl/sm/sm2/ec256"
)

// clientBlindSignContext 只能使用NewClientBlindSignContext来生成.
type clientBlindSignContext struct {
	k1                 *big.Int
	rand               io.Reader
	paillierPrivateKey *paillier.PrivateKey
}

// NewClientSignContext 生成客户端协同签名上下文
// pk没有用,可以传nil, rnd是随机数发生器. paillierPrivateKey可以临时产生密钥,也可以产生长期的保存在客户端.
// 注: pk实际是用来验证完成后的签名. 这里由调用者自己验证即可.
func NewClientBlindSignContext(pk *sm2.PublicKey, paillierPrivateKey *paillier.PrivateKey, r io.Reader) *clientBlindSignContext {
	c := &clientBlindSignContext{
		k1:                 new(big.Int),
		rand:               r,
		paillierPrivateKey: paillierPrivateKey,
	}
	return c
}

// Initial 客户端第一步, 返回e的同态密文和临时点P, 将他们同同态公钥一起发送给服务端.
//
// 客户端生成选择随机数k1 in [1, n - 1]，计算e = SM3(Z || M), P = [k1]G,
// 其中Z为32字节杂凑值，G为基点，向服务端发送签名请求并将 Enc_paillier(e), P, 发送给服务端。
// 注: 同态公钥可以使用客户端的部分私钥对应的公钥签名,服务端验证.
func (c *clientBlindSignContext) Initial(e []byte) (ecipher *paillier.Cipher, p *sm2.PublicKey, 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, 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())
	p = &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}
	paillierPublicKey := c.paillierPrivateKey.Public()
	ecipher, err = paillierPublicKey.Encrypt(new(big.Int).SetBytes(e), c.rand)
	if err != nil {
		return nil, nil, err
	}

	return ecipher, p, nil
}

// ServerSign 服务端计算. serverKey是服务端的部分密钥, pcipher,p和ppublickey是客户端发送过来的数据.
//
// 服务端生成随机数k2, k3, k4 in [1,n-1]，（k = k1 * k2 + k3）
// 计算(x1,y1)=[k2]P+[k3]G,
// rcipher = Enc(r) = Enc(e)*Enc(x1)
// s1 = k2*d2 mod n
// s2cipher = Enc(s2) = (Enc(r)*E(k_3))^ds * Enc(k4*n) = Enc(ds*(r+k3)+k4*n).
// 返回rcipher, s1, s2cipher, 并发送给客户端
func ServerBlindSign(serverKey *sm2.PrivateKey, ecipher *paillier.Cipher,
	p *sm2.PublicKey, ppublickey *paillier.PublicKey, rnd io.Reader) (rcipher *paillier.Cipher, s1 *big.Int, s2cipher *paillier.Cipher, err error) {
	var k2, k3, k4 *big.Int
	for {
		k2, err = rand.Int(rnd, sm2.OrderN())
		if err != nil {
			return nil, nil, nil, err
		}
		if k2.Sign() == 0 {
			continue
		}
		k3, err = rand.Int(rnd, sm2.OrderN())
		if err != nil {
			return nil, nil, nil, err
		}
		if k3.Sign() == 0 {
			continue
		}
		k4, err = rand.Int(rnd, sm2.OrderN())
		if err != nil {
			return nil, nil, nil, err
		}
		if k4.Sign() == 0 {
			continue
		}
		break
	}

	x1, _ := ec256.CombinedMult(p.X, p.Y, k3.Bytes(), k2.Bytes())
	x1cipher, err := ppublickey.Encrypt(x1, rnd)
	if err != nil {
		return nil, nil, nil, err
	}
	rcipher = x1cipher.HomomorphicAdd(x1cipher, ecipher, ppublickey)

	s1 = new(big.Int)
	s1.Mul(k2, serverKey.D)
	s1.Mod(s1, sm2.OrderN())

	s2cipher, err = ppublickey.Encrypt(k3, rnd)
	if err != nil {
		return nil, nil, nil, err
	}
	s2cipher.HomomorphicAdd(s2cipher, rcipher, ppublickey)
	s2cipher.HomomorphicScalarMul(s2cipher, serverKey.D, ppublickey)

	k4.Mul(k4, sm2.OrderN())
	k4Cipher, err := ppublickey.Encrypt(k4, rnd)
	if err != nil {
		return nil, nil, nil, err
	}
	s2cipher.HomomorphicAdd(s2cipher, k4Cipher, ppublickey)

	return rcipher, s1, s2cipher, nil
}

// Final 客户端计算签名
//
// 客户端解密rcipher, s2cipher得到r, s2,
// 计算s = d1*(k1*s1+s2)-r mod n,并检查s,r+s是否为0 mod q。
// 对消息M的签名为（r,s），客户端可验证签名是否正确.
func (c *clientBlindSignContext) Final(clientKey *sm2.PrivateKey,
	rcipher *paillier.Cipher, s1 *big.Int, s2cipher *paillier.Cipher) (*sm2.Signature, error) {
	r, err := c.paillierPrivateKey.Decrypt(rcipher)
	if err != nil {
		return nil, err
	}
	r.Mod(r, sm2.OrderN())
	if r.Sign() == 0 {
		return nil, errors.New("Blind sign failed, r = 0")
	}

	s2, err := c.paillierPrivateKey.Decrypt(s2cipher)
	s2.Mod(s2, sm2.OrderN())

	s := c.k1.Mul(c.k1, s1)
	s.Add(s, s2).Mod(s, sm2.OrderN()).Mul(s, clientKey.D).Sub(s, r).Mod(s, sm2.OrderN())
	if s.Sign() == 0 {
		return nil, errors.New("Blind sign failed, s = 0")
	}
	c.k1 = nil
	return &sm2.Signature{R: r, S: s}, nil
}