package sm2m

import (
	"errors"
	"fmt"
	"io"
	"math/big"

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

/*
服务端和用户端的秘钥生成
秘钥关系式： (1 + d)^(-1)= d1 * d2

签名(加密）秘钥双方生成流程
a0  客户端初始化context
a1	客户端(), 向服务端发送数据。

b1  服务端结收数据计算并发回数据（）

a2	客户端接收数据并计算，输出签名
a3  客户端清理context
*/

// clientKey = d1
// clientTempPublicKey = [d1^(-1)]*G
func ClientGenSignKey_one(rnd any) (clientKey *sm2.PrivateKey, clientTempPublicKey *sm2.PublicKey, err error) {
	if rnd == nil {
		rnd = grand.Reader
	}
	switch rnd := rnd.(type) {
	case io.Reader:
		clientKey, err := sm2.GenerateKey(sm2.Curve(), rnd)
		if err != nil {
			return nil, nil, err
		}
		clientTempPublicKey := &clientKey.PublicKey
		clientKey.D.ModInverse(clientKey.D, sm2.OrderN())
		return clientKey, clientTempPublicKey, nil
	case []byte:
		if rnd == nil || len(rnd) < 32 {
			return nil, nil, gerrors.WithAnnotating(ErrInvalidInput, "input rnd must be as 32 bytes")
		}
		clientKey = sm2.NewPrivateKey()
		_ = clientKey.SetBytes(rnd)
		clientTempPublicKey = sm2.GenPublicKey(clientKey)
		clientKey.D.ModInverse(clientKey.D, sm2.OrderN())
		err = nil
		return
	default:
		return nil, nil, errors.New("parameter rnd can only be nil, io.Reader or []byte (with length >= 32)")
	}
}

/*
serverkey = d2  in [1, n - 1]，
publicKey = [d2^(-1)]clientTempPublicKey - G,用户签名公钥P
serverTempPublicKey = [d2^(-1)]G，
*/
func ServerGenSignKey(clientTempPublicKey *sm2.PublicKey, rnd []byte) (
	serverKey *sm2.PrivateKey,
	serverTempPublicKey *sm2.PublicKey,
	publicKey *sm2.PublicKey,
	err error) {

	if rnd == nil || len(rnd) < 32 {
		return nil, nil, nil, gerrors.WithAnnotating(ErrInvalidInput, "input rnd must be as 32 bytes")
	}

	serverKey = sm2.NewPrivateKey()
	_ = serverKey.SetBytes(rnd)

	serverTempPublicKey = sm2.GenPublicKey(serverKey)
	x, y := sm2.Curve256.ScalarMult(clientTempPublicKey.X, clientTempPublicKey.Y, serverKey.D.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)
	publicKey = &sm2.PublicKey{X: x, Y: y}
	serverKey.D.ModInverse(serverKey.D, sm2.OrderN())
	err = nil
	return
}

// ClientGenSignKey_two 客户端对服务端发送来的公钥publicKey和根据clientKey,serverTempPublicKey计算出来的公钥比较。
// 若两者一致，则秘钥生成成功。
func ClientGenSignKey_two(clientKey *sm2.PrivateKey, serverTempPublicKey *sm2.PublicKey, publicKey *sm2.PublicKey) error {

	d1Inv := new(big.Int).ModInverse(clientKey.D, sm2.OrderN())
	x, y := sm2.Curve256.ScalarMult(serverTempPublicKey.X, serverTempPublicKey.Y, d1Inv.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)

	if x.Cmp(publicKey.X) != 0 || y.Cmp(publicKey.Y) != 0 {
		return fmt.Errorf("generate key failed")
	}

	return nil
}

// realPrivateKey compute the real private key, for test
func realPrivateKey(clientKey, serverKey *sm2.PrivateKey) *sm2.PrivateKey {

	d := new(big.Int).Mul(serverKey.D, clientKey.D)
	d.ModInverse(d, sm2.OrderN())
	d.Sub(d, gmath.BigInt1)
	sk := &sm2.PrivateKey{
		D: d,
	}
	sm2.GenPublicKey(sk)
	return sk
}