package sm9m

import (
	"encoding"

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

// SM9协同密钥解封装

type ClientDecapsulateContext struct {
	pack *sm9.KeyPackage
}

type ClientDecapsulateParam struct {
	c sm9.G1
}

type ServerDecapsulateParam struct {
	e sm9.GT // e(C, Ks)
}

func (ctx *ClientDecapsulateContext) GenerateParam(pack *sm9.KeyPackage) (*ClientDecapsulateParam, error) {
	if !pack.G1.IsValid() {
		return nil, ErrInvalidCipherC1
	}
	ctx.pack = pack
	param := &ClientDecapsulateParam{}
	param.c.Set(&pack.G1)
	return param, nil
}

func (ctx *ClientDecapsulateContext) DecapsulateFinal(id []byte, keylen int, clientEncKey *EncKeyClient, param *ServerDecapsulateParam) ([]byte, error) {
	w := sm9.Pairing(&ctx.pack.G1, &clientEncKey.p)
	w.Mul(w, &param.e) // e = e(C1, de)

	key := make([]byte, keylen)
	_ = sm9.Kdf(key, ctx.pack.G1.Marshal(), w.Marshal(), id)
	return key, nil
}

func ServerDecapsulate(clientParam *ClientDecapsulateParam, serverEncKey *EncKeyServer) (*ServerDecapsulateParam, error) {
	if !clientParam.c.IsValid() {
		return nil, ErrInvalidCipherC1
	}

	w := sm9.Pairing(&clientParam.c, &serverEncKey.p)
	return &ServerDecapsulateParam{e: *w}, nil
}

var _ encoding.BinaryMarshaler = &ClientDecapsulateParam{}
var _ encoding.BinaryUnmarshaler = &ClientDecapsulateParam{}

var _ encoding.BinaryMarshaler = &ServerDecapsulateParam{}
var _ encoding.BinaryUnmarshaler = &ServerDecapsulateParam{}

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

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

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

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