//go:build !wasm

package mls

import (
	"errors"
	"smesh.lol/web/common/crypto/chacha20poly1305"
)

// Exported test support functions for vector validation.
// Called from web/mlstest/ main package.

// TestTreeMath validates tree math against RFC 9420 Appendix C test vectors.
// Returns nil on success, error describing first failure.
func TestTreeMath() error {
	for _, tc := range treeMathVectors {
		n := numLeaves(tc.nLeaves)
		if w := n.width(); w != tc.nNodes {
			return errors.New("width mismatch")
		}
		if r := n.root(); uint(r) != tc.root {
			return errors.New("root mismatch")
		}
		for i, want := range tc.left {
			x := nodeIndex(i)
			got, ok := x.left()
			if want < 0 {
				if ok {
					return errors.New("left should be nil")
				}
			} else {
				if !ok || int(got) != want {
					return errors.New("left mismatch")
				}
			}
		}
		for i, want := range tc.right {
			x := nodeIndex(i)
			got, ok := x.right()
			if want < 0 {
				if ok {
					return errors.New("right should be nil")
				}
			} else {
				if !ok || int(got) != want {
					return errors.New("right mismatch")
				}
			}
		}
		for i, want := range tc.parent {
			x := nodeIndex(i)
			got, ok := n.parent(x)
			if want < 0 {
				if ok {
					return errors.New("parent should be nil")
				}
			} else {
				if !ok || int(got) != want {
					return errors.New("parent mismatch")
				}
			}
		}
		for i, want := range tc.sibling {
			x := nodeIndex(i)
			got, ok := n.sibling(x)
			if want < 0 {
				if ok {
					return errors.New("sibling should be nil")
				}
			} else {
				if !ok || int(got) != want {
					return errors.New("sibling mismatch")
				}
			}
		}
	}
	return nil
}

