package mls

// MLS key schedule crypto operations (RFC 9420 §8).
// Methods requiring CipherSuite.

import "errors"

func (ctx *groupContext) extractJoinerSecret(prevInitSecret, commitSecret []byte) ([]byte, error) {
	cs := ctx.cipherSuite
	extracted := cs.hkdfExtract(prevInitSecret, commitSecret)

	rawCtx, err := marshalRaw(ctx)
	if err != nil {
		return nil, err
	}
	return cs.expandWithLabel(extracted, []byte("joiner"), rawCtx, uint16(cs.ExtractSize()))
}

func (ctx *groupContext) extractEpochSecret(joinerSecret, pskSecret []byte) ([]byte, error) {
	cs := ctx.cipherSuite
	if pskSecret == nil {
		pskSecret = []byte{:cs.ExtractSize()}
	}
	extracted := cs.hkdfExtract(joinerSecret, pskSecret)

	rawCtx, err := marshalRaw(ctx)
	if err != nil {
		return nil, err
	}
	return cs.expandWithLabel(extracted, []byte("epoch"), rawCtx, uint16(cs.ExtractSize()))
}

func (ctx *groupContext) signConfirmationTag(epochSecret []byte) ([]byte, error) {
	cs := ctx.cipherSuite
	confirmationKey, err := cs.deriveSecret(epochSecret, secretLabelConfirm)
	if err != nil {
		return nil, err
	}
	return cs.signMAC(confirmationKey, ctx.confirmedTranscriptHash), nil
}

func extractWelcomeSecret(cs CipherSuite, joinerSecret, pskSecret []byte) ([]byte, error) {
	if pskSecret == nil {
		pskSecret = []byte{:cs.ExtractSize()}
	}
	extracted := cs.hkdfExtract(joinerSecret, pskSecret)
	return cs.deriveSecret(extracted, []byte("welcome"))
}

func deriveExporter(cs CipherSuite, exporterSecret, label, context []byte, length uint16) ([]byte, error) {
	derived, err := cs.deriveSecret(exporterSecret, label)
	if err != nil {
		return nil, err
	}
	contextHash := cs.hash(context)
	return cs.expandWithLabel(derived, []byte("exported"), contextHash, length)
}

func (input *confirmedTranscriptHashInput) hashValue(cs CipherSuite, interimTranscriptHashBefore []byte) ([]byte, error) {
	rawInput, err := marshalRaw(input)
	if err != nil {
		return nil, err
	}
	data := append(interimTranscriptHashBefore, rawInput...)
	return cs.hash(data), nil
}

func nextInterimTranscriptHash(cs CipherSuite, confirmedTranscriptHash, confirmationTag []byte) ([]byte, error) {
	var w Writer
	w.writeOpaqueVec(confirmationTag)
	rawInput, err := w.bytes()
	if err != nil {
		return nil, err
	}
	data := append(confirmedTranscriptHash, rawInput...)
	return cs.hash(data), nil
}

// extractPSKSecret derives the PSK secret from a list of PSKs (RFC 9420 §8.4).
func extractPSKSecret(cs CipherSuite, pskIDs []preSharedKeyID, psks [][]byte) ([]byte, error) {
	if len(pskIDs) != len(psks) {
		return nil, errors.New("mls: PSK ID count != PSK count")
	}

	zero := []byte{:cs.ExtractSize()}
	pskSecret := zero
	for i := range pskIDs {
		pskExtracted := cs.hkdfExtract(psks[i], zero)

		label := pskLabel{
			id:    pskIDs[i],
			index: uint16(i),
			count: uint16(len(pskIDs)),
		}
		rawLabel, err := marshalRaw(&label)
		if err != nil {
			return nil, err
		}

		pskInput, err := cs.expandWithLabel(pskExtracted, []byte("derived psk"), rawLabel, uint16(cs.ExtractSize()))
		if err != nil {
			return nil, err
		}
		pskSecret = cs.hkdfExtract(pskSecret, pskInput)
	}
	return pskSecret, nil
}

type pskLabel struct {
	id    preSharedKeyID
	index uint16
	count uint16
}

func (label *pskLabel) marshal(w *Writer) {
	label.id.marshal(w)
	w.addUint16(label.index)
	w.addUint16(label.count)
}