package node_test

import (
	"fmt"
	"sync"
	"sync/atomic"
	"testing"

	toki "git.toki-labs.com/toki/proto-socket/go"

	edgenode "iop/apps/edge/internal/node"
)

func TestRegistry_RegisterAndCount(t *testing.T) {
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{
		NodeID: "node-001",
		Alias:  "local-node",
		Client: &toki.TcpClient{},
	}
	reg.Register(entry)
	if reg.Count() != 1 {
		t.Fatalf("expected 1 node, got %d", reg.Count())
	}
}

func TestRegistryRegisterIfAbsentRejectsConcurrentDuplicate(t *testing.T) {
	reg := edgenode.NewRegistry()

	var accepted int32
	var wg sync.WaitGroup
	start := make(chan struct{})
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			if reg.RegisterIfAbsent(&edgenode.NodeEntry{
				NodeID: "node-dup",
				Alias:  "dup",
				Client: &toki.TcpClient{},
			}) {
				atomic.AddInt32(&accepted, 1)
			}
		}()
	}
	close(start)
	wg.Wait()

	if accepted != 1 {
		t.Fatalf("accepted registrations: got %d want 1", accepted)
	}
	if reg.Count() != 1 {
		t.Fatalf("registry count: got %d want 1", reg.Count())
	}
}

func TestRegistryUnregisterIfClientIgnoresStaleConnection(t *testing.T) {
	reg := edgenode.NewRegistry()
	live := &toki.TcpClient{}
	stale := &toki.TcpClient{}
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1", Client: live})

	if _, ok := reg.UnregisterIfClient("node-1", stale); ok {
		t.Fatal("stale client should not unregister live entry")
	}
	entry, ok := reg.Get("node-1")
	if !ok {
		t.Fatal("live entry was removed by stale client")
	}
	if entry.Client != live {
		t.Fatal("live entry client changed unexpectedly")
	}

	gen, ok := reg.UnregisterIfClient("node-1", live)
	if !ok {
		t.Fatal("live client should unregister entry")
	}
	if gen != entry.ConnectionGeneration {
		t.Fatalf("unregister returned generation %d, want the live owner's %d", gen, entry.ConnectionGeneration)
	}
	if reg.Count() != 0 {
		t.Fatalf("registry count after live unregister: got %d want 0", reg.Count())
	}
}

// TestRegistryAssignsMonotonicConnectionGeneration pins that each accepted
// registration of a node id draws a strictly higher generation, that a reconnect
// after unregister keeps climbing, and that a rejected duplicate never advances
// the counter or steals the live owner's generation.
func TestRegistryAssignsMonotonicConnectionGeneration(t *testing.T) {
	reg := edgenode.NewRegistry()

	first := &edgenode.NodeEntry{NodeID: "node-gen", Alias: "gen", Client: &toki.TcpClient{}}
	if !reg.RegisterIfAbsent(first) {
		t.Fatal("first registration should be accepted")
	}
	if first.ConnectionGeneration == 0 {
		t.Fatal("accepted registration must be assigned a non-zero generation")
	}
	if gen, ok := reg.CurrentGeneration("node-gen"); !ok || gen != first.ConnectionGeneration {
		t.Fatalf("CurrentGeneration = %d,%v; want %d", gen, ok, first.ConnectionGeneration)
	}
	if !reg.IsCurrentOwnerGeneration("node-gen", first.ConnectionGeneration) {
		t.Fatal("first owner generation should be current")
	}

	// A rejected duplicate must not advance the counter or become current.
	duplicate := &edgenode.NodeEntry{NodeID: "node-gen", Alias: "gen", Client: &toki.TcpClient{}}
	if reg.RegisterIfAbsent(duplicate) {
		t.Fatal("duplicate registration should be rejected")
	}
	if gen, _ := reg.CurrentGeneration("node-gen"); gen != first.ConnectionGeneration {
		t.Fatalf("rejected duplicate changed the live generation to %d", gen)
	}

	// Reconnect after unregister climbs strictly.
	if _, ok := reg.UnregisterIfClient("node-gen", first.Client); !ok {
		t.Fatal("live client should unregister")
	}
	if reg.IsCurrentOwnerGeneration("node-gen", first.ConnectionGeneration) {
		t.Fatal("no owner should be current after unregister")
	}
	second := &edgenode.NodeEntry{NodeID: "node-gen", Alias: "gen", Client: &toki.TcpClient{}}
	if !reg.RegisterIfAbsent(second) {
		t.Fatal("reconnect registration should be accepted")
	}
	if second.ConnectionGeneration <= first.ConnectionGeneration {
		t.Fatalf("reconnect generation %d must exceed previous %d", second.ConnectionGeneration, first.ConnectionGeneration)
	}
	// The old generation is no longer the current owner; the stale connection's
	// late callback cannot masquerade as the reconnect.
	if reg.IsCurrentOwnerGeneration("node-gen", first.ConnectionGeneration) {
		t.Fatal("old generation must not be current after reconnect")
	}
	if !reg.IsCurrentOwnerGeneration("node-gen", second.ConnectionGeneration) {
		t.Fatal("reconnect generation should be current")
	}
}

