iop/apps/edge/internal/service/model_queue_test.go

package service

import (
	"context"
	"errors"
	"net"
	"reflect"
	"sync"
	"testing"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

// waitForQueueLen polls until the group queue reaches wantLen or times out.
func waitForQueueLen(t *testing.T, m *modelQueueManager, groupKey string, wantLen int) {
	t.Helper()
	deadline := time.Now().Add(200 * time.Millisecond)
	for time.Now().Before(deadline) {
		m.mu.Lock()
		g, ok := m.groups[groupKey]
		got := 0
		if ok {
			got = len(g.queue)
		}
		m.mu.Unlock()
		if got >= wantLen {
			return
		}
		time.Sleep(1 * time.Millisecond)
	}
	t.Fatalf("timeout: queue for %q did not reach length %d", groupKey, wantLen)
}

// TestModelQueueFIFOOrdering verifies that queued items are dispatched in FIFO
// order when a slot becomes available.
func TestModelQueueFIFOOrdering(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-q1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-q1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(store)

	// Fill the only capacity slot.
	first, err := m.admit(context.Background(), "g-fifo", "", "", cands, defPolicy)
	if err != nil || first == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Manually insert two items in known order so we can verify FIFO.
	item1 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	item2 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-fifo", groupPolicy{})
	g.queue = append(g.queue, item1, item2)
	m.mu.Unlock()

	// Release one slot — item1 (head) must be dispatched first.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1", "", false)
	m.mu.Unlock()

	select {
	case res := <-item1.waitCh:
		if res.err != nil {
			t.Fatalf("item1 expected dispatch, got error: %v", res.err)
		}
		if res.candidate == nil || res.candidate.entry.NodeID != "node-q1" {
			t.Fatalf("item1: unexpected candidate %v", res.candidate)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item1 was not dispatched after slot release")
	}

	// item2 must still be waiting (no slot available yet).
	select {
	case <-item2.waitCh:
		t.Fatal("item2 should not have been dispatched yet")
	default:
	}

	// Release the slot item1 holds so item2 gets dispatched.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1", "", false)
	m.mu.Unlock()

	select {
	case res := <-item2.waitCh:
		if res.err != nil {
			t.Fatalf("item2 expected dispatch, got error: %v", res.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item2 was not dispatched after second slot release")
	}
}

// TestModelQueueOverflow verifies that admit rejects requests when the queue
// is at max_queue capacity.
func TestModelQueueOverflow(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ov1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-ov1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill the node's capacity and set max_queue=1.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-overflow", groupPolicy{})
	g.policy.maxQueue = 1
	g.inflight["node-ov1"] = 1
	m.mu.Unlock()

	// Queue one item — should succeed.
	item := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	m.groups["g-overflow"].queue = append(m.groups["g-overflow"].queue, item)
	m.mu.Unlock()

	// Second admit should fail immediately with errQueueFull.
	_, err := m.admit(context.Background(), "g-overflow", "", "", cands, groupPolicy{})
	if !errors.Is(err, errQueueFull) {
		t.Fatalf("expected errQueueFull, got: %v", err)
	}
}

// TestModelQueueTimeout verifies that queued items receive a timeout error
// when no slot becomes available before the deadline.
func TestModelQueueTimeout(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-to1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-to1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill capacity and set a very short queue timeout.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-timeout", groupPolicy{})
	g.policy.queueTimeout = 20 * time.Millisecond
	g.inflight["node-to1"] = 1
	m.mu.Unlock()

	start := time.Now()
	_, err := m.admit(context.Background(), "g-timeout", "", "", cands, groupPolicy{})
	elapsed := time.Since(start)

	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("expected errQueueTimeout, got: %v", err)
	}
	if elapsed < 15*time.Millisecond {
		t.Fatalf("timed out too fast: %v", elapsed)
	}
}

// TestModelQueueTerminalReleaseDispatchesNext verifies that a terminal run event
// releases the in-flight slot and dispatches the next queued item.
func TestModelQueueTerminalReleaseDispatchesNext(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-tr1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-tr1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	bus := edgeevents.NewBus()
	m := newModelQueueManager(store)
	stop := m.startEventWatcher(bus)
	defer stop()

	// Fill capacity and record inflight.
	selected, err := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy)
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	m.trackInflight("g-tr", "run-tr-001", selected.entry.NodeID, selected.providerID)

	// Queue a second item in a goroutine.
	resultCh := make(chan admitResult, 1)
	go func() {
		n, e := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy)
		resultCh <- admitResult{candidate: n, err: e}
	}()

	waitForQueueLen(t, m, "g-tr", 1)

	// Terminal event for the first run — should unblock the queued admit.
	for _, termType := range []string{"complete", "error", "cancelled"} {
		t.Run(termType, func(t *testing.T) {
			// Reset state for each sub-test.
			m.mu.Lock()
			g := m.getOrCreateGroupLocked("g-tr-"+termType, groupPolicy{})
			g.inflight["node-tr1"] = 1
			item := &queueItem{
				candidates: cands,
				waitCh:     make(chan admitResult, 1),
				deadline:   time.Now().Add(2 * time.Second),
			}
			g.queue = []*queueItem{item}
			runID := "run-tr-" + termType
			m.inflightByRun[runID] = inflightRec{groupKey: "g-tr-" + termType, nodeID: "node-tr1"}
			m.mu.Unlock()

			bus.PublishRun(&iop.RunEvent{RunId: runID, Type: termType})

			select {
			case res := <-item.waitCh:
				if res.err != nil {
					t.Fatalf("expected dispatch, got error: %v", res.err)
				}
				if res.candidate == nil {
					t.Fatal("expected non-nil candidate")
				}
			case <-time.After(500 * time.Millisecond):
				t.Fatalf("timeout: queued item not dispatched after %q terminal event", termType)
			}
		})
	}

	// Cleanup the goroutine from the outer admit.
	m.mu.Lock()
	if g, ok := m.groups["g-tr"]; ok && len(g.queue) > 0 {
		g.inflight["node-tr1"] = 1
		m.tryDispatchLocked(g)
	}
	m.mu.Unlock()
	select {
	case <-resultCh:
	case <-time.After(100 * time.Millisecond):
	}
}

// TestModelQueueNodeDisconnectReleasesInflight verifies that a node disconnect
// event removes the node from queued item candidate lists, and that subsequent
// dispatch goes to a remaining live candidate (not the disconnected node).
func TestModelQueueNodeDisconnectReleasesInflight(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd2",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry1 := &edgenode.NodeEntry{NodeID: "node-nd1"}
	entry2 := &edgenode.NodeEntry{NodeID: "node-nd2"}
	// Queued item that can use either node (both at capacity 1).
	bothCands := []candidateNode{
		{entry: entry1, capacity: 1},
		{entry: entry2, capacity: 1},
	}

	bus := edgeevents.NewBus()
	m := newModelQueueManager(store)
	stop := m.startEventWatcher(bus)
	defer stop()

	// Fill node-nd1 and node-nd2 capacities and record inflight.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-nd", groupPolicy{})
	g.inflight["node-nd1"] = 1
	g.inflight["node-nd2"] = 1
	m.inflightByRun["run-nd-x"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd1"}
	m.inflightByRun["run-nd-y"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd2"}
	// Queue an item that can use either node.
	item := &queueItem{
		candidates: bothCands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	g.queue = []*queueItem{item}
	m.mu.Unlock()

	// node-nd1 disconnects: its inflight is freed and it is removed from candidates.
	// node-nd2 is still full, so no dispatch yet.
	bus.PublishNode(&iop.EdgeNodeEvent{
		NodeId: "node-nd1",
		Type:   "node.disconnected",
	})

	// node-nd2's run terminates: now nd2 has capacity, dispatch goes to nd2.
	bus.PublishRun(&iop.RunEvent{RunId: "run-nd-y", Type: "complete"})

	select {
	case res := <-item.waitCh:
		if res.err != nil {
			t.Fatalf("expected dispatch after disconnect+terminal, got error: %v", res.err)
		}
		if res.candidate == nil {
			t.Fatal("expected non-nil candidate after disconnect release")
		}
		// Must be nd2 — nd1 was removed from candidates on disconnect.
		if res.candidate.entry.NodeID != "node-nd2" {
			t.Errorf("expected dispatch to node-nd2, got %q (nd1 was disconnected)", res.candidate.entry.NodeID)
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout: queued item not dispatched after node disconnect and terminal event")
	}

	// The inflight record for run-nd-x should be gone (removed at disconnect).
	m.mu.Lock()
	_, stillInFlight := m.inflightByRun["run-nd-x"]
	m.mu.Unlock()
	if stillInFlight {
		t.Error("run-nd-x should be removed from inflightByRun after node disconnect")
	}
}

// TestModelQueueUsesProviderCapacity verifies that the capacity supplied in the
// candidateNode (derived from adapter config) overrides the default of 1.
func TestModelQueueUsesProviderCapacity(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-pc1"}
	// Provider capacity = 2 (two concurrent slots on this node).
	cands := []candidateNode{{entry: entry, capacity: 2}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(nil)

	// First admit: inflight=0 < cap=2 → dispatched immediately.
	n1, err := m.admit(context.Background(), "g-pc", "", "", cands, defPolicy)
	if err != nil || n1 == nil {
		t.Fatalf("first admit: %v", err)
	}

	// Second admit: inflight=1 < cap=2 → still dispatched (not queued).
	n2, err := m.admit(context.Background(), "g-pc", "", "", cands, defPolicy)
	if err != nil || n2 == nil {
		t.Fatalf("second admit (capacity=2 should allow): %v", err)
	}

	// Third request must queue because inflight=2 == cap=2.
	item := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(200 * time.Millisecond),
	}
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-pc", groupPolicy{})
	g.queue = append(g.queue, item)
	m.mu.Unlock()

	// Releasing one slot should dispatch the queued item.
	m.releaseSlot("g-pc", "node-pc1")

	select {
	case res := <-item.waitCh:
		if res.err != nil {
			t.Fatalf("expected dispatch after slot release, got: %v", res.err)
		}
		if res.candidate == nil || res.candidate.entry.NodeID != "node-pc1" {
			t.Fatalf("unexpected candidate: %v", res.candidate)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout: item not dispatched after slot release")
	}
}

// TestRefreshProviderQueuePolicyUpdatesExistingGroup verifies that a tightened
// queue policy (e.g. provider max_queue lowered via config refresh) propagates
// to an already-created group on the next admission, while in-flight slots are
// preserved.
func TestRefreshProviderQueuePolicyUpdatesExistingGroup(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-rp"}
	cands := []candidateNode{{entry: entry, capacity: 2}}
	m := newModelQueueManager(nil)

	// Create the group with a generous policy and take one slot.
	initial := groupPolicy{maxQueue: 8, queueTimeout: 5 * time.Second}
	if _, err := m.admit(context.Background(), "g-rp", "", "", cands, initial); err != nil {
		t.Fatalf("first admit: %v", err)
	}
	m.mu.Lock()
	if got := m.groups["g-rp"].policy.maxQueue; got != 8 {
		m.mu.Unlock()
		t.Fatalf("group created with maxQueue=%d, want 8", got)
	}
	m.mu.Unlock()

	// Config refresh tightens the policy to maxQueue=1; the recomputed admit
	// carries it and must update the existing group's policy. capacity=2 means
	// this second admit is dispatched (not queued).
	refreshed := groupPolicy{maxQueue: 1, queueTimeout: 5 * time.Second}
	if _, err := m.admit(context.Background(), "g-rp", "", "", cands, refreshed); err != nil {
		t.Fatalf("second admit: %v", err)
	}

	m.mu.Lock()
	g := m.groups["g-rp"]
	if g.policy.maxQueue != 1 {
		m.mu.Unlock()
		t.Fatalf("existing group policy not refreshed: maxQueue=%d, want 1", g.policy.maxQueue)
	}
	// Both slots now in-flight (cap=2). Queue one item to reach the refreshed
	// maxQueue=1 so the next admission must overflow.
	g.queue = append(g.queue, &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(5 * time.Second),
	})
	m.mu.Unlock()

	if _, err := m.admit(context.Background(), "g-rp", "", "", cands, refreshed); !errors.Is(err, errQueueFull) {
		t.Fatalf("expected errQueueFull under refreshed maxQueue=1, got: %v", err)
	}
}

// TestModelQueueUsesProviderQueuePolicy verifies that max_queue and
// queue_timeout from the policy parameter are respected.
func TestModelQueueUsesProviderQueuePolicy(t *testing.T) {
	t.Run("maxQueue enforced", func(t *testing.T) {
		entry := &edgenode.NodeEntry{NodeID: "node-pq1"}
		cands := []candidateNode{{entry: entry, capacity: 1}}
		// Provider policy: maxQueue=1.
		policy := groupPolicy{maxQueue: 1, queueTimeout: 5 * time.Second}

		m := newModelQueueManager(nil)

		// Fill capacity and queue one item manually.
		m.mu.Lock()
		g := m.getOrCreateGroupLocked("g-pq-max", policy)
		g.inflight["node-pq1"] = 1
		g.queue = []*queueItem{{
			candidates: cands,
			waitCh:     make(chan admitResult, 1),
			deadline:   time.Now().Add(5 * time.Second),
		}}
		m.mu.Unlock()

		// Second admit should fail: queue already at maxQueue=1.
		_, err := m.admit(context.Background(), "g-pq-max", "", "", cands, policy)
		if !errors.Is(err, errQueueFull) {
			t.Fatalf("expected errQueueFull, got: %v", err)
		}
	})

	t.Run("queueTimeout enforced", func(t *testing.T) {
		entry := &edgenode.NodeEntry{NodeID: "node-pq2"}
		cands := []candidateNode{{entry: entry, capacity: 1}}
		// Provider policy: very short timeout.
		policy := groupPolicy{maxQueue: 16, queueTimeout: 20 * time.Millisecond}

		m := newModelQueueManager(nil)

		// Fill capacity.
		m.mu.Lock()
		g := m.getOrCreateGroupLocked("g-pq-to", policy)
		g.inflight["node-pq2"] = 1
		m.mu.Unlock()

		start := time.Now()
		_, err := m.admit(context.Background(), "g-pq-to", "", "", cands, policy)
		elapsed := time.Since(start)

		if !errors.Is(err, errQueueTimeout) {
			t.Fatalf("expected errQueueTimeout, got: %v", err)
		}
		if elapsed < 15*time.Millisecond {
			t.Fatalf("timed out too fast: %v", elapsed)
		}
	})
}

// TestModelQueueProviderInflightSelectionBeatsPriority verifies that in_flight
// level is the primary selection criterion: a candidate with lower in_flight is
// preferred even if it has a higher (worse) priority value.
func TestModelQueueProviderInflightSelectionBeatsPriority(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-ip-a"}
	entryB := &edgenode.NodeEntry{NodeID: "node-ip-b"}
	// A has priority=1 (better) but inflight=3.
	// B has priority=10 (worse) but inflight=0.
	cands := []candidateNode{
		{entry: entryA, capacity: 5, priority: 1, providerID: "prov-a"},
		{entry: entryB, capacity: 5, priority: 10, providerID: "prov-b"},
	}
	m := newModelQueueManager(nil)

	// Put A at inflight=3, B at inflight=0.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-ip", groupPolicy{})
	g.inflight["node-ip-a:prov-a"] = 3
	m.mu.Unlock()

	// Admit should pick B (inflight=0 < inflight=3), not A.
	sel, err := m.admit(context.Background(), "g-ip", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	if sel == nil || sel.entry.NodeID != "node-ip-b" {
		t.Fatalf("expected node-ip-b (lower in_flight), got %v", sel)
	}
}

// TestModelQueueProviderFullPriorityFallsThrough verifies that priority does
// not override capacity: a full high-priority provider is skipped and the next
// available priority tier receives the request.
func TestModelQueueProviderFullPriorityFallsThrough(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-full-a"}
	entryB := &edgenode.NodeEntry{NodeID: "node-next-b"}
	cands := []candidateNode{
		{entry: entryA, capacity: 4, priority: 1, providerID: "prov-fast"},
		{entry: entryB, capacity: 3, priority: 10, providerID: "prov-slow"},
	}
	m := newModelQueueManager(nil)

	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-full-priority", groupPolicy{})
	g.inflight["node-full-a:prov-fast"] = 4
	m.mu.Unlock()

	sel, err := m.admit(context.Background(), "g-full-priority", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	if sel == nil || sel.providerID != "prov-slow" {
		t.Fatalf("expected prov-slow after high-priority provider is full, got %v", sel)
	}
}

// TestModelQueueProviderLevelingPriorityAdmissionOrder fixes the intended
// provider-pool order for differently-sized providers: fill each in_flight
// level before moving to the next level, and use priority within the same
// level. With capacities 4/3/2 and priorities 0/1/2, the first nine admits are
// gx10, onex, mac, gx10, onex, mac, gx10, onex, gx10. The tenth has no
// available candidate until a slot is released.
func TestModelQueueProviderLevelingPriorityAdmissionOrder(t *testing.T) {
	cands := []candidateNode{
		{entry: &edgenode.NodeEntry{NodeID: "node-gx10"}, capacity: 4, priority: 0, providerID: "gx10-vllm"},
		{entry: &edgenode.NodeEntry{NodeID: "node-onex"}, capacity: 3, priority: 1, providerID: "onexplayer-lemonade"},
		{entry: &edgenode.NodeEntry{NodeID: "node-mac"}, capacity: 2, priority: 2, providerID: "mac-mlx-vllm"},
	}
	m := newModelQueueManager(nil)

	var got []string
	for i := 0; i < 9; i++ {
		sel, err := m.admit(context.Background(), "qwen3.6:35b", "", "", cands, groupPolicy{})
		if err != nil {
			t.Fatalf("admit %d: %v", i+1, err)
		}
		got = append(got, sel.providerID)
	}
	want := []string{
		"gx10-vllm",
		"onexplayer-lemonade",
		"mac-mlx-vllm",
		"gx10-vllm",
		"onexplayer-lemonade",
		"mac-mlx-vllm",
		"gx10-vllm",
		"onexplayer-lemonade",
		"gx10-vllm",
	}
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("admission order mismatch:\n got: %v\nwant: %v", got, want)
	}

	m.mu.Lock()
	g := m.getOrCreateGroupLocked("qwen3.6:35b", groupPolicy{})
	if candidate := m.findAvailableNodeLocked(g, cands, false); candidate != nil {
		t.Fatalf("expected no available candidate after filling capacity, got %s", candidate.providerID)
	}
	m.mu.Unlock()
}

// TestModelQueueProviderPriorityBreaksEqualInflightTie verifies that when
// in_flight counts are equal, the candidate with the lower priority value
// is selected.
func TestModelQueueProviderPriorityBreaksEqualInflightTie(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-pt-a"}
	entryB := &edgenode.NodeEntry{NodeID: "node-pt-b"}
	// Both have inflight=2, but A has priority=1 (better) than B (priority=5).
	cands := []candidateNode{
		{entry: entryA, capacity: 5, priority: 1, providerID: "prov-a"},
		{entry: entryB, capacity: 5, priority: 5, providerID: "prov-b"},
	}
	m := newModelQueueManager(nil)

	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-pt", groupPolicy{})
	g.inflight["node-pt-a:prov-a"] = 2
	g.inflight["node-pt-b:prov-b"] = 2
	m.mu.Unlock()

	// Admit should pick A (lower priority=1 when inflight is equal).
	sel, err := m.admit(context.Background(), "g-pt", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	if sel == nil || sel.entry.NodeID != "node-pt-a" {
		t.Fatalf("expected node-pt-a (lower priority tie-break), got %v", sel)
	}
}

// TestModelQueueProviderEqualInflightPriorityRotates verifies that when both
// in_flight and priority are equal, the rotation logic selects the next slot
// after the last selected one.
func TestModelQueueProviderEqualInflightPriorityRotates(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-r-a"}
	entryB := &edgenode.NodeEntry{NodeID: "node-r-b"}
	// Both have inflight=0 and priority=0.
	cands := []candidateNode{
		{entry: entryA, capacity: 5, priority: 0, providerID: "prov-a"},
		{entry: entryB, capacity: 5, priority: 0, providerID: "prov-b"},
	}
	m := newModelQueueManager(nil)

	// First admit: A wins deterministically.
	sel1, err := m.admit(context.Background(), "g-r", "", "", cands, groupPolicy{})
	if err != nil || sel1 == nil {
		t.Fatalf("first admit: %v", err)
	}
	if sel1.providerID != "prov-a" {
		t.Fatalf("expected prov-a on first admit, got %q", sel1.providerID)
	}

	// Release prov-a's slot.
	m.releaseSlot("g-r", "node-r-a", "prov-a")

	// Second admit: rotation should pick B (next after last selected A).
	sel2, err := m.admit(context.Background(), "g-r", "", "", cands, groupPolicy{})
	if err != nil || sel2 == nil {
		t.Fatalf("second admit: %v", err)
	}
	if sel2.providerID != "prov-b" {
		t.Fatalf("expected prov-b on second admit (rotation), got %q", sel2.providerID)
	}

	// Release prov-b's slot.
	m.releaseSlot("g-r", "node-r-b", "prov-b")

	// Third admit: rotation should cycle back to A.
	sel3, err := m.admit(context.Background(), "g-r", "", "", cands, groupPolicy{})
	if err != nil || sel3 == nil {
		t.Fatalf("third admit: %v", err)
	}
	if sel3.providerID != "prov-a" {
		t.Fatalf("expected prov-a on third admit (rotation cycle), got %q", sel3.providerID)
	}
}

func TestModelQueueProviderEqualIdleTieRotates(t *testing.T) {
	gx10 := &edgenode.NodeEntry{NodeID: "node-gx10-idle"}
	onex := &edgenode.NodeEntry{NodeID: "node-onex-idle"}
	cands := []candidateNode{
		{entry: gx10, capacity: 4, providerID: "gx10-vllm"},
		{entry: onex, capacity: 3, providerID: "onexplayer-lemonade"},
	}
	m := newModelQueueManager(nil)

	sel1, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("first admit: %v", err)
	}
	if sel1.providerID != "gx10-vllm" {
		t.Fatalf("first admit provider: got %q want gx10-vllm", sel1.providerID)
	}
	m.releaseSlot("g-idle-rotate", sel1.entry.NodeID, sel1.providerID)

	sel2, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("second admit: %v", err)
	}
	if sel2.providerID != "onexplayer-lemonade" {
		t.Fatalf("second admit provider: got %q want onexplayer-lemonade", sel2.providerID)
	}
	m.releaseSlot("g-idle-rotate", sel2.entry.NodeID, sel2.providerID)

	sel3, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("third admit: %v", err)
	}
	if sel3.providerID != "gx10-vllm" {
		t.Fatalf("third admit provider: got %q want gx10-vllm", sel3.providerID)
	}
}

// TestModelQueueProviderServedTargetRewrite verifies that the selected
// candidateNode carries its servedTarget so callers can rewrite req.Target.
func TestModelQueueProviderServedTargetRewrite(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-tr-rewrite"}
	cands := []candidateNode{{
		entry:        entry,
		capacity:     2,
		providerID:   "prov-vllm",
		servedTarget: "qwen3-72b-instruct",
	}}
	m := newModelQueueManager(nil)

	sel, err := m.admit(context.Background(), "g-tr-rewrite", "", "", cands, groupPolicy{})
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	if sel == nil {
		t.Fatal("expected non-nil candidate")
	}
	if sel.servedTarget != "qwen3-72b-instruct" {
		t.Errorf("servedTarget: got %q, want %q", sel.servedTarget, "qwen3-72b-instruct")
	}
	if sel.providerID != "prov-vllm" {
		t.Errorf("providerID: got %q, want %q", sel.providerID, "prov-vllm")
	}
}

// TestProviderStatusInflightTracking verifies that getStatsForProviderLocked
// correctly counts in-flight runs keyed by (nodeID, providerID).
func TestProviderStatusInflightTracking(t *testing.T) {
	m := newModelQueueManager(nil)

	m.mu.Lock()
	m.inflightByRun["run-p1"] = inflightRec{groupKey: "g-alias", nodeID: "node-x", providerID: "prov-1"}
	m.inflightByRun["run-p2"] = inflightRec{groupKey: "g-alias", nodeID: "node-x", providerID: "prov-1"}
	m.inflightByRun["run-p3"] = inflightRec{groupKey: "g-alias2", nodeID: "node-x", providerID: "prov-2"}
	m.mu.Unlock()

	m.mu.Lock()
	inf1, q1 := m.getStatsForProviderLocked("node-x", "prov-1")
	inf2, q2 := m.getStatsForProviderLocked("node-x", "prov-2")
	inf3, q3 := m.getStatsForProviderLocked("node-y", "prov-1")
	m.mu.Unlock()

	if inf1 != 2 || q1 != 0 {
		t.Errorf("prov-1 on node-x: inflight=%d queued=%d, want 2/0", inf1, q1)
	}
	if inf2 != 1 || q2 != 0 {
		t.Errorf("prov-2 on node-x: inflight=%d queued=%d, want 1/0", inf2, q2)
	}
	if inf3 != 0 || q3 != 0 {
		t.Errorf("prov-1 on node-y: inflight=%d queued=%d, want 0/0", inf3, q3)
	}
}

// TestModelQueueCancelledTerminalReleasesProviderInflight verifies that a
// cancelled terminal RunEvent releases provider in-flight accounting back to
// zero, the same as complete/error. This guards the BUG-2 recovery contract:
// once Edge OpenAI cancel propagation reaches Node and Node replies with a
// cancelled terminal event, the provider snapshot's in_flight must not stay
// stuck positive.
func TestModelQueueCancelledTerminalReleasesProviderInflight(t *testing.T) {
	m := newModelQueueManager(nil)
	bus := edgeevents.NewBus()
	stop := m.startEventWatcher(bus)
	defer stop()

	m.trackInflight("g-cancel", "run-cancel-1", "node-x", "prov-1")

	m.mu.Lock()
	inf, _ := m.getStatsForProviderLocked("node-x", "prov-1")
	m.mu.Unlock()
	if inf != 1 {
		t.Fatalf("inflight before cancel: got %d, want 1", inf)
	}

	bus.PublishRun(&iop.RunEvent{RunId: "run-cancel-1", Type: "cancelled"})

	deadline := time.Now().Add(200 * time.Millisecond)
	for {
		m.mu.Lock()
		inf, _ = m.getStatsForProviderLocked("node-x", "prov-1")
		m.mu.Unlock()
		if inf == 0 {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("inflight did not recover to 0 after cancelled terminal event, got %d", inf)
		}
		time.Sleep(1 * time.Millisecond)
	}
}

// TestModelQueueContextCancelRemovesQueuedItem verifies that cancelling the
// context of a queued admit removes the item and returns context.Canceled.
func TestModelQueueContextCancelRemovesQueuedItem(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-cc1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(nil)

	// Fill the only slot.
	n, err := m.admit(context.Background(), "g-cc", "", "", cands, defPolicy)
	if err != nil || n == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Queue a second admit with a cancellable context.
	ctx, cancel := context.WithCancel(context.Background())
	resultCh := make(chan error, 1)
	go func() {
		_, err := m.admit(ctx, "g-cc", "", "", cands, defPolicy)
		resultCh <- err
	}()

	waitForQueueLen(t, m, "g-cc", 1)
	cancel()

	select {
	case err := <-resultCh:
		if !errors.Is(err, context.Canceled) {
			t.Errorf("expected context.Canceled, got: %v", err)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout: context cancellation not handled")
	}

	// Queue must be empty after cancellation.
	m.mu.Lock()
	qLen := 0
	if g, ok := m.groups["g-cc"]; ok {
		qLen = len(g.queue)
	}
	m.mu.Unlock()
	if qLen != 0 {
		t.Errorf("queue should be empty after cancel, got %d items", qLen)
	}
}

// TestModelQueueProviderCapacityIsPerProviderSlot verifies that same-node
// multiple provider candidates each have independent capacity/in-flight
// accounting per REVIEW_API-2.
func TestModelQueueProviderCapacityIsPerProviderSlot(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-same"}
	// Two provider candidates on the same node, each capacity=1.
	cands := []candidateNode{
		{entry: entry, capacity: 1, providerID: "prov-a", adapter: "vllm"},
		{entry: entry, capacity: 1, providerID: "prov-b", adapter: "vllm"},
	}
	m := newModelQueueManager(nil)

	// First admit: pick prov-a (deterministic tie-break by providerID).
	sel1, err := m.admit(context.Background(), "g-same", "", "", cands, groupPolicy{})
	if err != nil || sel1 == nil {
		t.Fatalf("first admit: %v", err)
	}
	if sel1.providerID != "prov-a" {
		t.Fatalf("expected prov-a, got %s", sel1.providerID)
	}

	// Second admit: should pick prov-b independently (prov-a is full, prov-b has capacity).
	sel2, err := m.admit(context.Background(), "g-same", "", "", cands, groupPolicy{})
	if err != nil || sel2 == nil {
		t.Fatalf("second admit (prov-b should be available): %v", err)
	}
	if sel2.providerID != "prov-b" {
		t.Fatalf("expected prov-b, got %s", sel2.providerID)
	}

	// Third admit: both providers at capacity, should queue.
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	_, err = m.admit(ctx, "g-same", "", "", cands, groupPolicy{})
	if !errors.Is(err, errQueueTimeout) && !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("expected queue timeout or deadline exceeded, got: %v", err)
	}
}

// TestResolveProviderPoolCandidatesFiltersInvalidProviders verifies that
// unavailable providers, served model mismatch, empty adapter providers,
// and capacity zero/unknown providers are excluded from dispatch candidates.
// This test calls the actual Service.resolveProviderPoolCandidates method.
func TestResolveProviderPoolCandidatesFiltersInvalidProviders(t *testing.T) {
	// Build model catalog entry for "qwen3.6:35b".
	catalog := []config.ModelCatalogEntry{
		{
			ID: "qwen3.6:35b",
			Providers: map[string]string{
				"prov-available":   "served-qwen",
				"prov-unavailable": "served-qwen",
				"prov-mismatch":    "served-qwen",
				"prov-no-adapter":  "served-qwen",
				"prov-cap-zero":    "served-qwen",
				"prov-cap-unknown": "served-qwen",
			},
		},
	}

	// Build NodeStore with multiple providers.
	store := edgenode.NewNodeStore()

	// Valid provider: available, has adapter, has capacity, served model matches provider's models.
	store.Add(&edgenode.NodeRecord{
		ID:      "node-valid",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-available",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen", "served-llama"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Invalid: health = "unavailable".
	store.Add(&edgenode.NodeRecord{
		ID:      "node-bad-health",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-unavailable",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "unavailable",
				Capacity: 2,
			},
		},
	})

	// Invalid: served model not in provider's own models list.
	store.Add(&edgenode.NodeRecord{
		ID:      "node-mismatch",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-mismatch",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-llama"}, // does NOT include "served-qwen"
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Provider-first (empty adapter, health available, capacity > 0): now a valid candidate.
	// adapter key = provider ID "prov-no-adapter".
	store.Add(&edgenode.NodeRecord{
		ID:      "node-no-adapter",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-no-adapter",
				Type:     "vllm",
				Adapter:  "", // provider-first: adapter key derived from ID
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Invalid: capacity = 0.
	store.Add(&edgenode.NodeRecord{
		ID:      "node-cap-zero",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-cap-zero",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 0,
			},
		},
	})

	// Invalid: capacity < 0 (negative/unknown).
	store.Add(&edgenode.NodeRecord{
		ID:      "node-cap-unknown",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-cap-unknown",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: -1,
			},
		},
	})

	// Build a fake registry with all nodes.
	reg := edgenode.NewRegistry()
	allRecs := store.All()
	for _, rec := range allRecs {
		entry := &edgenode.NodeEntry{
			NodeID:         rec.ID,
			LifecycleState: edgenode.LifecycleConnected,
		}
		reg.Register(entry)
	}

	// Create Service with catalog and node store.
	svc := New(reg, nil)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// Call the actual resolveProviderPoolCandidates.
	req := SubmitRunRequest{
		ModelGroupKey: "qwen3.6:35b",
		ProviderPool:  true,
	}
	storeSnapshot, catalogSnapshot := svc.runtimeConfigSnapshot()
	candidates, policy, err := svc.resolveProviderPoolCandidates(req, storeSnapshot, catalogSnapshot)
	if err != nil {
		t.Fatalf("resolveProviderPoolCandidates: %v", err)
	}

	// prov-available (legacy/compat) and prov-no-adapter (provider-first) both pass.
	// All others are filtered: unavailable health, model mismatch, cap=0, cap<0.
	if len(candidates) != 2 {
		ids := make([]string, len(candidates))
		for i, c := range candidates {
			ids[i] = c.providerID
		}
		t.Fatalf("expected 2 candidates, got %d: %v", len(candidates), ids)
	}

	byID := map[string]*candidateNode{}
	for i := range candidates {
		byID[candidates[i].providerID] = &candidates[i]
	}

	// Legacy/compat candidate uses explicit adapter key.
	avail := byID["prov-available"]
	if avail == nil {
		t.Fatal("prov-available not in candidates")
	}
	if avail.adapter != "vllm-gpu" {
		t.Errorf("prov-available adapter: got %q, want %q", avail.adapter, "vllm-gpu")
	}
	if avail.servedTarget != "served-qwen" {
		t.Errorf("prov-available servedTarget: got %q", avail.servedTarget)
	}

	// Provider-first candidate uses provider ID as adapter key.
	noAdp := byID["prov-no-adapter"]
	if noAdp == nil {
		t.Fatal("prov-no-adapter not in candidates (provider-first dispatch must use provider ID as adapter key)")
	}
	if noAdp.adapter != "prov-no-adapter" {
		t.Errorf("prov-no-adapter adapter: got %q, want provider ID %q", noAdp.adapter, "prov-no-adapter")
	}
	if noAdp.servedTarget != "served-qwen" {
		t.Errorf("prov-no-adapter servedTarget: got %q", noAdp.servedTarget)
	}

	// Policy should use defaults since none of the providers set policy.
	if policy.maxQueue <= 0 {
		t.Errorf("policy.maxQueue: got %d, expected > 0", policy.maxQueue)
	}
	if policy.queueTimeout <= 0 {
		t.Errorf("policy.queueTimeout: got %v, expected > 0", policy.queueTimeout)
	}
}

// TestSubmitRunProviderPoolRewritesAdapterAndTarget verifies that provider-pool
// SubmitRun rewrites both req.Adapter and req.Target from the selected candidate.
// This test uses a fake TCP client via net.Pipe to capture the actual RunRequest
// sent by SubmitRun(ProviderPool=true), confirming that the winning candidate's
// adapter and servedTarget propagate correctly through the full service path.
func TestSubmitRunProviderPoolRewritesAdapterAndTarget(t *testing.T) {
	// Use net.Pipe to create a fake node connection that captures the RunRequest.
	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	// Capture the RunRequest received by the fake node.
	var capturedReq *iop.RunRequest
	var capturedMu sync.Mutex
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(req *iop.RunRequest) {
		capturedMu.Lock()
		capturedReq = req
		capturedMu.Unlock()
	})

	// Build the model catalog with provider references.
	catalog := []config.ModelCatalogEntry{
		{
			ID: "qwen3.6:35b",
			Providers: map[string]string{
				"prov-vllm-01": "served-qwen",
			},
		},
	}

	// Build NodeStore with a provider-pool provider.
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-pool",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-vllm-01",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-qwen"},
				Health:   "available",
				Capacity: 2,
			},
		},
	})

	// Build registry with the fake node.
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:         "node-pool",
		LifecycleState: edgenode.LifecycleConnected,
		Client:         edgeClient,
	})

	// Create Service with queue and catalog.
	// events bus must be non-nil to activate the queue path for provider-pool.
	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	// SubmitRun with ProviderPool=true.
	result, err := svc.SubmitRun(context.Background(), SubmitRunRequest{
		RunID:         "run-pool-test-001",
		ModelGroupKey: "qwen3.6:35b",
		ProviderPool:  true,
		Background:    true,
	})
	if err != nil {
		t.Fatalf("SubmitRun: %v", err)
	}
	if result == nil {
		t.Fatal("expected non-nil RunResult")
	}

	// Wait for the fake node to receive the request.
	time.Sleep(50 * time.Millisecond)

	capturedMu.Lock()
	defer capturedMu.Unlock()

	if capturedReq == nil {
		t.Fatal("no RunRequest captured from fake node; SubmitRun did not send")
	}

	// Verify that the adapter and target were rewritten from the provider-pool candidate.
	if capturedReq.GetAdapter() != "vllm-gpu" {
		t.Errorf("adapter: got %q, want %q", capturedReq.GetAdapter(), "vllm-gpu")
	}
	if capturedReq.GetTarget() != "served-qwen" {
		t.Errorf("target: got %q, want %q", capturedReq.GetTarget(), "served-qwen")
	}
	if capturedReq.GetRunId() != "run-pool-test-001" {
		t.Errorf("runID: got %q, want %q", capturedReq.GetRunId(), "run-pool-test-001")
	}
}

