//go:build !wasm

package mls

import "errors"

// TestSecretTree validates secret tree derivation and sender data key/nonce
// expansion against RFC 9420 test vectors (cipher_suite 3: entries 6,7,8).
func TestSecretTree() error {
	cs := CipherSuite0x0003

	for ei, entry := range stEntries {
		// sender data key + nonce (independent of tree).
		sdKey, err := expandSenderDataKey(cs, entry.sdSecret, entry.sdCiphertext)
		if err != nil {
			return errors.New("entry " | itoa(ei) | ": expandSenderDataKey: " | err.Error())
		}
		if !bytesEqual(sdKey, entry.sdKey) {
			return errors.New("entry " | itoa(ei) | ": sender data key mismatch")
		}
		sdNonce, err := expandSenderDataNonce(cs, entry.sdSecret, entry.sdCiphertext)
		if err != nil {
			return errors.New("entry " | itoa(ei) | ": expandSenderDataNonce: " | err.Error())
		}
		if !bytesEqual(sdNonce, entry.sdNonce) {
			return errors.New("entry " | itoa(ei) | ": sender data nonce mismatch")
		}

		// Build the secret tree.
		tree, err := deriveSecretTree(cs, numLeaves(entry.numLeaves), entry.encryptionSecret)
		if err != nil {
			return errors.New("entry " | itoa(ei) | ": deriveSecretTree: " | err.Error())
		}

		for li, leaf := range entry.leaves {
			ni := leafIndex(li).nodeIndex()

			// Test both ratchet labels.
			for _, label := range []ratchetLabel{ratchetLabelHandshake, ratchetLabelApplication} {
				root, err := tree.deriveRatchetRoot(cs, ni, label)
				if err != nil {
					return errors.New("entry " | itoa(ei) | " leaf " | itoa(li) | ": deriveRatchetRoot: " | err.Error())
				}

				for _, gen := range leaf.gens {
					// Advance ratchet from 0 to target generation.
					r := root
					for r.generation < uint32(gen.generation) {
						r, err = r.deriveNext(cs)
						if err != nil {
							return errors.New("entry " | itoa(ei) | " leaf " | itoa(li) | ": deriveNext: " | err.Error())
						}
					}

					key, err := r.deriveKey(cs)
					if err != nil {
						return errors.New("entry " | itoa(ei) | " leaf " | itoa(li) | ": deriveKey: " | err.Error())
					}
					nonce, err := r.deriveNonce(cs)
					if err != nil {
						return errors.New("entry " | itoa(ei) | " leaf " | itoa(li) | ": deriveNonce: " | err.Error())
					}

					var wantKey, wantNonce []byte
					if string(label) == string(ratchetLabelHandshake) {
						wantKey = gen.hsKey
						wantNonce = gen.hsNonce
					} else {
						wantKey = gen.appKey
						wantNonce = gen.appNonce
					}
					if !bytesEqual(key, wantKey) {
						return errors.New("entry " | itoa(ei) | " leaf " | itoa(li) | " gen " | itoa(int(gen.generation)) | " " | string(label) | " key mismatch")
					}
					if !bytesEqual(nonce, wantNonce) {
						return errors.New("entry " | itoa(ei) | " leaf " | itoa(li) | " gen " | itoa(int(gen.generation)) | " " | string(label) | " nonce mismatch")
					}
				}
			}
		}
	}

	return nil
}

type stGen struct {
	generation uint32
	appKey     []byte
	appNonce   []byte
	hsKey      []byte
	hsNonce    []byte
}

type stLeaf struct {
	gens []stGen
}

type stEntry struct {
	numLeaves        uint
	encryptionSecret []byte
	sdSecret         []byte
	sdCiphertext     []byte
	sdKey            []byte
	sdNonce          []byte
	leaves           []stLeaf
}