package sm2a

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

	"xdx.jelly/xgcl/gmath"
	"xdx.jelly/xgcl/grand/drng"
	"xdx.jelly/xgcl/sm/sm2"
)

type ClientSignKeyGenContext struct {
	X            [SecureParam]*big.Int
	Y            [SecureParam]*big.Int
	K            [SecureParam]byte
	ClientKey    *sm2.PrivateKey // 客户端私钥，与sm2私钥是同样的形式
	PubKey       *sm2.PublicKey  // 公钥
	ClientSubKey *sm2.PrivateKey
	Rand         io.Reader

	isClientKeyReady bool
}

func NewClientSignKeyGenContext(r io.Reader) *ClientSignKeyGenContext {
	c := new(ClientSignKeyGenContext)
	for i := 0; i < SecureParam; i++ {
		c.X[i] = new(big.Int)
		c.Y[i] = new(big.Int)
	}
	c.ClientKey = sm2.NewPrivateKey()
	c.ClientSubKey = sm2.NewPrivateKey()
	c.PubKey = sm2.NewPublicKey()
	c.Rand = r

	return c
}

func (c *ClientSignKeyGenContext) GetClientKey() (*sm2.PrivateKey, error) {
	if c.isClientKeyReady {
		return c.ClientKey, nil
	}
	return nil, fmt.Errorf("Key agreement not finished")
}

func (c *ClientSignKeyGenContext) GetClientSubKey() (*sm2.PrivateKey, error) {
	if c.isClientKeyReady {
		return c.ClientSubKey, nil
	}
	return nil, fmt.Errorf("Key agreement not finished")
}

func (c *ClientSignKeyGenContext) GetPublicKey() (*sm2.PublicKey, error) {
	if c.isClientKeyReady {
		return c.PubKey, nil
	}
	return nil, fmt.Errorf("Key agreement not finished")
}

// ClientKeyGen_one 客户端第一步，根据服务端第一步计算数据计算客户端的中间数据
// in = X_1||Y_1||X_2||Y_2||....||X_128||Y_128
// out = (h_i0, h_i1, F_ix, F_iy), ... , for i = 1...128
func (c *ClientSignKeyGenContext) ClientKeyGen_one(in []byte) (out []byte, err error) {
	c.isClientKeyReady = false
	buf := make([]byte, ((SecureParam + 7) >> 3))
	if n, err := c.Rand.Read(buf); n != len(buf) || err != nil {
		return nil, fmt.Errorf("Generate random number error")
	}
	for i := 0; i < SecureParam; i++ {
		c.K[i] = (buf[i>>3] >> uint(i&7)) & 1
	}
	var fx, fy *big.Int
	xx := new(big.Int)
	xy := new(big.Int)
	r := new(big.Int)
	a := new(big.Int)
	buf = make([]byte, sm2.ByteSize())
	out = make([]byte, sm2.ByteSize()*4*SecureParam)
	outPos := 0

	// 使用基于SM3的随机比特生成器生成客户端签名密钥分量
	c.ClientKey.D, err = rand.Int(drng.SM3Rng, sm2.OrderN())
	if err != nil {
		return nil, err
	}

	var n int
	var i int
	d := c.ClientKey.Get()
	d.SetInt64(0)
	// 计算给协同签名服务端的的中间数据
	for i = 0; i < SecureParam; i++ {
		if n, err = c.Rand.Read(buf); n != len(buf) || err != nil {
			fmt.Printf("%d,%d, %x\n", i, n, len(buf))
			return nil, fmt.Errorf("Generate random number error1")
		}
		r.SetBytes(buf)
		if n, err = c.Rand.Read(buf); n != len(buf) || err != nil {
			fmt.Printf("%d,%d, %x, %v\n", i, n, buf, err)
			return nil, fmt.Errorf("Generate random number error2")
		}
		a.SetBytes(buf)
		d.Add(d, a)
		if n, err = c.Rand.Read(buf); n != len(buf) || err != nil {
			fmt.Printf("%d,%d, %x, %v\n", i, n, buf, err)
			return nil, fmt.Errorf("Generate random number error3")
		}
		fx, fy = sm2.Curve256.ScalarBaseMult(buf)
		xbuf := in[2*i*sm2.ByteSize():]
		xx.SetBytes(xbuf[:sm2.ByteSize()])
		xy.SetBytes(xbuf[sm2.ByteSize() : 2*sm2.ByteSize()])
		c.X[i], c.Y[i] = sm2.Curve256.ScalarMult(xx, xy, buf)
		if c.K[i] == 1 {
			fx, fy = sm2.Curve().Add(fx, fy, xx, xy)
			outPos += marshal(out[outPos:], r, a, fx, fy, sm2.ByteSize())
		} else {
			outPos += marshal(out[outPos:], a, r, fx, fy, sm2.ByteSize())
		}
	}
	return out, nil
}

// ClientKeyGen_two 客户端第二步，根据服务端第二步计算数据计算公钥和中间数据并发送给服务端
// in = (C_i0, C_i1), ...., for i = 1..128, 服务端计算的中间数据
// out = [u^{-1}]·G, 给服务端的中间数据
func (c *ClientSignKeyGenContext) ClientKeyGen_two(in []byte) (out []byte, err error) {

	out = make([]byte, 2*sm2.ByteSize())
	u := new(big.Int)
	tmp1 := new(big.Int)
	tmp2 := new(big.Int)
	u.SetInt64(0)
	for i := 0; i < SecureParam; i++ {
		pos := (i*2 + int(c.K[i])) * sm2.ByteSize()
		tmp1.SetBytes(in[pos : pos+sm2.ByteSize()])
		pointToInt(tmp2, c.X[i], c.Y[i])
		tmp1.Sub(tmp1, tmp2)
		u.Add(u, tmp1)
		u.Mod(u, sm2.Curve().Params().N)
	}

	u.ModInverse(u, sm2.Curve().Params().N)

	tmp1, tmp2 = sm2.Curve256.ScalarBaseMult(u.Bytes())
	copy(out, gmath.BigIntToNByte(tmp1, sm2.ByteSize()))
	copy(out[sm2.ByteSize():], gmath.BigIntToNByte(tmp2, sm2.ByteSize()))

	tmp1.SetBytes(in[2*SecureParam*sm2.ByteSize() : (2*SecureParam+1)*sm2.ByteSize()])
	tmp2.SetBytes(in[(2*SecureParam+1)*sm2.ByteSize() : (2*SecureParam+2)*sm2.ByteSize()])
	x, y := sm2.Curve256.ScalarMult(tmp1, tmp2, u.Bytes())
	y.Sub(sm2.Curve().Params().P, y)
	// 客户端计算公钥
	x, y = sm2.Curve().Add(x, y, sm2.Curve().Params().Gx, sm2.Curve().Params().Gy)
	y.Sub(sm2.Curve().Params().P, y)
	c.PubKey.X.Set(x)
	c.PubKey.Y.Set(y)

	c.ClientSubKey.SetBigInt(u)
	c.isClientKeyReady = true
	return out, nil
}

func marshal(out []byte, a, b, c, d *big.Int, n int) int {
	outPos := 0
	outPos += copy(out[outPos:], gmath.BigIntToNByte(a, n))
	outPos += copy(out[outPos:], gmath.BigIntToNByte(b, n))
	outPos += copy(out[outPos:], gmath.BigIntToNByte(c, n))
	outPos += copy(out[outPos:], gmath.BigIntToNByte(d, n))
	return outPos
}