// TestResolveProviderPoolCandidatesAdapterInstanceValidation verifies that the
// defensive isProviderAdapterInstanceValid check in resolveProviderPoolCandidates
// excludes providers whose adapter name resolves to a disabled exact instance,
// an ambiguous type route, a zero-instance type route, or an unknown adapter key,
// while valid named instances and single enabled type routes remain as candidates.
func TestResolveProviderPoolCandidatesAdapterInstanceValidation(t *testing.T) {
	catalog := []config.ModelCatalogEntry{
		{
			ID: "ai-group",
			Providers: map[string]string{
				"prov-disabled":  "ai-model",
				"prov-ambiguous": "ai-model",
				"prov-valid":     "ai-model",
				"prov-zero-type": "ai-model",
				"prov-missing":   "ai-model",
			},
		},
	}

	store := edgenode.NewNodeStore()

	// prov-disabled: adapter="vllm-gpu" but VllmInstance "vllm-gpu" is disabled.
	store.Add(&edgenode.NodeRecord{
		ID: "node-disabled",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: false, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-disabled", Adapter: "vllm-gpu", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
		},
	})

	// prov-ambiguous: adapter="openai_compat" type route with 2 enabled instances → ambiguous.
	store.Add(&edgenode.NodeRecord{
		ID: "node-ambiguous",
		Adapters: config.AdaptersConf{
			OpenAICompatInstances: []config.OpenAICompatInstanceConf{
				{Name: "compat-a", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				{Name: "compat-b", Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-ambiguous", Adapter: "openai_compat", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
		},
	})

	// prov-valid: adapter="vllm-gpu" and VllmInstance "vllm-gpu" is enabled.
	store.Add(&edgenode.NodeRecord{
		ID: "node-valid",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-valid", Adapter: "vllm-gpu", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
		},
	})

	// prov-zero-type: adapter="ollama" type route but no enabled ollama instances → zero-instance type route excluded.
	store.Add(&edgenode.NodeRecord{
		ID: "node-zero-type",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-disabled", Enabled: false, BaseURL: "http://127.0.0.1:11434"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-zero-type", Adapter: "ollama", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
		},
	})

	// prov-missing: adapter="unknown-missing" not in any instance list and not a known type key → unknown adapter excluded.
	store.Add(&edgenode.NodeRecord{
		ID: "node-missing",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:9000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-missing", Adapter: "unknown-missing", Models: []string{"ai-model"}, Health: "available", Capacity: 2},
		},
	})

	reg := edgenode.NewRegistry()
	allRecs := store.All()
	for _, rec := range allRecs {
		reg.Register(&edgenode.NodeEntry{NodeID: rec.ID, LifecycleState: edgenode.LifecycleConnected})
	}

	svc := New(reg, nil)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	req := SubmitRunRequest{ModelGroupKey: "ai-group", ProviderPool: true}
	storeSnap, catalogSnap := svc.runtimeConfigSnapshot()
	candidates, _, err := svc.resolveProviderPoolCandidates(req, storeSnap, catalogSnap)
	if err != nil {
		t.Fatalf("resolveProviderPoolCandidates: %v", err)
	}

	// Only prov-valid should be a candidate.
	got := make(map[string]bool)
	for _, c := range candidates {
		got[c.providerID] = true
	}
	if len(candidates) != 1 {
		t.Fatalf("expected 1 candidate (prov-valid), got %d: %v", len(candidates), got)
	}
	if !got["prov-valid"] {
		t.Error("prov-valid (enabled exact instance) must be a candidate")
	}
	if got["prov-disabled"] {
		t.Error("prov-disabled (disabled exact instance) must be excluded")
	}
	if got["prov-ambiguous"] {
		t.Error("prov-ambiguous (ambiguous type route) must be excluded")
	}
	if got["prov-zero-type"] {
		t.Error("prov-zero-type (zero-instance type route) must be excluded")
	}
	if got["prov-missing"] {
		t.Error("prov-missing (unknown adapter) must be excluded")
	}
}

