package outsource

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

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

// 外包服务计算

type OSSignContext struct {
	k1p *big.Int
}

func NewOSSignContext() *OSSignContext {
	return &OSSignContext{}
}

// Step1 需要临时保存k1
func (o *OSSignContext) Step1(rnd io.Reader) (pp *sm2.PublicKey, err error) {
	for {
		k, err := rand.Int(rnd, sm2.OrderN())
		if err != nil {
			return nil, err
		}
		if k.Sign() == 0 {
			continue
		}

		x, y := sm2.Curve256.ScalarBaseMult(k.Bytes())

		o.k1p = k
		pp = &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}
		return pp, nil
	}
}

// Step2 input s1 in [1, N-1]
func (o *OSSignContext) Step2(s1 *big.Int, s2 *big.Int, encryptedKey *paillier.Cipher, evalKey *paillier.PublicKey) (w *paillier.Cipher, err error) {
	if o.k1p == nil || o.k1p.Sign() == 0 {
		return nil, errors.New("call Prepare first")
	}
	if s1 == nil || s2 == nil {
		return nil, errors.New("s1 or s2 is nil")
	}

	// s1, s2不能为0 mod n, 否则外包服务可能被套取到用户的密钥.
	if s1.Sign() <= 0 || s1.Cmp(sm2.OrderN()) >= 0 || s2.Sign() <= 0 || s2.Cmp(sm2.OrderN()) >= 0 {
		return nil, errors.New("s1 or s2 not in range [1, N-1]")
	}

	s := new(big.Int).Mul(o.k1p, s1)
	s.Add(s, s2).Mod(s, sm2.OrderN()) // s = k*s1 + s2 mod N

	// check s = 0?

	// 计算 w = encryptedClientKey^s mod n^2 = Enc(ds * s)
	w = (&paillier.Cipher{}).HomomorphicScalarMul(encryptedKey, s, evalKey)

	// blind the result
	if true {
		k, err := rand.Int(grand.Reader, sm2.OrderN())
		if err != nil {
			return nil, err
		}
		k.Mul(k, sm2.OrderN())
		w.HomomorphicAddInt(w, k, evalKey)
	} else {
		kn, err := rand.Int(grand.Reader, sm2.OrderN())
		if err != nil {
			return nil, err
		}
		kn.Mul(kn, sm2.OrderN())
		blind, err := paillier.Encrypt(kn, evalKey, grand.Reader)
		if err != nil {
			return nil, err
		}
		w.HomomorphicAdd(w, blind, evalKey) // c = Enc(dc*(k1*s1+s2) + k*N)
	}
	return w, nil
}

// outsourcingContext 需要序列化后保存到redis中
func (o *OSSignContext) Marshal() ([]byte, error) {
	buf := make([]byte, sm2.ByteSize())
	return o.k1p.FillBytes(buf), nil
}

func (o *OSSignContext) Unmarshal(b []byte) error {
	if len(b) != sm2.ByteSize() {
		return errors.New("input data invalid")
	}
	if o.k1p == nil {
		o.k1p = new(big.Int)
	}
	o.k1p.SetBytes(b)
	return nil
}

// clientSignContext 外包服务客户端(与JS的outsource.ClientSigner一致)
type clientSignContext struct {
	pprivateKey *paillier.PrivateKey
	r           *big.Int
	k1pp        *big.Int
	rnd         io.Reader
}

// NewClientSignContext simple factory for creating clientSignContext.
func NewClientSignContext(pprivateKey *paillier.PrivateKey, rnd io.Reader) *clientSignContext {
	ctx := &clientSignContext{
		pprivateKey: pprivateKey,
		rnd:         rnd,
	}
	if rnd == nil {
		ctx.rnd = grand.Reader
	}
	return ctx
}

// Step1 客户端签名第一步 返回p给虎符服务端
//
// pp是外包服务第一步返回的P'=k1p*G
// 返回p = k1p * k2pp * G
func (c *clientSignContext) Step1(e []byte, pp *sm2.PublicKey) (p *sm2.PublicKey, err error) {
	c.k1pp, err = rand.Int(c.rnd, sm2.OrderN())
	if err != nil {
		return nil, err
	}
	x, y := sm2.Curve256.ScalarMult(pp.X, pp.Y, c.k1pp.Bytes())

	p = &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}
	return p, nil
}

// Step2 客户端签名第二步, 收到虎符服务端返回的r, s1, s2, 计算s1p, 把s1p, s2p给外包服务
func (c *clientSignContext) Step2(r *big.Int, s1 *big.Int, s2 *big.Int) (s1p *big.Int, s2p *big.Int, err error) {
	c.r = r
	if r.Sign() == 0 {
		return nil, nil, errors.New("Outsource sign failed, r = 0")
	}
	s1p = c.k1pp
	c.k1pp = nil
	s1p.Mul(s1p, s1).Mod(s1p, sm2.OrderN())
	s2p = s2
	return s1p, s2p, nil
}

// Step3 客户端签名第三步, 收到外包服务返回的w = Enc(s+r), 输出签名
func (c *clientSignContext) Step3(w *paillier.Cipher) (*sm2.Signature, error) {
	s, err := c.pprivateKey.Decrypt(w)
	if err != nil {
		return nil, err
	}
	if s.Sign() == 0 {
		return nil, errors.New("Outsource sign failed, s = 0")
	}
	s.Sub(s, c.r).Mod(s, sm2.OrderN())
	r := c.r
	c.r = nil
	return &sm2.Signature{R: r, S: s}, nil
}