// TestCryptoBasics validates suite 0x0003 crypto primitives against
// RFC 9420 crypto-basics test vectors (cipher_suite: 3).
// Returns nil on success, error describing first failure.
func TestCryptoBasics() error {
	cs := CipherSuite0x0003

	// AEAD round-trip sanity (not a vector, but isolates AEAD from HPKE).
	{
		var k [32]byte
		var n [12]byte
		for i := 0; i < 32; i++ {
			k[i] = byte(i)
		}
		for i := 0; i < 12; i++ {
			n[i] = byte(i + 1)
		}
		pt := []byte("Hello, MLS!")
		ct := chacha20poly1305.Seal(k, n, pt, nil)
		got, ok := chacha20poly1305.Open(k, n, ct, nil)
		if !ok {
			return errors.New("aead roundtrip: Open returned !ok")
		}
		if !bytesEqual(got, pt) {
			return errors.New("aead roundtrip: plaintext mismatch")
		}
	}

	// RFC 8439 §2.8.2 full ChaCha20-Poly1305 test vector.
	{
		key := hexb("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f")
		nonce := hexb("070000004041424344454647")
		aad := hexb("50515253c0c1c2c3c4c5c6c7")
		pt := hexb("4c616469657320616e642047656e746c656d656e206f662074686520636c617373206f66202739393a204966204920636f756c64206f6666657220796f75206f6e6c79206f6e652074697020666f7220746865206675747572652c2073756e73637265656e20776f756c642062652069742e")
		var k [32]byte
		var n [12]byte
		copy(k[:], key)
		copy(n[:], nonce)
		expCT := hexb("d31a8d34648e60db7b86afbc53ef7ec2a4aded51296e08fea9e2b5a736ee62d63dbea45e8ca9671282fafb69da92728b1a71de0a9e060b2905d6a5b67ecd3b3692ddbd7f2d778b8c9803aee328091b58fab324e4fad675945585808b4831d7bc3ff4def08e4b7a9de576d26586cec64b6116")
		expTag := hexb("1ae10b594f09e26a7e902ecbd0600691")
		expSealed := append(expCT, expTag...)
		got := chacha20poly1305.Seal(k, n, pt, aad)
		if !bytesEqual(got, expSealed) {
			return errors.New("rfc8439 AEAD Seal mismatch")
		}
		dec, ok := chacha20poly1305.Open(k, n, got, aad)
		if !ok {
			return errors.New("rfc8439 AEAD Open failed")
		}
		if !bytesEqual(dec, pt) {
			return errors.New("rfc8439 AEAD Open mismatch")
		}
	}

	// refHash
	{
		out, err := cs.refHash(
			[]byte("RefHash"),
			hexb("4f0c86f9c82fba0a896bd7eecf79a29856e98a7e4f13b9f841ae285d70ed8b68"),
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("f11019703c8b630060839b12a475fd39c6a30f8a866790ff46a35f9c65e1df3c")) {
			return errors.New("refHash mismatch")
		}
	}

	// expandWithLabel
	{
		out, err := cs.expandWithLabel(
			hexb("55aa3ae5242564782567ce097beafe19510230660008b2cc064a78387fa16f36"),
			[]byte("ExpandWithLabel"),
			hexb("2e07148f4340c62a55e7608c20d73fddf1f3b8dafb2c7ef24eceb70e136c0d8c"),
			32,
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("1df5ba7996a34f75d717916a094a14083c03a75e80f0330a8095f5f11cfe1e1f")) {
			return errors.New("expandWithLabel mismatch")
		}
	}

	// deriveSecret
	{
		out, err := cs.deriveSecret(
			hexb("cae460c779ebaa3e81c061a371486dff1ed1ff273bea369cc0fc46550b83c407"),
			[]byte("DeriveSecret"),
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("aad859818ca5f2a9896d4d3ee2dccc0cefcd69b666bdb16b52f1de15fb1a5567")) {
			return errors.New("deriveSecret mismatch")
		}
	}

	// deriveTreeSecret
	{
		out, err := deriveTreeSecret(cs,
			hexb("c994e257b53f726087ddd7121876f558f1fbd6f807e5ff010830d618d7bab6f2"),
			[]byte("DeriveTreeSecret"),
			2694881440,
			32,
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("2095d6a81ab87095d1df26f6bdf012ec06f197e418381c1795a7b758603c936d")) {
			return errors.New("deriveTreeSecret mismatch")
		}
	}

	// signWithLabel: verify reference signature
	{
		ok := cs.verifyWithLabel(
			hexb("18275f892ee0ca6f4687ff26c990776387502646ff658c3f572b324faecb05c5"),
			[]byte("SignWithLabel"),
			hexb("df308cf2dbf471edf2c29d30e3daf161b5b87d350ee3b2c715c298ec3d10d432"),
			hexb("4f56851c2c47f5115a61ff0ab6121b4a4732d4e94805fc7135a5132f87d5ca5f1dc7408816c1ea4f25887725cf5914b48c427a52cabcfeb746a2b8a12e821f08"),
		)
		if !ok {
			return errors.New("signWithLabel: reference signature verify failed")
		}
	}

	// signWithLabel: sign + verify round-trip
	{
		sig, err := cs.signWithLabel(
			hexb("4e312160ee4981358db479aa877412847abc7f7054b5605511256c395404d054"),
			[]byte("SignWithLabel"),
			hexb("df308cf2dbf471edf2c29d30e3daf161b5b87d350ee3b2c715c298ec3d10d432"),
		)
		if err != nil {
			return err
		}
		ok := cs.verifyWithLabel(
			hexb("18275f892ee0ca6f4687ff26c990776387502646ff658c3f572b324faecb05c5"),
			[]byte("SignWithLabel"),
			hexb("df308cf2dbf471edf2c29d30e3daf161b5b87d350ee3b2c715c298ec3d10d432"),
			sig,
		)
		if !ok {
			return errors.New("signWithLabel: round-trip verify failed")
		}
	}

	// decryptWithLabel: decrypt reference ciphertext
	{
		pt, err := cs.decryptWithLabel(
			hexb("9d122ad4638fcb301b6eb5f4073414afb44bb34d37b4ddee9975b2941d700edb"),
			[]byte("EncryptWithLabel"),
			hexb("0d6a5cf9ee88b1f8c79d8512477d9bfc5496c207c8173f8dcac0368b4dba7407"),
			hexb("f26e9e5a94396a90f85a5f72eedf3dacfb1b7f4164e0573edeb9c6c912e1cb49"),
			hexb("40dd09ad4c5dc29d373f814bf054c9359cb75a468bc4d2c8bbcffb072a73105c4d9416ebd4fafeb62e59a9dea55da3cd"),
		)
		if err != nil {
			return errors.New("decryptWithLabel error: " | err.Error())
		}
		expected := hexb("1dd4c1904996ce7d42cee7de68881459fa7a345da59a02040ade37103505baf6")
		if !bytesEqual(pt, expected) {
			return errors.New("decryptWithLabel mismatch")
		}
	}

	return nil
}