// TestGetSnapshotForNodeCatalogFirstNoDuplicates verifies that when a node has
// both adapter instances and a providers[] catalog, getSnapshotForNode returns
// only catalog snapshots (no adapter duplicates).
func TestGetSnapshotForNodeCatalogFirstNoDuplicates(t *testing.T) {
	rec := &edgenode.NodeRecord{
		ID: "node-catalog",
		Adapters: config.AdaptersConf{
			CLI: config.CLIConf{Enabled: true},
			OpenAICompatInstances: []config.OpenAICompatInstanceConf{
				{Name: "vllm-gpu", Enabled: true},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-a", Adapter: "vllm-gpu", Models: []string{"model-a"}, Health: "available", Capacity: 2},
			{ID: "prov-b", Adapter: "vllm-gpu", Models: []string{"model-b"}, Health: "available", Capacity: 3},
		},
		Runtime: config.RuntimeConf{Concurrency: 1},
	}
	store := edgenode.NewNodeStore()
	store.Add(rec)

	m := newModelQueueManager(store)

	snaps := m.getSnapshotForNode("node-catalog", rec)

	// Must return exactly the catalog providers — no CLI or openai_compat adapter duplicates.
	if len(snaps) != 2 {
		t.Fatalf("expected 2 catalog snapshots, got %d: %+v", len(snaps), snaps)
	}
	ids := map[string]bool{}
	for _, s := range snaps {
		ids[s.Id] = true
	}
	if !ids["prov-a"] || !ids["prov-b"] {
		t.Errorf("expected catalog provider IDs prov-a and prov-b, got: %v", ids)
	}
	// No adapter-only snapshots (no empty Id entries that correspond to adapter sections).
	for _, s := range snaps {
		if s.Id == "" {
			t.Errorf("unexpected adapter-only snapshot (empty Id) in catalog-first result: %+v", s)
		}
	}
}

// TestResolveProviderPoolCandidatesSkipsDisabledProvider verifies that a
// provider with enabled=false is excluded from dispatch candidates.
func TestResolveProviderPoolCandidatesSkipsDisabledProvider(t *testing.T) {
	disabled := false
	catalog := []config.ModelCatalogEntry{
		{
			ID: "model-x",
			Providers: map[string]string{
				"prov-active":   "served-x",
				"prov-disabled": "served-x",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-d",
		Providers: []config.NodeProviderConf{
			{ID: "prov-active", Type: "vllm", Adapter: "vllm-gpu", Models: []string{"served-x"}, Health: "available", Capacity: 2},
			{ID: "prov-disabled", Type: "vllm", Adapter: "vllm-gpu", Models: []string{"served-x"}, Health: "available", Capacity: 2, Enabled: &disabled},
		},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"}},
		},
	})
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-d", LifecycleState: edgenode.LifecycleConnected})

	svc := New(reg, nil)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	storeSnap, catalogSnap := svc.runtimeConfigSnapshot()
	candidates, _, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{ModelGroupKey: "model-x", ProviderPool: true}, storeSnap, catalogSnap)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(candidates) != 1 {
		t.Fatalf("expected 1 candidate (disabled filtered), got %d", len(candidates))
	}
	if candidates[0].providerID != "prov-active" {
		t.Errorf("expected prov-active, got %q", candidates[0].providerID)
	}
}