func TestRegistry_Resolve_ByAliasOrID(t *testing.T) {
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{
		NodeID: "node-001",
		Alias:  "local-node",
	}
	reg.Register(entry)

	// Resolve by ID
	if e, err := reg.Resolve("node-001"); err != nil || e.NodeID != "node-001" {
		t.Errorf("failed to resolve by ID: %v", err)
	}

	// Resolve by Alias
	if e, err := reg.Resolve("local-node"); err != nil || e.NodeID != "node-001" {
		t.Errorf("failed to resolve by Alias: %v", err)
	}

	// Resolve non-existent
	if _, err := reg.Resolve("unknown"); err == nil {
		t.Error("expected error for unknown node")
	}
}

func TestRegistry_Resolve_ByDisplayLabel(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-001", Alias: "local-node"})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-002", Alias: "remote-node"})

	if e, err := reg.Resolve("node0"); err != nil || e.NodeID != "node-001" {
		t.Fatalf("Resolve(node0) = %+v, %v; want node-001", e, err)
	}
	if e, err := reg.Resolve("node1"); err != nil || e.NodeID != "node-002" {
		t.Fatalf("Resolve(node1) = %+v, %v; want node-002", e, err)
	}
}

func TestRegistry_Resolve_ImplicitSingleNodeOnly(t *testing.T) {
	reg := edgenode.NewRegistry()

	// Empty
	if _, err := reg.Resolve(""); err == nil {
		t.Error("expected error on empty registry")
	}

	// Single node
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
	if e, err := reg.Resolve(""); err != nil || e.NodeID != "node-1" {
		t.Errorf("failed implicit resolve for single node: %v", err)
	}

	// Multiple nodes
	reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})
	if _, err := reg.Resolve(""); err == nil {
		t.Error("expected error for implicit resolve with multiple nodes")
	}
}

func TestRegistryRegisterDefaultsLifecycle(t *testing.T) {
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{NodeID: "node-001", Alias: "local-node"}
	reg.Register(entry)

	if entry.LifecycleState != edgenode.LifecycleConnected {
		t.Fatalf("lifecycle default: got %q want %q", entry.LifecycleState, edgenode.LifecycleConnected)
	}
}

func TestRegistry_UnregisterRemovesAlias(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"})

	reg.Unregister("node-1")
	if reg.Count() != 0 {
		t.Fatalf("expected 0 nodes, got %d", reg.Count())
	}

	if _, err := reg.Resolve("alias-1"); err == nil {
		t.Error("alias should be removed after unregister")
	}
}

func TestRegistryUpdateLifecycle(t *testing.T) {
	reg := edgenode.NewRegistry()

	if ok := reg.UpdateLifecycle("non-existent", edgenode.LifecycleOnline); ok {
		t.Error("expected false for updating non-existent node")
	}

	entry := &edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"}
	reg.Register(entry)

	if ok := reg.UpdateLifecycle("node-1", edgenode.LifecycleOnline); !ok {
		t.Error("expected true for updating existing node")
	}

	retrieved, ok := reg.Get("node-1")
	if !ok {
		t.Fatal("failed to get node")
	}
	if retrieved.LifecycleState != edgenode.LifecycleOnline {
		t.Errorf("lifecycle: got %q, want %q", retrieved.LifecycleState, edgenode.LifecycleOnline)
	}
}

