package sm2a

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

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

type ServerSignKeyGenContext struct {
	B [SecureParam]*big.Int
	X [SecureParam]*big.Int

	ServerKey    *sm2.PrivateKey
	PubKey       *sm2.PublicKey
	ServerSubKey *sm2.PrivateKey // ServerSubKey * ClientSubKey = 1+d

	isServerKeyReady bool
}

func NewServerSignKeyGenContext() *ServerSignKeyGenContext {
	s := new(ServerSignKeyGenContext)
	for i := 0; i < SecureParam; i++ {
		s.B[i] = new(big.Int)
		s.X[i] = new(big.Int)
	}
	s.ServerKey = sm2.NewPrivateKey()
	s.PubKey = sm2.NewPublicKey()
	s.ServerSubKey = sm2.NewPrivateKey()

	return s
}

func (s *ServerSignKeyGenContext) GetServerKey() (*sm2.PrivateKey, error) {
	if s.isServerKeyReady {
		return s.ServerKey, nil
	}
	return nil, fmt.Errorf("Key agreement not finished")
}

func (s *ServerSignKeyGenContext) GetServerSubKey() (*sm2.PrivateKey, error) {
	if s.isServerKeyReady {
		return s.ServerSubKey, nil
	}
	return nil, fmt.Errorf("Key agreement not finished")
}

func (s *ServerSignKeyGenContext) GetPublicKey() (*sm2.PublicKey, error) {
	if s.isServerKeyReady {
		return s.PubKey, nil
	}
	return nil, fmt.Errorf("Key agreement not finished")
}

// MarshalBinary 保存ServerSignKeyGenContext
func (s *ServerSignKeyGenContext) MarshalBinary() ([]byte, error) {
	out := make([]byte, 0, 8432)
	size := sm2.ByteSize()
	N := sm2.OrderN()

	// binary.BigEndian.PutUint32(out[:4], uint32(SecureParam))
	for i := 0; i < SecureParam; i++ {
		s.B[i].Mod(s.B[i], N)
		out = append(out, gmath.BigIntToNByte(s.B[i], size)...)
		s.X[i].Mod(s.X[i], N)
		out = append(out, gmath.BigIntToNByte(s.X[i], size)...)

	}
	// s.ServerSubKey.Mod(s.ServerSubKey, N)
	// out = append(out, gmath.BigIntToNByte(s.ServerSubKey, size)...)

	var buf4 [4]byte
	buf, err := s.ServerSubKey.MarshalBinary()
	if err != nil {
		return nil, err
	}
	binary.BigEndian.PutUint32(buf4[:], uint32(len(buf)))
	out = append(out, buf4[:]...)
	out = append(out, buf...)

	buf, err = s.ServerKey.MarshalBinary()
	if err != nil {
		return nil, err
	}
	binary.BigEndian.PutUint32(buf4[:], uint32(len(buf)))
	out = append(out, buf4[:]...)
	out = append(out, buf...)

	buf, err = s.PubKey.MarshalBinary()
	if err != nil {
		return nil, err
	}
	binary.BigEndian.PutUint32(buf4[:], uint32(len(buf)))
	out = append(out, buf4[:]...)
	out = append(out, buf...)

	if s.isServerKeyReady {
		out = append(out, 1)
	} else {
		out = append(out, 0)
	}
	return out, nil
}

// UnmarshalBinary 恢复ServerSignKeyGenContext
func (s *ServerSignKeyGenContext) UnmarshalBinary(b []byte) error {
	size := sm2.ByteSize()

	for i := 0; i < SecureParam; i++ {
		s.B[i].SetBytes(b[:size])
		b = b[size:]
		s.X[i].SetBytes(b[:size])
		b = b[size:]
	}

	// s.ServerSubKey.SetBytes(b[:size])
	// b = b[size:]

	len := binary.BigEndian.Uint32(b)
	b = b[4:]
	if err := s.ServerSubKey.UnmarshalBinary(b[:len]); err != nil {
		return err
	}
	b = b[len:]

	len = binary.BigEndian.Uint32(b)
	b = b[4:]
	if err := s.ServerKey.UnmarshalBinary(b[:len]); err != nil {
		return err
	}
	b = b[len:]

	len = binary.BigEndian.Uint32(b)
	b = b[4:]
	if err := s.PubKey.UnmarshalBinary(b[:len]); err != nil {
		return err
	}
	b = b[len:]

	if b[0] != 0 {
		s.isServerKeyReady = true
	} else {
		s.isServerKeyReady = false
	}

	return nil
}