// TestResolveProviderPoolCandidatesUsesProviderIDWhenAdapterOmitted verifies
// that a provider-first provider (empty Adapter field) uses its ID as the
// dispatch adapter key.
func TestResolveProviderPoolCandidatesUsesProviderIDWhenAdapterOmitted(t *testing.T) {
	catalog := []config.ModelCatalogEntry{
		{
			ID: "model-y",
			Providers: map[string]string{
				"prov-first": "served-y",
			},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-pf",
		Providers: []config.NodeProviderConf{
			{ID: "prov-first", Type: "vllm", Adapter: "", Models: []string{"served-y"}, Health: "available", Capacity: 3},
		},
	})
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-pf", LifecycleState: edgenode.LifecycleConnected})

	svc := New(reg, nil)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	storeSnap, catalogSnap := svc.runtimeConfigSnapshot()
	candidates, _, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{ModelGroupKey: "model-y", ProviderPool: true}, storeSnap, catalogSnap)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(candidates) != 1 {
		t.Fatalf("expected 1 candidate, got %d", len(candidates))
	}
	if candidates[0].adapter != "prov-first" {
		t.Errorf("provider-first adapter key: got %q, want provider ID %q", candidates[0].adapter, "prov-first")
	}
	if candidates[0].providerID != "prov-first" {
		t.Errorf("providerID: got %q", candidates[0].providerID)
	}
}

