package sm9m

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

	"xdx.jelly/xgcl/gerrors"
	"xdx.jelly/xgcl/gmath"
	"xdx.jelly/xgcl/sm/sm9"
)

// 一、签名密钥分割方式，ds= [Ks] * Kc,
// 客户端掌握Kc, 服务端掌握 Ks \in G1

// SignKeyClient 客户端部分签名密钥
type SignKeyClient struct {
	p sm9.G1
}

// SignKeyServer 服务端部分签名密钥
type SignKeyServer struct {
	k big.Int
}

// GenerateKeys 从完整的SM9用户密钥ds生成客户端和服务端部分密钥。
// ds为KGC下发的用户密钥，应在调用完成后清除。
func GenerateSignKeys(ds *sm9.UserSignKey, rnd io.Reader) (*SignKeyClient, *SignKeyServer, error) {
	i, j, err := generatePairInt(rnd, sm9.N)
	if err != nil {
		return nil, nil, gerrors.WithAnnotating(err, "GenerateSignKeys failed")
	}
	kc := &SignKeyClient{}
	ks := &SignKeyServer{}
	ks.k.Set(i)
	kc.p.ScalarMult(&ds.G1, j)
	return kc, ks, nil
}

// 二、加密密钥分割方式，de= Kc + Ks,
// 客户端掌握Kc, 服务端掌握 Ks

// EncKeyClient 客户端部分签名密钥
type EncKeyClient struct {
	p sm9.G2
}

// EncKeyServer 服务端部分签名密钥
type EncKeyServer struct {
	p sm9.G2
}

// GenerateEncKeys 从完整的SM9用户密钥de生成客户端和服务端部分密钥。
// de为KGC下发的用户密钥，应在调用完成后清除。
func GenerateEncKeys(de *sm9.UserEncKey, rnd io.Reader) (*EncKeyClient, *EncKeyServer, error) {
	var k *big.Int
	var err error
	for {
		k, err = rand.Int(rnd, sm9.N)
		if err != nil {
			return nil, nil, err
		}
		if k.Sign() > 0 {
			break
		}
	}

	kc := &EncKeyClient{}
	ks := &EncKeyServer{}
	kc.p.ScalarBaseMult(k)
	ks.p.Neg(&kc.p)
	ks.p.Add(&de.G2, &ks.p)
	return kc, ks, nil
}

// generatePairInt 生成k,s满足k * s = 1 mod n, 0 < k,s < n.
func generatePairInt(rnd io.Reader, n *big.Int) (k *big.Int, s *big.Int, err error) {
	for {
		k, err = rand.Int(rnd, n)
		if err != nil {
			return nil, nil, err
		}
		if k.Sign() > 0 {
			break
		}
	}
	s = new(big.Int).ModInverse(k, n)
	return k, s, nil
}

var _ encoding.BinaryMarshaler = &EncKeyClient{}
var _ encoding.BinaryUnmarshaler = &EncKeyClient{}
var _ encoding.BinaryMarshaler = &EncKeyServer{}
var _ encoding.BinaryUnmarshaler = &EncKeyServer{}

var _ encoding.BinaryMarshaler = &SignKeyClient{}
var _ encoding.BinaryUnmarshaler = &SignKeyClient{}
var _ encoding.BinaryMarshaler = &SignKeyServer{}
var _ encoding.BinaryUnmarshaler = &SignKeyServer{}

// UnmarshalBinary implements encoding.BinaryUnmarshaler
func (k *EncKeyServer) UnmarshalBinary(data []byte) error {
	_, err := k.p.Unmarshal(data)
	return err
}

// MarshalBinary implements encoding.BinaryMarshaler
func (k *EncKeyServer) MarshalBinary() (data []byte, err error) {
	return k.p.Marshal(), nil
}

// UnmarshalBinary implements encoding.BinaryUnmarshaler
func (k *EncKeyClient) UnmarshalBinary(data []byte) error {
	_, err := k.p.Unmarshal(data)
	return err
}

// MarshalBinary implements encoding.BinaryMarshaler
func (k *EncKeyClient) MarshalBinary() (data []byte, err error) {
	return k.p.Marshal(), nil
}

// UnmarshalBinary implements encoding.BinaryUnmarshaler
func (k *SignKeyClient) UnmarshalBinary(data []byte) error {
	_, err := k.p.Unmarshal(data)
	return err
}

// MarshalBinary implements encoding.BinaryMarshaler
func (k *SignKeyClient) MarshalBinary() (data []byte, err error) {
	return k.p.Marshal(), nil
}

// UnmarshalBinary implements encoding.BinaryUnmarshaler
func (k *SignKeyServer) UnmarshalBinary(data []byte) error {
	if len(data) > sm9.ByteSize() {
		return ErrInvalidInput
	}
	k.k.SetBytes(data)
	if k.k.Cmp(sm9.N) >= 0 {
		return ErrKeyTooBig
	}
	if k.k.Sign() == 0 {
		return ErrKeyIsZero
	}
	return nil
}

// MarshalBinary implements encoding.BinaryMarshaler
func (k *SignKeyServer) MarshalBinary() (data []byte, err error) {
	data = make([]byte, sm9.ByteSize())
	if err := gmath.FillBytes(&k.k, data); err != nil {
		return nil, err
	}
	return data, nil
}