// TestRegistryRegisterIfAbsentPendingUntilReady pins the two-phase transport
// handshake: an accepted registration claims the id (so duplicates are rejected)
// but stays out of every dispatch-ready lookup until MarkDispatchReadyIfClient
// flips it, at which point it becomes visible to AllReady/ResolveReady/GetReady.
func TestRegistryRegisterIfAbsentPendingUntilReady(t *testing.T) {
	reg := edgenode.NewRegistry()
	client := &toki.TcpClient{}
	entry := &edgenode.NodeEntry{NodeID: "node-ready", Alias: "ready", Client: client}
	if !reg.RegisterIfAbsent(entry) {
		t.Fatal("expected RegisterIfAbsent to accept the first registration")
	}
	if entry.DispatchReady {
		t.Fatal("RegisterIfAbsent entry must be pending, not dispatch-ready")
	}

	// Pending: present for ownership (Get), absent from every ready lookup.
	if _, ok := reg.Get("node-ready"); !ok {
		t.Fatal("pending entry must be present for duplicate-ownership checks")
	}
	if _, ok := reg.GetReady("node-ready"); ok {
		t.Fatal("pending entry must not be dispatch-ready via GetReady")
	}
	if len(reg.AllReady()) != 0 {
		t.Fatalf("pending entry must not appear in AllReady, got %d", len(reg.AllReady()))
	}
	if _, err := reg.ResolveReady("node-ready"); err == nil {
		t.Fatal("pending entry must not resolve via ResolveReady")
	}
	if _, err := reg.ResolveReady(""); err == nil {
		t.Fatal("single pending node must not satisfy implicit ready resolve")
	}

	gen, transitioned, ok := reg.MarkDispatchReadyIfClient("node-ready", client)
	if !ok || !transitioned {
		t.Fatalf("first ready must transition the current owner: ok=%v transitioned=%v", ok, transitioned)
	}
	if gen != entry.ConnectionGeneration {
		t.Fatalf("ready generation = %d, want owner's %d", gen, entry.ConnectionGeneration)
	}

	// Ready: now visible everywhere.
	if _, ok := reg.GetReady("node-ready"); !ok {
		t.Fatal("ready entry must be visible via GetReady")
	}
	if len(reg.AllReady()) != 1 {
		t.Fatalf("ready entry must appear in AllReady, got %d", len(reg.AllReady()))
	}
	if got, err := reg.ResolveReady("node-ready"); err != nil || got.NodeID != "node-ready" {
		t.Fatalf("ready entry must resolve via ResolveReady: got=%v err=%v", got, err)
	}
	if got, err := reg.ResolveReady(""); err != nil || got.NodeID != "node-ready" {
		t.Fatalf("single ready node must satisfy implicit ready resolve: got=%v err=%v", got, err)
	}
}