// TestProviderSnapshotsReportDisabledProvider verifies that a disabled
// provider appears in snapshots with status=disabled and effective capacity 0.
func TestProviderSnapshotsReportDisabledProvider(t *testing.T) {
	disabled := false
	rec := &edgenode.NodeRecord{
		ID: "node-snap",
		Providers: []config.NodeProviderConf{
			{ID: "prov-on", Type: "vllm", Health: "available", Capacity: 4},
			{ID: "prov-off", Type: "vllm", Health: "available", Capacity: 4, Enabled: &disabled},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(rec)
	m := newModelQueueManager(store)

	snaps := m.getSnapshotForNode("node-snap", rec)
	if len(snaps) != 2 {
		t.Fatalf("expected 2 snapshots, got %d", len(snaps))
	}

	byID := map[string]*iop.ProviderSnapshot{}
	for _, s := range snaps {
		byID[s.Id] = s
	}

	on := byID["prov-on"]
	if on == nil {
		t.Fatal("prov-on snapshot missing")
	}
	if on.Status != "available" {
		t.Errorf("prov-on status: got %q, want available", on.Status)
	}
	if on.Capacity != 4 {
		t.Errorf("prov-on capacity: got %d, want 4", on.Capacity)
	}

	off := byID["prov-off"]
	if off == nil {
		t.Fatal("prov-off snapshot missing")
	}
	if off.Status != "disabled" {
		t.Errorf("prov-off status: got %q, want disabled", off.Status)
	}
	if off.Health != "disabled" {
		t.Errorf("prov-off health: got %q, want disabled", off.Health)
	}
	if off.Capacity != 0 {
		t.Errorf("prov-off effective capacity: got %d, want 0", off.Capacity)
	}
}

// TestStatusProviderProviderFirstNoAdapterDuplicates verifies that provider-first
// providers (empty Adapter) appear once in snapshots with no adapter duplicate.
func TestStatusProviderProviderFirstNoAdapterDuplicates(t *testing.T) {
	rec := &edgenode.NodeRecord{
		ID: "node-pf-snap",
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{{Name: "vllm-gpu", Enabled: true}},
		},
		Providers: []config.NodeProviderConf{
			{ID: "prov-first-a", Type: "vllm", Adapter: "", Models: []string{"model-a"}, Health: "available", Capacity: 2},
			{ID: "prov-first-b", Type: "vllm", Adapter: "", Models: []string{"model-b"}, Health: "available", Capacity: 2},
		},
	}
	store := edgenode.NewNodeStore()
	store.Add(rec)
	m := newModelQueueManager(store)

	snaps := m.getSnapshotForNode("node-pf-snap", rec)
	// catalog-first: exactly 2 provider snapshots, no adapter duplicate.
	if len(snaps) != 2 {
		t.Fatalf("expected 2 provider-first snapshots, got %d: %+v", len(snaps), snaps)
	}
	for _, s := range snaps {
		if s.Id == "" {
			t.Errorf("unexpected adapter-only snapshot (empty Id): %+v", s)
		}
		if s.Adapter != "" {
			t.Errorf("provider-first snapshot should have empty Adapter, got %q for id=%q", s.Adapter, s.Id)
		}
	}
}

// TestResolveProviderPoolCandidatesPropagatesPriority verifies that the actual
// resolveProviderPoolCandidates path propagates config provider Priority into
// candidateNode.priority.  This is the missing evidence from G06: the previous
// tests used hand-written candidateNode structs and never exercised the
// config -> resolver -> candidate.priority chain.
func TestResolveProviderPoolCandidatesPropagatesPriority(t *testing.T) {
	catalog := []config.ModelCatalogEntry{
		{
			ID: "model-prio",
			Providers: map[string]string{
				"prov-high-prio":   "served-model-prio",
				"prov-low-prio":    "served-model-prio",
				"prov-no-priority": "served-model-prio",
			},
		},
	}

	store := edgenode.NewNodeStore()
	// High-priority provider (priority=3).
	store.Add(&edgenode.NodeRecord{
		ID:      "node-prio",
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-high-prio",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-model-prio"},
				Health:   "available",
				Capacity: 2,
				Priority: 3,
			},
			{
				ID:       "prov-low-prio",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-model-prio"},
				Health:   "available",
				Capacity: 2,
				Priority: 10,
			},
			{
				ID:       "prov-no-priority",
				Adapter:  "vllm-gpu",
				Models:   []string{"served-model-prio"},
				Health:   "available",
				Capacity: 2,
				// Priority omitted → defaults to 0.
			},
		},
	})

	reg := edgenode.NewRegistry()
	for _, rec := range store.All() {
		reg.Register(&edgenode.NodeEntry{NodeID: rec.ID, LifecycleState: edgenode.LifecycleConnected})
	}

	svc := New(reg, nil)
	svc.SetNodeStore(store)
	svc.SetModelCatalog(catalog)

	storeSnap, catalogSnap := svc.runtimeConfigSnapshot()
	candidates, _, err := svc.resolveProviderPoolCandidates(
		SubmitRunRequest{ModelGroupKey: "model-prio", ProviderPool: true},
		storeSnap, catalogSnap,
	)
	if err != nil {
		t.Fatalf("resolveProviderPoolCandidates: %v", err)
	}
	if len(candidates) != 3 {
		t.Fatalf("expected 3 candidates, got %d", len(candidates))
	}

	byID := map[string]*candidateNode{}
	for i := range candidates {
		byID[candidates[i].providerID] = &candidates[i]
	}

	// prov-high-prio: fixture Priority=3 → candidate.priority must be 3.
	hp := byID["prov-high-prio"]
	if hp == nil {
		t.Fatal("prov-high-prio not in candidates")
	}
	if hp.priority != 3 {
		t.Errorf("prov-high-prio: got priority %d, want 3", hp.priority)
	}

	// prov-low-prio: fixture Priority=10 → candidate.priority must be 10.
	lp := byID["prov-low-prio"]
	if lp == nil {
		t.Fatal("prov-low-prio not in candidates")
	}
	if lp.priority != 10 {
		t.Errorf("prov-low-prio: got priority %d, want 10", lp.priority)
	}

	// prov-no-priority: Priority omitted → defaults to 0.
	np := byID["prov-no-priority"]
	if np == nil {
		t.Fatal("prov-no-priority not in candidates")
	}
	if np.priority != 0 {
		t.Errorf("prov-no-priority: got priority %d, want 0", np.priority)
	}
}

