package concentration

import (
	"encoding/binary"
	"fmt"
	"hash/crc32"
	"math/big"

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

var (
	ErrServerKeyGen     = fmt.Errorf("GenerateAgreementData error")
	ErrInputDataInvalid = fmt.Errorf("input data invalid")
	ErrServerSignError  = fmt.Errorf("server sign error")
)

type ServerKeyGenContext struct {
	d          *sm2.PrivateKey
	u          *sm2.PublicKey
	clientPPub *sm2.PublicKey

	//sk *sm2.PrivateKey
	//pk *sm2.PublicKey
}

func NewServerKeyGenContext() *ServerKeyGenContext {
	return &ServerKeyGenContext{
		d: sm2.NewPrivateKey(),
		u: sm2.NewPublicKey(),
	}
}

// GenerateAgreementData 根据客户端发送的数据生成临时数据并发送给客户度
func (s *ServerKeyGenContext) GenerateAgreementData(rnd32 []byte) ([]byte, error) {
	var data []byte
	var err error

	if len(rnd32) < sm2.ByteSize() {
		rnd32 = grand.GetRandom(sm2.ByteSize())
	}
	if s.d, err = sm2.GenPrivateKey(rnd32); err != nil {
		return []byte{}, err
	}
	s.u = sm2.GenPublicKey(s.d)
	if data, err = s.u.MarshalBinary(); err != nil {
		return []byte{}, ErrServerKeyGen
	}
	return data, nil
}

// ComputeKeyPair clientData = ppub_x || ppub_y || W_x || W_y || t || check
// out key pair
func (s *ServerKeyGenContext) ComputeKeyPair(clientData []byte, clientID []byte) (*sm2.PrivateKey, *sm2.PublicKey, error) {
	if len(clientData) < 64+64+32+4 {
		return nil, nil, ErrInputDataInvalid
	}
	checkSumShouldBe := crc32.ChecksumIEEE(clientData[:64+64+32])
	checkSum := binary.BigEndian.Uint32(clientData[128+32:])
	if checkSum != checkSumShouldBe {
		return nil, nil, ErrInputDataInvalid
	}

	sk := sm2.NewPrivateKey()
	sk.SetBytes(clientData[128 : 128+32])
	sk.D.Add(sk.D, s.d.D)
	sk.D.Mod(sk.D, sm2.OrderN())
	if gmath.IsBigInt0(sk.D) {
		return nil, nil, ErrServerKeyGen
	}
	pk := sm2.GenPublicKey(sk)
	//
	//ppub := sm2.NewPublicKey()
	//if err := ppub.UnmarshalBinary(clientData); err != nil{
	//	return nil, nil, ErrServerKeyGen
	//}
	//
	//digest := sm3.Sum(clientData[64:128],sm2.PreComputeWithIdAndPubkey(clientID, ppub))
	//PublicKey := sm2.Curve256.ScalarMult(ppub.X, ppub.Y, digest[:])
	return sk, pk, nil
}

func (s *ServerKeyGenContext) Clear() {
	s.d.Clear()
}

type ServerSignContext struct {
	k1         *sm2.PrivateKey
	u          *sm2.PublicKey
	clientPPub *sm2.PublicKey
}

func NewServerSignContext() *ServerSignContext {
	return &ServerSignContext{
		sm2.NewPrivateKey(),
		sm2.NewPublicKey(),
		sm2.NewPublicKey(),
	}
}

func (s *ServerSignContext) MarshalBinary() ([]byte, error) {
	if s.k1 == nil {
		s.k1 = sm2.NewPrivateKey()
	}
	if s.u == nil {
		s.u = sm2.NewPublicKey()
	}

	if s.clientPPub == nil {
		s.clientPPub = sm2.NewPublicKey()
	}

	r := make([]byte, 0, 3*4+(4+sm2.ECCRefMaxLen)+(4+2*sm2.ECCRefMaxLen)+(4+2*sm2.ECCRefMaxLen))
	var buf []byte
	var uintBuf = make([]byte, 4)
	var err error

	if buf, err = s.k1.MarshalBinary(); err != nil {
		return nil, err
	}
	binary.BigEndian.PutUint32(uintBuf, uint32(len(buf)))
	r = append(r, uintBuf...)
	r = append(r, buf...)

	if buf, err = s.u.MarshalBinary(); err != nil {
		return nil, err
	}
	binary.BigEndian.PutUint32(uintBuf, uint32(len(buf)))
	r = append(r, uintBuf...)
	r = append(r, buf...)

	if buf, err = s.clientPPub.MarshalBinary(); err != nil {
		return nil, err
	}
	binary.BigEndian.PutUint32(uintBuf, uint32(len(buf)))
	r = append(r, uintBuf...)
	r = append(r, buf...)
	return r, nil
}

func (s *ServerSignContext) UnmarshalBinary(data []byte) error {

	if s.k1 == nil {
		s.k1 = sm2.NewPrivateKey()
	}
	if s.u == nil {
		s.u = sm2.NewPublicKey()
	}

	if s.clientPPub == nil {
		s.clientPPub = sm2.NewPublicKey()
	}

	var err error

	if len(data) < 4 {
		return fmt.Errorf("invalid input of ServerSignContext.UnmarshalBinary")
	}
	size := binary.BigEndian.Uint32(data)
	data = data[4:]
	if uint32(len(data)) < size {
		return fmt.Errorf("invalid input of ServerSignContext.UnmarshalBinary")
	}
	if err = s.k1.UnmarshalBinary(data[:size]); err != nil {
		return err
	}
	data = data[size:]

	if len(data) < 4 {
		return fmt.Errorf("invalid input of ServerSignContext.UnmarshalBinary")
	}
	size = binary.BigEndian.Uint32(data)
	data = data[4:]
	if uint32(len(data)) < size {
		return fmt.Errorf("invalid input of ServerSignContext.UnmarshalBinary")
	}
	if err = s.u.UnmarshalBinary(data[:size]); err != nil {
		return err
	}
	data = data[size:]

	if len(data) < 4 {
		return fmt.Errorf("invalid input of ServerSignContext.UnmarshalBinary")
	}
	size = binary.BigEndian.Uint32(data)
	data = data[4:]
	if uint32(len(data)) < size {
		return fmt.Errorf("invalid input of ServerSignContext.UnmarshalBinary")
	}
	if err = s.clientPPub.UnmarshalBinary(data[:size]); err != nil {
		return err
	}

	return nil
}

func (s *ServerSignContext) GenerateSignData(rnd32 []byte) ([]byte, error) {
	var data []byte = make([]byte, 0, 64)
	var err error

	if len(rnd32) < sm2.ByteSize() {
		rnd32 = grand.GetRandom(sm2.ByteSize())
	}
	if s.k1, err = sm2.GenPrivateKey(rnd32); err != nil {

	}
	s.u = sm2.GenPublicKey(s.k1)
	data = append(data, gmath.BigIntToNByte(s.u.X, 32)...)
	data = append(data, gmath.BigIntToNByte(s.u.Y, 32)...)
	return data, nil
}

// Sign generate signature
// input: clientData 32 bytes
func (ssc *ServerSignContext) Sign(clientData []byte, e []byte, sk *sm2.PrivateKey) (*sm2.Signature, error) {
	k := new(big.Int)
	k.SetBytes(clientData)
	k.Mul(k, ssc.k1.D)
	k.Mod(k, sm2.OrderN())
	x, _ := sm2.Curve256.ScalarBaseMult(k.Bytes())
	r := new(big.Int)
	r.SetBytes(e)
	r.Add(x, r)
	r.Mod(r, sm2.OrderN())

	s := new(big.Int)
	s.Set(sk.D)
	s.Add(s, gmath.BigInt1)
	s.ModInverse(s, sm2.OrderN())
	k.Add(k, r)
	s.Mul(s, k)
	s.Sub(s, r)
	s.Mod(s, sm2.OrderN())
	sig := &sm2.Signature{
		R: r,
		S: s,
	}
	return sig, nil
}

func (ssc *ServerSignContext) Clear() {
	ssc.k1.Clear()
}