func TestRegistryReadyOwnerSnapshot(t *testing.T) {
	reg := edgenode.NewRegistry()
	firstClient := &toki.TcpClient{}
	first := &edgenode.NodeEntry{
		NodeID:                       "workspace-node",
		Alias:                        "workspace",
		Client:                       firstClient,
		CredentialRecipientPublicKey: []byte("first-key"),
	}
	if !reg.RegisterIfAbsent(first) {
		t.Fatal("initial registration was rejected")
	}
	if snapshot, ok := reg.ReadyOwnerSnapshot(first.NodeID); ok || snapshot != nil {
		t.Fatalf("pending snapshot = %#v, %v; want nil, false", snapshot, ok)
	}
	if _, transitioned, ok := reg.MarkDispatchReadyIfClient(first.NodeID, firstClient); !ok || !transitioned {
		t.Fatalf("MarkDispatchReadyIfClient = transitioned=%v ok=%v, want true,true", transitioned, ok)
	}

	snapshot, ok := reg.ReadyOwnerSnapshot(first.NodeID)
	if !ok || snapshot == nil {
		t.Fatal("ready owner snapshot was unavailable")
	}
	if snapshot == first {
		t.Fatal("ready owner snapshot retained the registry entry")
	}
	if snapshot.ConnectionGeneration != first.ConnectionGeneration || !snapshot.DispatchReady {
		t.Fatalf("snapshot = %#v, want ready generation %d", snapshot, first.ConnectionGeneration)
	}
	snapshot.CredentialRecipientPublicKey[0] = 'X'
	snapshot.Alias = "mutated"
	current, ok := reg.GetReady(first.NodeID)
	if !ok {
		t.Fatal("ready owner disappeared after snapshot mutation")
	}
	if current.Alias != "workspace" || string(current.CredentialRecipientPublicKey) != "first-key" {
		t.Fatalf("snapshot mutation leaked into registry: %#v", current)
	}

	if _, ok := reg.UnregisterIfClient(first.NodeID, firstClient); !ok {
		t.Fatal("initial owner did not unregister")
	}
	secondClient := &toki.TcpClient{}
	second := &edgenode.NodeEntry{NodeID: first.NodeID, Alias: "workspace", Client: secondClient}
	if !reg.RegisterIfAbsent(second) {
		t.Fatal("reconnect registration was rejected")
	}
	if _, transitioned, ok := reg.MarkDispatchReadyIfClient(second.NodeID, secondClient); !ok || !transitioned {
		t.Fatalf("reconnect ready = transitioned=%v ok=%v, want true,true", transitioned, ok)
	}
	reconnected, ok := reg.ReadyOwnerSnapshot(second.NodeID)
	if !ok || reconnected.ConnectionGeneration <= snapshot.ConnectionGeneration {
		t.Fatalf("reconnect snapshot generation = %#v, want greater than %d", reconnected, snapshot.ConnectionGeneration)
	}
	if reg.IsCurrentOwnerGeneration(second.NodeID, snapshot.ConnectionGeneration) {
		t.Fatal("old snapshot generation remained current after reconnect")
	}

	// Synchronize snapshot readers with repeated reconnects. Each transition has
	// an explicit unavailable and ready rendezvous, so this cannot pass merely
	// because one goroutine happened to finish before the other observed it.
	t.Run("concurrent reconnect snapshots remain self-consistent", func(t *testing.T) {
		concurrent := edgenode.NewRegistry()
		client := &toki.TcpClient{}
		concurrent.Register(&edgenode.NodeEntry{NodeID: "workspace-node", Alias: "workspace", Client: client})

		snapshotStep := make(chan int)
		transitionDone := make(chan int)
		errCh := make(chan error, 1)
		var wg sync.WaitGroup
		wg.Add(2)
		go func() {
			defer wg.Done()
			snapshot, ok := concurrent.ReadyOwnerSnapshot("workspace-node")
			if !ok || snapshot == nil || snapshot.NodeID != "workspace-node" || !snapshot.DispatchReady || snapshot.ConnectionGeneration == 0 {
				errCh <- fmt.Errorf("initial ready snapshot: %#v", snapshot)
				return
			}
			generation := snapshot.ConnectionGeneration
			for i := 0; i < 64; i++ {
				snapshotStep <- i
				if got := <-transitionDone; got != i {
					errCh <- fmt.Errorf("step %d unavailable transition=%d", i, got)
					return
				}
				if snapshot, ok := concurrent.ReadyOwnerSnapshot("workspace-node"); ok || snapshot != nil {
					errCh <- fmt.Errorf("step %d unavailable snapshot=%#v, %v", i, snapshot, ok)
					return
				}

				snapshotStep <- i
				if got := <-transitionDone; got != i {
					errCh <- fmt.Errorf("step %d ready transition=%d", i, got)
					return
				}
				snapshot, ok = concurrent.ReadyOwnerSnapshot("workspace-node")
				if !ok || snapshot == nil || snapshot.NodeID != "workspace-node" || !snapshot.DispatchReady || snapshot.ConnectionGeneration <= generation {
					errCh <- fmt.Errorf("step %d invalid ready snapshot: %#v", i, snapshot)
					return
				}
				generation = snapshot.ConnectionGeneration
			}
		}()
		go func() {
			defer wg.Done()
			for i := 0; i < 64; i++ {
				if step := <-snapshotStep; step != i {
					errCh <- fmt.Errorf("step %d reconnect request=%d", i, step)
					return
				}
				if _, ok := concurrent.UnregisterIfClient("workspace-node", client); !ok {
					errCh <- fmt.Errorf("reconnect %d failed to unregister current owner", i)
					return
				}
				transitionDone <- i
				if step := <-snapshotStep; step != i {
					errCh <- fmt.Errorf("step %d ready request=%d", i, step)
					return
				}
				client = &toki.TcpClient{}
				if !concurrent.RegisterIfAbsent(&edgenode.NodeEntry{NodeID: "workspace-node", Alias: "workspace", Client: client}) {
					errCh <- fmt.Errorf("reconnect %d registration rejected", i)
					return
				}
				if _, transitioned, ok := concurrent.MarkDispatchReadyIfClient("workspace-node", client); !ok || !transitioned {
					errCh <- fmt.Errorf("reconnect %d ready = transitioned=%v ok=%v", i, transitioned, ok)
					return
				}
				transitionDone <- i
			}
		}()
		wg.Wait()
		select {
		case err := <-errCh:
			t.Fatal(err)
		default:
		}
	})
}