// TestGetSnapshotForNodeLegacyAdapterFallback verifies that a node with no
// providers[] catalog returns adapter snapshots (CLI, OllamaInstances, etc.)
// via the legacy adapter fallback path.
func TestGetSnapshotForNodeLegacyAdapterFallback(t *testing.T) {
	rec := &edgenode.NodeRecord{
		ID: "node-legacy",
		Adapters: config.AdaptersConf{
			CLI: config.CLIConf{Enabled: true},
			OllamaInstances: []config.OllamaInstanceConf{
				{Name: "ollama-local", Enabled: true, Capacity: 3},
			},
		},
		Runtime: config.RuntimeConf{Concurrency: 2},
	}
	store := edgenode.NewNodeStore()
	store.Add(rec)

	m := newModelQueueManager(store)

	snaps := m.getSnapshotForNode("node-legacy", rec)

	// Must return adapter snapshots: CLI + ollama-local = 2 entries.
	if len(snaps) != 2 {
		t.Fatalf("expected 2 adapter snapshots (cli + ollama-local), got %d: %+v", len(snaps), snaps)
	}
	adapters := map[string]bool{}
	for _, s := range snaps {
		adapters[s.Adapter] = true
	}
	if !adapters["cli"] {
		t.Error("expected cli adapter snapshot in legacy fallback")
	}
	if !adapters["ollama-local"] {
		t.Error("expected ollama-local adapter snapshot in legacy fallback")
	}
}

