package outsource

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

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

// 客户端加密clientKey^{-1}给外包服务器保存，用作同态解密密钥
func EncryptClientDecKey(clientKey *sm2.PrivateKey, k *paillier.PublicKey, rnd io.Reader) (*paillier.Cipher, error) {
	d := clientKey.D
	di := new(big.Int).ModInverse(d, sm2.OrderN())
	return k.Encrypt(di, rnd)
}

type ClientKeyGenerator struct {
	dc1, dc1Inv *big.Int
}

// NewClientKeyGen 客户端第一步
// 返回数据发送给OS服务端
func (c *ClientKeyGenerator) Step1(k *paillier.PublicKey, rnd io.Reader) (*paillier.Cipher, *paillier.Cipher, error) {
	var d *big.Int
	var err error
	for {
		d, err = rand.Int(rnd, sm2.OrderN())
		if err != nil {
			return nil, nil, err
		}
		if d.Sign() > 0 {
			break
		}
	}
	di := new(big.Int).ModInverse(d, sm2.OrderN())
	c.dc1 = d
	c.dc1Inv = di

	a1, err := k.Encrypt(d, rnd)
	if err != nil {
		return nil, nil, err
	}
	a2, err := k.Encrypt(di, rnd)
	if err != nil {
		return nil, nil, err
	}
	return a1, a2, nil

}

// GenerateDataToServer 客户端第二步
// 返回的sm2.PublicKey发送给协同签名服务端。
func (c *ClientKeyGenerator) Step2(P *sm2.PublicKey) (*sm2.PublicKey, error) {
	x, y := sm2.Curve256.ScalarMult(P.X, P.Y, c.dc1Inv.Bytes())
	return &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}, nil
}

// GenerateDataToServer 客户端第三步
// 收到协同签名服务端返回数据
func (c *ClientKeyGenerator) Step3(R *sm2.PublicKey) (*sm2.PublicKey, error) {
	x, y := sm2.Curve256.ScalarMult(R.X, R.Y, c.dc1Inv.Bytes())
	return &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}, nil
}

// GenerateDataToServer 客户端第四步
// 收到协同签名服务端返回数据
func (c *ClientKeyGenerator) Step4(T1 *sm2.PublicKey) (*sm2.PublicKey, error) {
	x, y := sm2.Curve256.ScalarMult(T1.X, T1.Y, c.dc1Inv.Bytes())
	x, y = sm2.Curve256.Add(x, y, nGx, nGy)
	return &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}, nil
}

// OS服务端
type OSKeyGenerator struct {
	dc2, dc2Inv *big.Int
}

func (o *OSKeyGenerator) Marshal() ([]byte, error) {
	data := make([]byte, 64)
	if o.dc2 != nil {
		gmath.FillBytes(o.dc2, data)
	}

	if o.dc2Inv != nil {
		gmath.FillBytes(o.dc2, data[32:])
	}

	return data, nil
}

func (o *OSKeyGenerator) Unmarshal(b []byte) error {
	if len(b) < 64 {
		return errors.New("input too short")
	}
	o.dc2 = new(big.Int).SetBytes(b[:32])
	o.dc2Inv = new(big.Int).SetBytes(b[32:64])
	return nil
}

// OSGenerateKey 外包服务器计算， 返回B1, B2, P
// 其中，保存B1, B2, B1用于签名，B2用于解密加密密钥保护结构。
// P发送给客户端。
func (o *OSKeyGenerator) Step1(a1 *paillier.Cipher, a2 *paillier.Cipher, pk *paillier.PublicKey, rnd io.Reader) (*paillier.Cipher, *paillier.Cipher, *sm2.PublicKey, error) {
	var d *big.Int
	var err error
	for {
		d, err = rand.Int(rnd, sm2.OrderN())
		if err != nil {
			return nil, nil, nil, err
		}
		if d.Sign() > 0 {
			break
		}
	}
	di := new(big.Int).ModInverse(d, sm2.OrderN())
	o.dc2 = d
	o.dc2Inv = di

	x, y := sm2.Curve256.ScalarBaseMult(di.Bytes())
	b1 := new(paillier.Cipher).HomomorphicScalarMul(a1, d, pk)
	b2 := new(paillier.Cipher).HomomorphicScalarMul(a2, di, pk)
	return b1, b2, &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}, nil
}

func (o *OSKeyGenerator) Step2(R *sm2.PublicKey, S *sm2.PublicKey) (*sm2.PublicKey, *sm2.PublicKey, error) {
	x1, y1 := sm2.Curve256.ScalarMult(R.X, R.Y, o.dc2Inv.Bytes())
	x2, y2 := sm2.Curve256.ScalarMult(S.X, S.Y, o.dc2Inv.Bytes())
	x, y := sm2.Curve256.Add(x2, y2, nGx, nGy)
	return &sm2.PublicKey{Curve: sm2.Curve(), X: x1, Y: y1}, &sm2.PublicKey{Curve: sm2.Curve(), X: x, Y: y}, nil
}