// TestCryptoBasics0x0001 validates suite 0x0001
// (DHKEM(X25519) + AES-128-GCM + Ed25519) against the RFC 9420 crypto-basics
// test vectors for cipher_suite=1. 0x0001 shares KEM and signature with 0x0003;
// only the AEAD and hpkeSuiteID differ.
func TestCryptoBasics0x0001() error {
	cs := CipherSuite0x0001

	// refHash (uses SHA-256 — same for both suites, but verify dispatch works).
	{
		out, err := cs.refHash(
			[]byte("RefHash"),
			hexb("40312db83f651883c05ab26fa12c6af61930015c81947cfd0f129e6d99210bb2"),
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("e8027fffc5f9bb469f29172538dc0f3a78f14f323495bbd2217eba7a77fb242a")) {
			return errors.New("0x0001 refHash mismatch")
		}
	}

	// expandWithLabel (length=16 exercises AEAD-key-sized expand).
	{
		out, err := cs.expandWithLabel(
			hexb("1499360a561335f4ef51d0a1b0d586900dc8007ae405b1ab79bf4207bb3d67e4"),
			[]byte("ExpandWithLabel"),
			hexb("2ff8c1f9d9c1248f82e372ddb5791c771695e01882abca6a64097bd2f04c971f"),
			16,
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("c1e8eb360391526c0c64039f13e0c5b1")) {
			return errors.New("0x0001 expandWithLabel mismatch")
		}
	}

	// deriveSecret.
	{
		out, err := cs.deriveSecret(
			hexb("1a9ce178a53f8752d2513c27efe9c85133f6c0a97f7b35ac200695024a77228e"),
			[]byte("DeriveSecret"),
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("3b08c195a246c4ad469c1d11c10e62890d8fa6b684494ff925409efdb1ff0464")) {
			return errors.New("0x0001 deriveSecret mismatch")
		}
	}

	// deriveTreeSecret.
	{
		out, err := deriveTreeSecret(cs,
			hexb("5133c6f8bad297f5d3beacdf477f0c45ec51b02de659d305220c5f9385c6eb43"),
			[]byte("DeriveTreeSecret"),
			2694881440,
			32,
		)
		if err != nil {
			return err
		}
		if !bytesEqual(out, hexb("8461f3ccc603eae52149a23a4134d29c880a1ad1ba70441e5d586e3521ec7b25")) {
			return errors.New("0x0001 deriveTreeSecret mismatch")
		}
	}

	// signWithLabel: verify reference signature (Ed25519 — same primitive as
	// 0x0003, routed via 0x0001 dispatch).
	{
		ok := cs.verifyWithLabel(
			hexb("85600e54e5c2919ccbd0742126e5d837cf7a2ba50d75a69b3f35dcfe4a50ffe2"),
			[]byte("SignWithLabel"),
			hexb("cd289cc7ba2869f64f3c32ffd133f500d17abace919a5ffe7faa974200d81932"),
			hexb("996bd223ddb4d55a2b57d85cb2944f21facc95696053ddf66d590060fdc719f4a26c6212ce605414e0d5e66a55921dd99d11122218c35bc23408b0076e8bc40b"),
		)
		if !ok {
			return errors.New("0x0001 signWithLabel: reference signature verify failed")
		}
	}

	// signWithLabel: sign + verify round-trip.
	{
		sig, err := cs.signWithLabel(
			hexb("a2f640dd5005fcad6adb8e9bd8b60d70946bb802e1e788307929fdac81e1ec74"),
			[]byte("SignWithLabel"),
			hexb("cd289cc7ba2869f64f3c32ffd133f500d17abace919a5ffe7faa974200d81932"),
		)
		if err != nil {
			return err
		}
		ok := cs.verifyWithLabel(
			hexb("85600e54e5c2919ccbd0742126e5d837cf7a2ba50d75a69b3f35dcfe4a50ffe2"),
			[]byte("SignWithLabel"),
			hexb("cd289cc7ba2869f64f3c32ffd133f500d17abace919a5ffe7faa974200d81932"),
			sig,
		)
		if !ok {
			return errors.New("0x0001 signWithLabel: round-trip verify failed")
		}
	}

	// decryptWithLabel: decrypt reference HPKE ciphertext (exercises AES-128-GCM
	// within HPKE-base with 0x0001 suite_id).
	{
		pt, err := cs.decryptWithLabel(
			hexb("fb1ade7939987ff12a9d620772b1f9f7caeba26f8a3ecea9617d9402cd862444"),
			[]byte("EncryptWithLabel"),
			hexb("26347dd7f218d1de8673d6a66646ce06ac5fd3aa8d5c33f65d86aeefdcf4a31e"),
			hexb("0a144e8fbf2d6dcf6fe9d2e2b8aeca5461ff5b0ea9c0ede1040c3dc7ed1dfd1c"),
			hexb("15c80ea2bc37db221baa530ef5aea88650f0ce0f262803d6f78f3a1392f7ccd960eff94ca081ee54efa4c3acfa0eb591"),
		)
		if err != nil {
			return errors.New("0x0001 decryptWithLabel error: " | err.Error())
		}
		expected := hexb("8f55dd30f03d64335c22b53ea7670bb1becf49b04021f706368fe93eeb358f46")
		if !bytesEqual(pt, expected) {
			return errors.New("0x0001 decryptWithLabel plaintext mismatch")
		}
	}

	return nil
}