// longSlotCount reads the group's long in-flight counter for a provider slot.
func longSlotCount(m *modelQueueManager, groupKey, nodeID, providerID string) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	g, ok := m.groups[groupKey]
	if !ok {
		return -1
	}
	return g.longInflight[nodeID+":"+providerID]
}

func inflightSlotCount(m *modelQueueManager, groupKey, nodeID, providerID string) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	g, ok := m.groups[groupKey]
	if !ok {
		return -1
	}
	return g.inflight[nodeID+":"+providerID]
}

// TestModelQueueLongDispatchReservesBothCounters verifies S04: a long request
// occupies both the normal in-flight counter and the long in-flight counter, and
// every terminal reason (complete/error/cancelled) recovers both.
func TestModelQueueLongDispatchReservesBothCounters(t *testing.T) {
	for _, reason := range []string{"complete", "error", "cancelled"} {
		t.Run(reason, func(t *testing.T) {
			entry := &edgenode.NodeEntry{NodeID: "node-long"}
			cands := []candidateNode{{entry: entry, capacity: 3, providerID: "prov-long", longContextCapacity: 2}}
			m := newModelQueueManager(nil)

			sel, err := m.admit(context.Background(), "g-long", "", "", cands, groupPolicy{}, true)
			if err != nil || sel == nil {
				t.Fatalf("long admit: %v", err)
			}
			if got := inflightSlotCount(m, "g-long", "node-long", "prov-long"); got != 1 {
				t.Errorf("normal in-flight after admit = %d, want 1", got)
			}
			if got := longSlotCount(m, "g-long", "node-long", "prov-long"); got != 1 {
				t.Errorf("long in-flight after admit = %d, want 1", got)
			}

			longReserved := sel.longContextCapacity > 0
			m.trackInflight("g-long", "run-long-1", sel.entry.NodeID, sel.providerID, longReserved)
			if got := m.longInFlightForProvider("node-long", "prov-long"); got != 1 {
				t.Errorf("longInFlightForProvider after track = %d, want 1", got)
			}

			m.releaseRun("run-long-1", reason)
			if got := inflightSlotCount(m, "g-long", "node-long", "prov-long"); got != 0 {
				t.Errorf("normal in-flight after %s = %d, want 0", reason, got)
			}
			if got := longSlotCount(m, "g-long", "node-long", "prov-long"); got != 0 {
				t.Errorf("long in-flight after %s = %d, want 0", reason, got)
			}
			if got := m.longInFlightForProvider("node-long", "prov-long"); got != 0 {
				t.Errorf("longInFlightForProvider after %s = %d, want 0", reason, got)
			}
		})
	}
}

