package mls

// MLS secret tree crypto operations (RFC 9420 §9).

func deriveSecretTree(cs CipherSuite, n numLeaves, encryptionSecret []byte) (v [][]byte, err2 error) {
	tree := [][]byte{:int32(n.width())}
	secretTreeSet(tree, n.root(), encryptionSecret)
	err := secretTreeDeriveChildren(cs, tree, n.root())
	return tree, err
}

func secretTreeDeriveChildren(cs CipherSuite, tree [][]byte, x nodeIndex) (err error) {
	l, r, ok := x.children()
	if !ok {
		return nil
	}

	parentSecret := secretTreeGet(tree, x)
	nh := uint16(cs.ExtractSize())
	leftSecret, err := cs.expandWithLabel(parentSecret, []byte("tree"), []byte("left"), nh)
	if err != nil {
		return err
	}
	rightSecret, err := cs.expandWithLabel(parentSecret, []byte("tree"), []byte("right"), nh)
	if err != nil {
		return err
	}

	secretTreeSet(tree, l, leftSecret)
	secretTreeSet(tree, r, rightSecret)

	if e := secretTreeDeriveChildren(cs, tree, l); e != nil {
		return e
	}
	return secretTreeDeriveChildren(cs, tree, r)
}

func secretTreeDeriveRatchetRoot(tree [][]byte, cs CipherSuite, ni nodeIndex, label []byte) (v ratchetSecret, err2 error) {
	nh := uint16(cs.ExtractSize())
	root, err := cs.expandWithLabel(secretTreeGet(tree, ni), label, nil, nh)
	return ratchetSecret{root, 0}, err
}

func (secret *ratchetSecret) deriveNonce(cs CipherSuite) (out []byte, err error) {
	return deriveTreeSecret(cs, secret.secret, []byte("nonce"), secret.generation, uint16(cs.AEADNonceSize()))
}

func (secret *ratchetSecret) deriveKey(cs CipherSuite) (out []byte, err error) {
	return deriveTreeSecret(cs, secret.secret, []byte("key"), secret.generation, uint16(cs.AEADKeySize()))
}

func (secret *ratchetSecret) deriveNext(cs CipherSuite) (v ratchetSecret, err2 error) {
	nh := uint16(cs.ExtractSize())
	next, err := deriveTreeSecret(cs, secret.secret, []byte("secret"), secret.generation, nh)
	return ratchetSecret{next, secret.generation + 1}, err
}

func deriveTreeSecret(cs CipherSuite, secret, label []byte, generation uint32, length uint16) (out []byte, err error) {
	// context = I2OSP(generation, 4)
	context := []byte{byte(generation >> 24), byte(generation >> 16), byte(generation >> 8), byte(generation)}
	return cs.expandWithLabel(secret, label, context, length)
}