// TestRegistryMarkDispatchReadyIdempotentForCurrentOwner pins that a duplicate
// ready for an already-ready owner reports ok=true, transitioned=false so the
// transport acks success without repeating pump/event, while a stale client is
// rejected outright.
func TestRegistryMarkDispatchReadyIdempotentForCurrentOwner(t *testing.T) {
	reg := edgenode.NewRegistry()
	owner := &toki.TcpClient{}
	stale := &toki.TcpClient{}
	reg.RegisterIfAbsent(&edgenode.NodeEntry{NodeID: "node-1", Client: owner})

	if _, transitioned, ok := reg.MarkDispatchReadyIfClient("node-1", owner); !ok || !transitioned {
		t.Fatalf("first ready: ok=%v transitioned=%v, want true/true", ok, transitioned)
	}
	if _, transitioned, ok := reg.MarkDispatchReadyIfClient("node-1", owner); !ok || transitioned {
		t.Fatalf("duplicate ready: ok=%v transitioned=%v, want true/false", ok, transitioned)
	}
	if _, _, ok := reg.MarkDispatchReadyIfClient("node-1", stale); ok {
		t.Fatal("stale client ready must be rejected (ok=false)")
	}
	if _, _, ok := reg.MarkDispatchReadyIfClient("missing", owner); ok {
		t.Fatal("ready for an unregistered id must be rejected (ok=false)")
	}
}

// TestRegistryReadyOnlyLookupsExcludePending pins that with a mix of pending and
// ready connections, the ready-only lookups return only the ready one and the
// implicit resolve is unambiguous.
func TestRegistryReadyOnlyLookupsExcludePending(t *testing.T) {
	reg := edgenode.NewRegistry()
	readyClient := &toki.TcpClient{}
	readyEntry := &edgenode.NodeEntry{NodeID: "node-ready", Alias: "ready-alias", Client: readyClient}
	reg.RegisterIfAbsent(readyEntry)
	reg.MarkDispatchReadyIfClient("node-ready", readyClient)
	reg.RegisterIfAbsent(&edgenode.NodeEntry{NodeID: "node-pending", Alias: "pending-alias", Client: &toki.TcpClient{}})

	ready := reg.AllReady()
	if len(ready) != 1 || ready[0].NodeID != "node-ready" {
		t.Fatalf("AllReady = %v, want only node-ready", ready)
	}
	if got, err := reg.ResolveReady(""); err != nil || got.NodeID != "node-ready" {
		t.Fatalf("implicit ready resolve should pick the only ready node: got=%v err=%v", got, err)
	}
	if _, err := reg.ResolveReady("pending-alias"); err == nil {
		t.Fatal("resolving a pending node by alias via ResolveReady must fail")
	}
	// Non-ready lookups still see both.
	if len(reg.All()) != 2 {
		t.Fatalf("All must include pending entries, got %d", len(reg.All()))
	}
}

func TestRegistryMarkDispatchReadyOwnerAndWithCurrentOwner(t *testing.T) {
	reg := edgenode.NewRegistry()
	client1 := &toki.TcpClient{}
	client2 := &toki.TcpClient{}

	entry := &edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1", Client: client1}
	if !reg.RegisterIfAbsent(entry) {
		t.Fatal("expected RegisterIfAbsent to succeed")
	}

	// 1. Verify MarkDispatchReadyOwner transitions and returns a clone.
	snap, transitioned, ok := reg.MarkDispatchReadyOwner("node-1", client1)
	if !ok || !transitioned || snap == nil {
		t.Fatalf("expected ok=true, transitioned=true: ok=%v transitioned=%v snap=%v", ok, transitioned, snap)
	}
	if snap.NodeID != "node-1" || snap.ConnectionGeneration != entry.ConnectionGeneration || !snap.DispatchReady {
		t.Errorf("incorrect snapshot fields: %+v", snap)
	}

	// 2. Verify WithCurrentOwner runs function for matching snapshot.
	run := false
	ok = reg.WithCurrentOwner(snap, func() {
		run = true
	})
	if !ok || !run {
		t.Errorf("expected WithCurrentOwner to execute callback: ok=%v run=%v", ok, run)
	}

	// 3. Verify WithCurrentOwner is no-op if owner unregistered.
	reg.Unregister("node-1")
	run = false
	ok = reg.WithCurrentOwner(snap, func() {
		run = true
	})
	if ok || run {
		t.Errorf("expected WithCurrentOwner to skip callback after unregister: ok=%v run=%v", ok, run)
	}

	// 4. Verify WithCurrentOwner is no-op for stale generation after reconnect.
	entry2 := &edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1", Client: client2}
	if !reg.RegisterIfAbsent(entry2) {
		t.Fatal("expected RegisterIfAbsent for reconnect to succeed")
	}
	// Note: entry2 has a higher connection generation now.
	run = false
	ok = reg.WithCurrentOwner(snap, func() {
		run = true
	})
	if ok || run {
		t.Errorf("expected WithCurrentOwner to skip callback for stale generation: ok=%v run=%v", ok, run)
	}
}