// ServerGenKey_one 服务端第一步，生成协同密钥中间数据给客户端
// 计算(X_i,Y_i) = [x_i]·G
// 输出 X_1||Y_1||X_2||Y_2||....||X_128||Y_128
func (s *ServerSignKeyGenContext) ServerGenKey_one(rnd io.Reader) ([]byte, error) {
	var err error
	s.isServerKeyReady = false
	buf := make([]byte, sm2.ByteSize())
	out := make([]byte, 2*SecureParam*sm2.ByteSize())
	defer func() {
		for i := range buf {
			buf[i] = 0 // 清除临时使用的内存数据
		}
	}()

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

	d := s.ServerKey.Get()
	d.SetInt64(0)
	// 服务端计算中间数据
	for i := 0; i < SecureParam; i++ {
		if n, err := rnd.Read(buf); n != len(buf) || err != nil {
			return nil, fmt.Errorf("Generate random number error")
		}
		s.B[i].SetBytes(buf)
		d.Add(d, s.B[i])
		if n, err := rnd.Read(buf); n != len(buf) || err != nil {
			return nil, fmt.Errorf("Generate random number error")
		}
		s.X[i].SetBytes(buf)
		x, y := sm2.Curve256.ScalarBaseMult(buf)
		copy(out[(i*2+1)*sm2.ByteSize()-len(x.Bytes()):], x.Bytes())
		copy(out[(i*2+2)*sm2.ByteSize()-len(y.Bytes()):], y.Bytes())
	}
	d.Mod(d, sm2.OrderN())

	return out, nil
}

// ServerGenKey_two 服务端第二步，收到客户端数据in，计算中间数据并返回给客户端
// 其中in = (h_i0, h_i1, F_ix, F_iy), ... , for i = 1...128
// 输出 (C_i0, C_i1), ...., for i = 1..128
func (s *ServerSignKeyGenContext) ServerGenKey_two(in []byte, rand io.Reader) ([]byte, error) {
	tmp := new(big.Int)
	h0 := new(big.Int)
	h1 := new(big.Int)
	fx := new(big.Int)
	fy := new(big.Int)
	out := make([]byte, 2*sm2.ByteSize()*SecureParam+2*sm2.ByteSize())
	v := new(big.Int)
	buf := make([]byte, sm2.ByteSize())
	if n, err := rand.Read(buf); n != len(buf) || err != nil {
		return nil, fmt.Errorf("Generate random number error")
	}
	v.SetBytes(buf)
	v.Mod(v, sm2.Curve().Params().P)
	pos := 0
	for i := 0; i < SecureParam; i++ {
		h0.SetBytes(in[pos : pos+sm2.ByteSize()])
		pos += sm2.ByteSize()
		h1.SetBytes(in[pos : pos+sm2.ByteSize()])
		pos += sm2.ByteSize()
		fx.SetBytes(in[pos : pos+sm2.ByteSize()])
		pos += sm2.ByteSize()
		fy.SetBytes(in[pos : pos+sm2.ByteSize()])
		pos += sm2.ByteSize()

		x, y := sm2.Curve256.ScalarMult(fx, fy, s.X[i].Bytes())
		h0.Add(h0, s.B[i])
		h0.Mul(h0, v)
		pointToInt(tmp, x, y)
		h0.Add(h0, tmp)
		h0.Mod(h0, sm2.Curve().Params().N)

		x, y = sm2.Curve256.ScalarBaseMult(s.X[i].Bytes())
		y.Sub(sm2.Curve().Params().P, y)
		fx, fy = sm2.Curve().Add(fx, fy, x, y)
		x, y = sm2.Curve256.ScalarMult(fx, fy, s.X[i].Bytes())
		h1.Add(h1, s.B[i])
		h1.Mul(h1, v)
		pointToInt(tmp, x, y)
		h1.Add(h1, tmp)
		h1.Mod(h1, sm2.Curve().Params().N)

		copy(out[i*2*sm2.ByteSize():], gmath.BigIntToNByte(h0, sm2.ByteSize()))
		copy(out[(i*2+1)*sm2.ByteSize():], gmath.BigIntToNByte(h1, sm2.ByteSize()))
	}
	fx, fy = sm2.Curve256.ScalarBaseMult(v.Bytes())
	copy(out[SecureParam*2*sm2.ByteSize():], gmath.BigIntToNByte(fx, sm2.ByteSize()))
	copy(out[(SecureParam*2+1)*sm2.ByteSize():], gmath.BigIntToNByte(fy, sm2.ByteSize()))
	s.ServerSubKey.SetBigInt(v)
	return out, nil
}

// ServerGenKey_three 服务端第三步，服务端计算公钥
// in = [u^{-1}]·G
func (s *ServerSignKeyGenContext) ServerGenKey_three(in []byte) (err error) {
	x := new(big.Int)
	y := new(big.Int)
	x.SetBytes(in[:sm2.ByteSize()])
	y.SetBytes(in[sm2.ByteSize() : 2*sm2.ByteSize()])
	x, y = sm2.Curve256.ScalarMult(x, y, s.ServerSubKey.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)
	s.PubKey.X.Set(x)
	s.PubKey.Y.Set(y)
	s.isServerKeyReady = true
	return nil
}