package embed

import (
	"net"
	"testing"
)

func mustCIDR(t *testing.T, s string) *net.IPNet {
	t.Helper()
	_, n, err := net.ParseCIDR(s)
	if err != nil {
		t.Fatalf("ParseCIDR(%q): %v", s, err)
	}
	return n
}

func TestDistribute(t *testing.T) {
	cases := []struct {
		total, k int
		want     []int
	}{
		// from the doc table
		{19, 1, []int{19}}, // /48 1 field — would exceed MaxFieldNibbles, see error test below
		{19, 2, []int{10, 9}},
		{19, 3, []int{7, 6, 6}},
		{17, 1, []int{17}},
		{17, 2, []int{9, 8}},
		{17, 3, []int{6, 6, 5}},
		{15, 1, []int{15}},
		{15, 2, []int{8, 7}},
		{15, 3, []int{5, 5, 5}},
		{11, 1, []int{11}},
		{11, 2, []int{6, 5}},
		{11, 3, []int{4, 4, 3}},
		{7, 1, []int{7}},
		{7, 2, []int{4, 3}},
		{7, 3, []int{3, 2, 2}},
	}
	for _, c := range cases {
		got, err := distribute(c.total, c.k)
		if c.total > MaxFieldNibbles && c.k == 1 {
			if err == nil {
				t.Errorf("distribute(%d,%d): expected MaxFieldNibbles error", c.total, c.k)
			}
			continue
		}
		if err != nil {
			t.Errorf("distribute(%d,%d): %v", c.total, c.k, err)
			continue
		}
		if !equal(got, c.want) {
			t.Errorf("distribute(%d,%d) = %v, want %v", c.total, c.k, got, c.want)
		}
	}
}

func equal(a, b []int) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func TestEmbed_Slash64Deterministic(t *testing.T) {
	// /64 with 3 fields: 5+5+5+1 nibbles = 64-bit IID.
	net64 := mustCIDR(t, "2602:817:3000:f001::/64")
	addr, err := Embed(net64,
		[]Field{FieldNamespace, FieldPod, FieldImage},
		Values{Namespace: "mail", Pod: "stalwart-0", ImageFallback: "container-abc"},
		0xe,
	)
	if err != nil {
		t.Fatalf("Embed: %v", err)
	}
	// Property: same inputs → same output (twice).
	addr2, err := Embed(net64,
		[]Field{FieldNamespace, FieldPod, FieldImage},
		Values{Namespace: "mail", Pod: "stalwart-0", ImageFallback: "container-abc"},
		0xe,
	)
	if err != nil {
		t.Fatal(err)
	}
	if !addr.Equal(addr2) {
		t.Fatalf("non-deterministic: %s vs %s", addr, addr2)
	}
	// Property: prefix preserved.
	if !net64.Contains(addr) {
		t.Fatalf("addr %s outside network %s", addr, net64)
	}
	// Property: last nibble is exactly N.
	if got := addr[len(addr)-1] & 0x0F; got != 0xe {
		t.Fatalf("last nibble = %x, want e", got)
	}
}

func TestEmbed_DifferentInputsDifferentOutputs(t *testing.T) {
	net64 := mustCIDR(t, "2602:817:3000:f001::/64")
	a, _ := Embed(net64, []Field{FieldNamespace, FieldPod}, Values{Namespace: "ns1", Pod: "p1"}, 0)
	b, _ := Embed(net64, []Field{FieldNamespace, FieldPod}, Values{Namespace: "ns2", Pod: "p1"}, 0)
	if a.Equal(b) {
		t.Fatalf("different namespace produced identical IID: %s", a)
	}
}

func TestEmbed_NRandomizesLowNibble(t *testing.T) {
	net64 := mustCIDR(t, "2602:817:3000:f001::/64")
	a, _ := Embed(net64, []Field{FieldNamespace}, Values{Namespace: "x"}, 0x1)
	b, _ := Embed(net64, []Field{FieldNamespace}, Values{Namespace: "x"}, 0x2)
	if a.Equal(b) {
		t.Fatalf("changing N did not change address")
	}
	// And the only difference should be the last nibble.
	if a[15]>>4 != b[15]>>4 {
		t.Fatalf("upper nibble of last byte changed unexpectedly: %x vs %x", a[15], b[15])
	}
}

func TestEmbed_RejectsBadInputs(t *testing.T) {
	net64 := mustCIDR(t, "2602:817:3000:f001::/64")
	if _, err := Embed(net64, nil, Values{}, 0); err == nil {
		t.Fatalf("expected error for empty fields")
	}
	odd := &net.IPNet{IP: net.ParseIP("2602:817:3000::"), Mask: net.CIDRMask(63, 128)}
	if _, err := Embed(odd, []Field{FieldNamespace}, Values{Namespace: "x"}, 0); err == nil {
		t.Fatalf("expected error for /63 (not nibble-aligned)")
	}
	v4 := &net.IPNet{IP: net.ParseIP("10.0.0.0").To4(), Mask: net.CIDRMask(8, 32)}
	if _, err := Embed(v4, []Field{FieldNamespace}, Values{Namespace: "x"}, 0); err == nil {
		t.Fatalf("expected error for IPv4 network")
	}
}

func TestEmbed_ImageDigestVsFallback(t *testing.T) {
	net64 := mustCIDR(t, "2602:817:3000:f001::/64")
	digest := "sha256:abcdef0123456789aabbccddeeff00112233445566778899aabbccddeeff0011"
	a, err := Embed(net64, []Field{FieldImage}, Values{Image: digest}, 0)
	if err != nil {
		t.Fatalf("Embed digest: %v", err)
	}
	b, err := Embed(net64, []Field{FieldImage}, Values{ImageFallback: "ctr-xyz"}, 0)
	if err != nil {
		t.Fatalf("Embed fallback: %v", err)
	}
	if a.Equal(b) {
		t.Fatalf("digest and fallback produced same IID")
	}
}