func hexb(s string) []byte {
	n := len(s) / 2
	out := []byte{:n}
	for i := 0; i < n; i++ {
		out[i] = hexdig(s[2*i])<<4 | hexdig(s[2*i+1])
	}
	return out
}

var hexchars = []byte("0123456789abcdef")

func hexenc(b []byte) string {
	out := []byte{:len(b)*2}
	for i := 0; i < len(b); i++ {
		out[2*i] = hexchars[b[i]>>4]
		out[2*i+1] = hexchars[b[i]&0x0f]
	}
	return string(out)
}

func hexdig(c byte) byte {
	if c >= '0' && c <= '9' {
		return c - '0'
	}
	if c >= 'a' && c <= 'f' {
		return c - 'a' + 10
	}
	return 0
}

// treeMathVector holds one test case. -1 means null.
type treeMathVector struct {
	nLeaves uint
	nNodes  uint
	root    uint
	left    []int
	right   []int
	parent  []int
	sibling []int
}

// RFC 9420 tree-math vectors for n_leaves = 1, 2, 4, 8.
var treeMathVectors = []treeMathVector{
	{
		nLeaves: 1, nNodes: 1, root: 0,
		left:    []int{-1},
		right:   []int{-1},
		parent:  []int{-1},
		sibling: []int{-1},
	},
	{
		nLeaves: 2, nNodes: 3, root: 1,
		left:    []int{-1, 0, -1},
		right:   []int{-1, 2, -1},
		parent:  []int{1, -1, 1},
		sibling: []int{2, -1, 0},
	},
	{
		nLeaves: 4, nNodes: 7, root: 3,
		left:    []int{-1, 0, -1, 1, -1, 4, -1},
		right:   []int{-1, 2, -1, 5, -1, 6, -1},
		parent:  []int{1, 3, 1, -1, 5, 3, 5},
		sibling: []int{2, 5, 0, -1, 6, 1, 4},
	},
	{
		nLeaves: 8, nNodes: 15, root: 7,
		left:    []int{-1, 0, -1, 1, -1, 4, -1, 3, -1, 8, -1, 9, -1, 12, -1},
		right:   []int{-1, 2, -1, 5, -1, 6, -1, 11, -1, 10, -1, 13, -1, 14, -1},
		parent:  []int{1, 3, 1, 7, 5, 3, 5, -1, 9, 11, 9, 7, 13, 11, 13},
		sibling: []int{2, 5, 0, 11, 6, 1, 4, -1, 10, 13, 8, 3, 14, 9, 12},
	},
}