func TestCurrentOwnerForClient(t *testing.T) {
	reg := edgenode.NewRegistry()
	client1 := &toki.TcpClient{}
	client2 := &toki.TcpClient{}

	if _, ok := reg.CurrentOwnerForClient(nil); ok {
		t.Fatal("nil client should return false")
	}
	if _, ok := reg.CurrentOwnerForClient(client1); ok {
		t.Fatal("unregistered client should return false")
	}

	entry1 := &edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1", Client: client1}
	reg.RegisterIfAbsent(entry1)

	got, ok := reg.CurrentOwnerForClient(client1)
	if !ok || got == nil {
		t.Fatal("registered client1 should return entry and true")
	}
	if got.NodeID != "node-1" || got.ConnectionGeneration != entry1.ConnectionGeneration {
		t.Fatalf("unexpected entry for client1: %+v", got)
	}

	if _, ok := reg.CurrentOwnerForClient(client2); ok {
		t.Fatal("unregistered client2 should return false")
	}

	// Reconnect with client2 for same node ID
	reg.UnregisterIfClient("node-1", client1)
	entry2 := &edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1", Client: client2}
	reg.RegisterIfAbsent(entry2)

	if _, ok := reg.CurrentOwnerForClient(client1); ok {
		t.Fatal("stale client1 should return false after reconnect")
	}
	got2, ok := reg.CurrentOwnerForClient(client2)
	if !ok || got2 == nil {
		t.Fatal("reconnected client2 should return entry and true")
	}
	if got2.ConnectionGeneration <= got.ConnectionGeneration {
		t.Fatalf("reconnected generation %d must exceed previous %d", got2.ConnectionGeneration, got.ConnectionGeneration)
	}
}

func TestRegistryRegisterIfAbsentRejectsClientRebinding(t *testing.T) {
	reg := edgenode.NewRegistry()
	client := &toki.TcpClient{}

	first := &edgenode.NodeEntry{NodeID: "node-a", Alias: "alias-a", Client: client}
	if !reg.RegisterIfAbsent(first) {
		t.Fatal("first registration should succeed")
	}
	genA := first.ConnectionGeneration

	second := &edgenode.NodeEntry{NodeID: "node-b", Alias: "alias-b", Client: client}
	if reg.RegisterIfAbsent(second) {
		t.Fatal("second registration with same non-nil client should be rejected")
	}

	if reg.Count() != 1 {
		t.Fatalf("registry count: got %d want 1", reg.Count())
	}
	gotA, ok := reg.Get("node-a")
	if !ok || gotA.Client != client || gotA.ConnectionGeneration != genA {
		t.Fatalf("node-a owner/generation altered after rejected client rebinding attempt: %+v", gotA)
	}
	if _, ok := reg.Get("node-b"); ok {
		t.Fatal("node-b should not exist in registry")
	}
}

func TestCurrentOwnerForClientFailsClosedForAmbiguousClient(t *testing.T) {
	reg := edgenode.NewRegistry()
	client := &toki.TcpClient{}

	// Construct an ambiguous state directly via unconditional Register helper.
	reg.Register(&edgenode.NodeEntry{NodeID: "node-a", Alias: "alias-a", Client: client})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-b", Alias: "alias-b", Client: client})

	if reg.Count() != 2 {
		t.Fatalf("registry count: got %d want 2", reg.Count())
	}
	if got, ok := reg.CurrentOwnerForClient(client); ok || got != nil {
		t.Fatalf("CurrentOwnerForClient for ambiguous client should fail closed (nil, false), got %+v, %v", got, ok)
	}
}