// TestModelQueueNormalRequestDoesNotReserveLongSlot verifies a normal request
// occupies only the normal in-flight counter, never the long counter, even on a
// provider that declares long capacity.
func TestModelQueueNormalRequestDoesNotReserveLongSlot(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-n"}
	cands := []candidateNode{{entry: entry, capacity: 3, providerID: "prov-n", longContextCapacity: 2}}
	m := newModelQueueManager(nil)

	sel, err := m.admit(context.Background(), "g-n", "", "", cands, groupPolicy{}) // long omitted → normal
	if err != nil || sel == nil {
		t.Fatalf("normal admit: %v", err)
	}
	if got := inflightSlotCount(m, "g-n", "node-n", "prov-n"); got != 1 {
		t.Errorf("normal in-flight = %d, want 1", got)
	}
	if got := longSlotCount(m, "g-n", "node-n", "prov-n"); got != 0 {
		t.Errorf("long in-flight for normal request = %d, want 0", got)
	}
}

// TestModelQueueLongSlotFullExcludesLongKeepsNormal verifies S05 and D06: a
// provider whose long slots are full is excluded from long candidates (the long
// request moves to another provider or queues), but the same provider remains a
// valid candidate for normal requests while its normal capacity has room.
func TestModelQueueLongSlotFullExcludesLongKeepsNormal(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-a"}
	entryB := &edgenode.NodeEntry{NodeID: "node-b"}
	cands := []candidateNode{
		{entry: entryA, capacity: 3, priority: 0, providerID: "prov-a", longContextCapacity: 1},
		{entry: entryB, capacity: 3, priority: 1, providerID: "prov-b", longContextCapacity: 1},
	}
	m := newModelQueueManager(nil)

	// Long #1 → prov-a (priority 0 preferred), filling prov-a's single long slot.
	sel1, err := m.admit(context.Background(), "g-x", "", "", cands, groupPolicy{}, true)
	if err != nil || sel1 == nil {
		t.Fatalf("long1 admit: %v", err)
	}
	if sel1.providerID != "prov-a" {
		t.Fatalf("long1 provider = %q, want prov-a", sel1.providerID)
	}

	// Long #2 must skip the long-full prov-a and land on prov-b.
	sel2, err := m.admit(context.Background(), "g-x", "", "", cands, groupPolicy{}, true)
	if err != nil || sel2 == nil {
		t.Fatalf("long2 admit: %v", err)
	}
	if sel2.providerID != "prov-b" {
		t.Fatalf("long2 provider = %q, want prov-b (prov-a long slot full)", sel2.providerID)
	}

	// A normal request restricted to the long-full prov-a is still admitted,
	// because normal admission ignores long slot state (D06).
	normalCands := []candidateNode{cands[0]}
	selN, err := m.admit(context.Background(), "g-x", "", "", normalCands, groupPolicy{})
	if err != nil || selN == nil {
		t.Fatalf("normal admit on long-full provider: %v", err)
	}
	if selN.providerID != "prov-a" {
		t.Fatalf("normal provider = %q, want prov-a", selN.providerID)
	}

	// A further long request restricted to the long-full prov-a cannot be
	// admitted immediately; it queues and times out.
	longOnlyA := []candidateNode{cands[0]}
	_, err = m.admit(context.Background(), "g-x", "", "", longOnlyA, groupPolicy{maxQueue: 1, queueTimeout: 20 * time.Millisecond}, true)
	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("long request on long-full provider: expected queue timeout, got %v", err)
	}
}