package mls

// MLS key package crypto operations (RFC 9420 §10).

import "errors"

var (
	errKeyPackageVersionMismatch = errors.New("mls: key package version mismatch")
	errKeyPackageSuiteMismatch   = errors.New("mls: key package cipher suite mismatch")
	errKeyPackageBadSource       = errors.New("mls: key package leaf node invalid source")
	errKeyPackageBadSignature    = errors.New("mls: invalid key package signature")
	errKeyPackageKeysMatch       = errors.New("mls: key package init key equals encryption key")
)

func (pkg *KeyPackage) sign(signerPriv signaturePrivateKey) error {
	var w Writer
	pkg.marshalTBS(&w)
	rawTBS, err := w.bytes()
	if err != nil {
		return err
	}
	sig, err := pkg.cipherSuite.signWithLabel(signerPriv, []byte("KeyPackageTBS"), rawTBS)
	if err != nil {
		return err
	}
	pkg.signature = sig
	return nil
}

func (pkg *KeyPackage) verifySignature() bool {
	var w Writer
	pkg.marshalTBS(&w)
	rawTBS, err := w.bytes()
	if err != nil {
		return false
	}
	return pkg.cipherSuite.verifyWithLabel(pkg.leafNode.signatureKey, []byte("KeyPackageTBS"), rawTBS, pkg.signature)
}

func (pkg *KeyPackage) verify(ctx *groupContext) error {
	if pkg.version != ctx.version {
		return errKeyPackageVersionMismatch
	}
	if pkg.cipherSuite != ctx.cipherSuite {
		return errKeyPackageSuiteMismatch
	}
	if pkg.leafNode.leafNodeSource != leafNodeSourceKeyPackage {
		return errKeyPackageBadSource
	}
	if !pkg.verifySignature() {
		return errKeyPackageBadSignature
	}
	if bytesEqual(pkg.leafNode.encryptionKey, pkg.initKey) {
		return errKeyPackageKeysMatch
	}
	return nil
}

func (pkg *KeyPackage) GenerateRef() (KeyPackageRef, error) {
	var w Writer
	pkg.marshal(&w)
	raw, err := w.bytes()
	if err != nil {
		return nil, err
	}
	h, err := pkg.cipherSuite.refHash([]byte("MLS 1.0 KeyPackage Reference"), raw)
	if err != nil {
		return nil, err
	}
	return KeyPackageRef(h), nil
}

// GenerateKeyPairPackage creates a new key pair package for suite 0x0003.
// nowUnix is the current time as Unix seconds (caller provides from DOM/SW context).
func GenerateKeyPairPackage(cs CipherSuite, credential *Credential, nowUnix int64) (*KeyPairPackage, error) {
	return GenerateKeyPairPackageWithOptions(cs, credential, nil, nowUnix)
}

// GenerateKeyPairPackageWithOptions creates a key pair package with options.
func GenerateKeyPairPackageWithOptions(cs CipherSuite, credential *Credential, opts *KeyPackageOptions, nowUnix int64) (*KeyPairPackage, error) {
	initPub, initPriv, err := cs.generateEncryptionKeyPair()
	if err != nil {
		return nil, err
	}
	encPub, encPriv, err := cs.generateEncryptionKeyPair()
	if err != nil {
		return nil, err
	}
	sigPub, sigPriv, err := cs.generateSignatureKeyPair()
	if err != nil {
		return nil, err
	}

	var capExts []extensionType
	var leafExts []extension
	var kpExts []extension
	if opts != nil {
		capExts = opts.CapabilityExtensions
		leafExts = opts.LeafExtensions
		kpExts = opts.KeyPackageExtensions
	}
	if capExts == nil {
		capExts = []extensionType{extensionTypeRatchetTree}
	}

	// Lifetime: now - 1h to now + 84 days
	notBefore := nowUnix - 3600
	notAfter := nowUnix + 84*24*3600

	keyPkg := KeyPackage{
		version:     protocolVersionMLS10,
		cipherSuite: cs,
		initKey:     initPub,
		leafNode: leafNode{
			encryptionKey:  encPub,
			signatureKey:   sigPub,
			leafNodeSource: leafNodeSourceKeyPackage,
			credential:     *credential,
			capabilities: capabilities{
				versions:     []protocolVersion{protocolVersionMLS10},
				cipherSuites: []CipherSuite{cs},
				extensions:   capExts,
				proposals:    []proposalType{proposalTypeAdd, proposalTypeUpdate, proposalTypeRemove},
				credentials:  []credentialType{credentialTypeBasic},
			},
			lifetime:   &lifetime{notBefore: uint64(notBefore), notAfter: uint64(notAfter)},
			extensions: leafExts,
		},
		extensions: kpExts,
	}

	if err := keyPkg.leafNode.sign(cs, nil, 0, sigPriv); err != nil {
		return nil, err
	}
	if err := keyPkg.sign(sigPriv); err != nil {
		return nil, err
	}

	return &KeyPairPackage{
		Public: keyPkg,
		Private: PrivateKeyPackage{
			InitKey:       initPriv,
			EncryptionKey: encPriv,
			SignatureKey:  sigPriv,
		},
	}, nil
}