package mls

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

func deriveSecretTree(cs CipherSuite, n numLeaves, encryptionSecret []byte) (secretTree, error) {
	tree := secretTree([][]byte{:int(n.width())})
	tree.set(n.root(), encryptionSecret)
	err := tree.deriveChildren(cs, n.root())
	return tree, err
}

func (tree secretTree) deriveChildren(cs CipherSuite, x nodeIndex) error {
	l, r, ok := x.children()
	if !ok {
		return nil
	}

	parentSecret := tree.get(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
	}

	tree.set(l, leftSecret)
	tree.set(r, rightSecret)

	if err := tree.deriveChildren(cs, l); err != nil {
		return err
	}
	return tree.deriveChildren(cs, r)
}

func (tree secretTree) deriveRatchetRoot(cs CipherSuite, ni nodeIndex, label ratchetLabel) (ratchetSecret, error) {
	nh := uint16(cs.ExtractSize())
	root, err := cs.expandWithLabel(tree.get(ni), []byte(label), nil, nh)
	return ratchetSecret{root, 0}, err
}

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

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

func (secret ratchetSecret) deriveNext(cs CipherSuite) (ratchetSecret, 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) ([]byte, 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)
}