package service

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"

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

// 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, nil)
	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{}, nil)
	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{}, nil)
	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{}

	m := newModelQueueManager(store)

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

	// Queue a second item in a goroutine.
	resultCh := make(chan admitResult, 1)
	go func() {
		n, e := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy, nil)
		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.mu.Unlock()
			trackTestLease(m, "g-tr-"+termType, runID, "node-tr1", "", false)

			// The transport lifecycle hook calls this directly; the event bus is
			// not in the release path.
			m.releaseRun(runID, 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.pumpAllLocked()
	}
	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},
	}

	m := newModelQueueManager(store)

	// 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
	// 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()
	trackTestLease(m, "g-nd", "run-nd-x", "node-nd1", "", false)
	trackTestLease(m, "g-nd", "run-nd-y", "node-nd2", "", false)

	// node-nd1 disconnects: its inflight is freed and it is removed from candidates.
	// node-nd2 is still full, so no dispatch yet.
	m.releaseNode("node-nd1", "disconnected")

	// node-nd2's run terminates: now nd2 has capacity, dispatch goes to nd2.
	m.releaseRun("run-nd-y", "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 lease for run-nd-x should be gone (released at disconnect).
	m.mu.Lock()
	_, stillInFlight := m.leaseByRun["run-nd-x"]
	m.mu.Unlock()
	if stillInFlight {
		t.Error("run-nd-x lease should be released after node disconnect")
	}
}

// TestRunTerminalReleasesLeaseWhenEventSubscriberFull pins that lease accounting
// survives a lossy observability fanout. The bus drops into full subscriber
// channels by design, so the transport calls the service lifecycle hook directly;
// this test saturates the bus, confirms it is actually dropping, and then checks
// the terminal event still returns the slot.
func TestRunTerminalReleasesLeaseWhenEventSubscriberFull(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-drop",
		Providers: []config.NodeProviderConf{
			{ID: "prov-drop", Capacity: 1, Adapter: "vllm"},
		},
	})

	bus := edgeevents.NewBus()
	svc := New(edgenode.NewRegistry(), bus)
	svc.SetNodeStore(store)

	// A slow observer with a one-deep buffer: the fanout starts dropping as soon
	// as it stops reading, which is exactly the delivery correctness must not
	// depend on.
	_, unsub := bus.SubscribeAllRuns(1)
	defer unsub()

	entry := &edgenode.NodeEntry{NodeID: "node-drop"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-drop"}}
	admitted, _, err := svc.queue.admitWithReason(context.Background(), "g-drop", "", "", cands, groupPolicy{}, nil)
	if err != nil || admitted == nil {
		t.Fatalf("admit: %v", err)
	}
	svc.queue.trackLease(admitted.leaseID, "run-drop-1")

	for i := 0; i < 8; i++ {
		bus.PublishRun(&iop.RunEvent{RunId: "run-drop-1", Type: "chunk"})
	}
	if dropped := bus.Stats().DroppedRunEvents; dropped == 0 {
		t.Fatal("expected the saturated subscriber to drop run events; the test is not exercising the drop path")
	}

	// The terminal event is delivered the way the transport delivers it.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: "run-drop-1", Type: "complete"})

	svc.queue.mu.Lock()
	inflight, _ := svc.queue.getStatsForProviderLocked("node-drop", "prov-drop")
	svc.queue.mu.Unlock()
	if inflight != 0 {
		t.Errorf("in-flight after terminal event with a dropping bus = %d, want 0", inflight)
	}
	if got := leaseCount(svc.queue); got != 0 {
		t.Errorf("live leases after terminal event = %d, want 0", got)
	}
}

// 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) {
	entry := &edgenode.NodeEntry{NodeID: "node-x"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-1"}}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-x",
		Providers: []config.NodeProviderConf{
			{ID: "prov-1", Capacity: 1, Adapter: "vllm"},
		},
	})
	m := newModelQueueManager(store)
	m.setStore(store)

	selected, err := m.admit(context.Background(), "g-cancel", "", "", cands, groupPolicy{}, nil)
	if err != nil || selected == nil {
		t.Fatalf("admit failed: %v", err)
	}

	m.trackLease(selected.leaseID, "run-cancel-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)
	}

	m.releaseRun("run-cancel-1", "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, nil)
	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, nil)
		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)
	}
}

func TestProviderResourceCapacitySharedAcrossModelGroups(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-shared"}
	cands := []candidateNode{
		{entry: entry, capacity: 1, providerID: "prov-shared"},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-shared",
		Providers: []config.NodeProviderConf{
			{ID: "prov-shared", Capacity: 1, Adapter: "vllm"},
		},
	})
	m := newModelQueueManager(store)
	m.setStore(store)

	start := make(chan struct{})
	resCh := make(chan concurrentAdmitResult, 2)

	for i := 0; i < 2; i++ {
		go func(id int) {
			<-start
			ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
			defer cancel()
			groupKey := "g-shared-" + string(rune('a'+id))
			sel, err := m.admit(ctx, groupKey, "", "", cands, groupPolicy{maxQueue: 1, queueTimeout: 30 * time.Millisecond}, nil)
			resCh <- concurrentAdmitResult{candidate: sel, err: err}
		}(i)
	}

	close(start)

	var results []concurrentAdmitResult
	for i := 0; i < 2; i++ {
		results = append(results, <-resCh)
	}

	successCount := 0
	failCount := 0
	for _, res := range results {
		if res.err == nil {
			successCount++
		} else {
			failCount++
		}
	}

	if successCount != 1 || failCount != 1 {
		t.Fatalf("expected exactly 1 success and 1 failure under concurrent admission, got success=%d fail=%d", successCount, failCount)
	}

	// Verify in-flight count peak is 1
	m.mu.Lock()
	res := m.resources[providerResourceKey{nodeID: "node-shared", providerID: "prov-shared"}]
	if res == nil {
		m.mu.Unlock()
		t.Fatal("expected provider resource state to exist")
	}
	if res.inFlight != 1 {
		t.Errorf("expected inFlight peak to be 1, got %d", res.inFlight)
	}
	m.mu.Unlock()

	// Release slot
	m.releaseSlot("g-shared-a", "node-shared", "prov-shared")

	// Verify count is 0
	m.mu.Lock()
	if res.inFlight != 0 {
		t.Errorf("expected inFlight to recover to 0, got %d", res.inFlight)
	}
	m.mu.Unlock()
}

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

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID:      "node-snap-read",
		Runtime: config.RuntimeConf{Concurrency: 5},
		Providers: []config.NodeProviderConf{
			{
				ID:                  "snap-prov-read",
				Adapter:             "vllm",
				Capacity:            3,
				LongContextCapacity: 2,
				Enabled:             nil,
			},
		},
	}
	store.Add(rec)

	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)

	cands := []candidateNode{
		{entry: reg.All()[0], capacity: 3, providerID: "snap-prov-read", longContextCapacity: 2},
	}

	sel, err := svc.queue.admit(context.Background(), "g-snap-read", "", "", cands, groupPolicy{}, nil, true)
	if err != nil || sel == nil {
		t.Fatalf("admit failed: %v", err)
	}

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot")
	}
	p := snaps[0].ProviderSnapshots[0]
	if p.InFlight != 1 {
		t.Errorf("expected InFlight 1, got %d", p.InFlight)
	}
	if p.LongInFlight != 1 {
		t.Errorf("expected LongInFlight 1, got %d", p.LongInFlight)
	}

	svc.queue.releaseSlotWithLong("g-snap-read", "node-snap-read", "snap-prov-read", true)

	snaps2 := svc.ListNodeSnapshots()
	p2 := snaps2[0].ProviderSnapshots[0]
	if p2.InFlight != 0 {
		t.Errorf("expected InFlight 0 after release, got %d", p2.InFlight)
	}
	if p2.LongInFlight != 0 {
		t.Errorf("expected LongInFlight 0 after release, got %d", p2.LongInFlight)
	}
}

func TestProviderResourceCapacityRefreshPreservesInflight(t *testing.T) {
	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-refresh",
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-refresh",
				Capacity: 2,
				Enabled:  nil, // defaults to true
			},
		},
	}
	store.Add(rec)

	m := newModelQueueManager(store)
	m.setStore(store)
	entry := &edgenode.NodeEntry{NodeID: "node-refresh"}
	cands := []candidateNode{
		{entry: entry, capacity: 2, providerID: "prov-refresh"},
	}

	// Admit 2 items (matching capacity)
	sel1, err := m.admit(context.Background(), "group-1", "", "", cands, groupPolicy{}, nil)
	if err != nil || sel1 == nil {
		t.Fatalf("admit 1 failed: %v", err)
	}
	sel2, err := m.admit(context.Background(), "group-2", "", "", cands, groupPolicy{}, nil)
	if err != nil || sel2 == nil {
		t.Fatalf("admit 2 failed: %v", err)
	}

	// Verify inFlight = 2
	m.mu.Lock()
	res := m.resources[providerResourceKey{nodeID: "node-refresh", providerID: "prov-refresh"}]
	if res.inFlight != 2 {
		t.Fatalf("expected inFlight = 2, got %d", res.inFlight)
	}
	m.mu.Unlock()

	// Update store with decreased capacity (2 -> 1)
	newStore := edgenode.NewNodeStore()
	newRec := &edgenode.NodeRecord{
		ID: "node-refresh",
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-refresh",
				Capacity: 1,
				Enabled:  nil,
			},
		},
	}
	newStore.Add(newRec)

	// Refresh store
	m.setStore(newStore)

	// Verify capacity is updated but inFlight is preserved
	m.mu.Lock()
	if res.capacity != 1 {
		t.Errorf("expected updated capacity to be 1, got %d", res.capacity)
	}
	if res.inFlight != 2 {
		t.Errorf("expected inFlight to remain 2, got %d", res.inFlight)
	}
	m.mu.Unlock()

	// Third admit should fail because capacity is full (capacity 1, inflight 2)
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	_, err = m.admit(ctx, "group-3", "", "", cands, groupPolicy{queueTimeout: 10 * time.Millisecond}, nil)
	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("expected queue timeout for third admit, got: %v", err)
	}
}

// TestProviderResourceRemovedProviderDoesNotReadmitQueuedCandidate verifies
// that a stale queued candidate cannot re-admit a removed provider. After the
// last release deletes the orphan resource, a queue pump must not recreate an
// active resource from the stale candidate's old capacity.
//
// Arrange order is critical: the second admit must be in-flight in the queue
// BEFORE the release that triggers pumpAllLocked. This exercises the real
// release-triggered queue-pump path where releaseSlotLocked → pumpAllLocked
// would otherwise hand a stale candidate to a deleted provider.
func TestProviderResourceRemovedProviderDoesNotReadmitQueuedCandidate(t *testing.T) {
	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-stale",
		Providers: []config.NodeProviderConf{
			{ID: "prov-stale", Capacity: 1, Adapter: "vllm"},
		},
	}
	store.Add(rec)

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

	m := newModelQueueManager(store)
	m.setStore(store)

	// 1. Admit the first request — fills the only capacity slot.
	first, err := m.admit(context.Background(), "g-stale", "", "", cands, groupPolicy{}, nil)
	if err != nil || first == nil {
		t.Fatalf("initial admit: %v", err)
	}

	key := providerResourceKey{nodeID: "node-stale", providerID: "prov-stale"}
	m.mu.Lock()
	res := m.resources[key]
	m.mu.Unlock()
	if res == nil || res.orphan {
		t.Fatal("expected active non-orphan resource after first admit")
	}

	// 2. Start the second admit in a goroutine so it blocks in the queue
	//    while the first request is still running. Use a short timeout so the
	//    test can observe the terminal result quickly.
	resultCh := make(chan admitResult, 1)
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
		defer cancel()
		c, _, e := m.admitWithReason(ctx, "g-stale", "", "", cands, groupPolicy{maxQueue: 16, queueTimeout: 500 * time.Millisecond}, nil)
		resultCh <- admitResult{candidate: c, err: e}
	}()

	// 3. Confirm the second request actually entered the queue.
	waitForQueueLen(t, m, "g-stale", 1)

	// 4. Remove the provider from the store so the resource becomes orphan.
	rec.Providers = nil
	m.setStore(store)

	m.mu.Lock()
	res, ok := m.resources[key]
	m.mu.Unlock()
	if !ok || res == nil || !res.orphan {
		t.Fatalf("expected resource to be orphan after store update, ok=%t, res=%v", ok, res)
	}

	// 5. Release the first slot to trigger the queue pump. The orphan
	//    resource must be deleted (inFlight=0) and the queue pump must
	//    NOT recreate it via pumpAllLocked → findAvailableNodeLocked.
	m.releaseSlotWithLong("g-stale", "node-stale", "prov-stale", false)

	// 6. The second request must NOT succeed — the stale candidate must
	//    not be dispatched and the resource must not be recreated.
	select {
	case res := <-resultCh:
		if res.err == nil {
			t.Fatal("expected error/timeout for queued candidate of removed provider, got nil error")
		}
		if res.candidate != nil {
			t.Fatalf("expected no candidate dispatch, got %+v", res.candidate)
		}
	case <-time.After(1 * time.Second):
		t.Fatal("timeout: second admit did not return within expected window")
	}

	// 7. Verify the resource was not recreated as an active entry.
	m.mu.Lock()
	_, exists := m.resources[key]
	m.mu.Unlock()
	if exists {
		t.Fatal("expected orphan resource to remain deleted after queue pump")
	}
}

// queueItemForTest builds a pending item with a generous deadline so the test
// controls when it resolves.
func queueItemForTest(cands []candidateNode, long bool) *queueItem {
	return &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(5 * time.Second),
		long:       long,
	}
}

// enqueueForTest registers an item in the given group through the same helper
// admission uses, so it receives a real global enqueue sequence. The optional
// resolveCandidates closure is forwarded to the manager; nil keeps the legacy
// snapshot path for existing fixtures.
func enqueueForTest(m *modelQueueManager, groupKey string, item *queueItem, resolveCandidates func() ([]candidateNode, error)) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.enqueueItemLocked(m.getOrCreateGroupLocked(groupKey, groupPolicy{}), item, resolveCandidates)
}

// TestGlobalPumpWakesDifferentModelGroup verifies the cross-group wake-up
// contract (S03 / cross-group-wakeup): provider resources are shared across
// model groups, so a lease released by one group must wake a waiter queued in
// another. Before the global pump, the waiter slept until its own group saw an
// event or its queue timeout expired.
func TestGlobalPumpWakesDifferentModelGroup(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-wake"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-wake"}}
	m := newModelQueueManager(nil)

	// Group A takes the provider's only slot.
	first, err := m.admit(context.Background(), "g-wake-a", "", "", cands, groupPolicy{}, nil)
	if err != nil || first == nil {
		t.Fatalf("group A admit: %v", err)
	}

	// Group B queues a request for the same provider.
	waiter := queueItemForTest(cands, false)
	enqueueForTest(m, "g-wake-b", waiter, nil)

	// Releasing group A's lease must dispatch group B's waiter.
	m.releaseLease(first.leaseID, "complete")

	select {
	case res := <-waiter.waitCh:
		if res.err != nil {
			t.Fatalf("group B waiter expected dispatch, got error: %v", res.err)
		}
		if res.candidate == nil || res.candidate.providerID != "prov-wake" {
			t.Fatalf("group B waiter: unexpected candidate %v", res.candidate)
		}
	default:
		t.Fatal("group B waiter was not woken by the group A release (cross-group wake-up missing)")
	}
}

// TestGlobalPumpPreservesEarliestDispatchableSequence verifies that the pump
// orders waiters by global arrival, not by group: an item blocked on a full
// provider is skipped, and among the items that can dispatch the earliest
// enqueued one wins even when it sits in a different group than the release.
func TestGlobalPumpPreservesEarliestDispatchableSequence(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-seq"}
	blockedCands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-seq-blocked"}}
	openCands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-seq-open"}}
	m := newModelQueueManager(nil)

	// Fill both providers so every queued item starts out blocked.
	if _, err := m.admit(context.Background(), "g-seq-a", "", "", blockedCands, groupPolicy{}, nil); err != nil {
		t.Fatalf("fill blocked provider: %v", err)
	}
	openHolder, err := m.admit(context.Background(), "g-seq-a", "", "", openCands, groupPolicy{}, nil)
	if err != nil {
		t.Fatalf("fill open provider: %v", err)
	}

	// Arrival order: the earliest item can never dispatch (its provider is never
	// released); the next two compete for the one slot that does free.
	blocked := queueItemForTest(blockedCands, false)
	enqueueForTest(m, "g-seq-a", blocked, nil)
	earlier := queueItemForTest(openCands, false)
	enqueueForTest(m, "g-seq-b", earlier, nil)
	later := queueItemForTest(openCands, false)
	enqueueForTest(m, "g-seq-a", later, nil)

	m.releaseLease(openHolder.leaseID, "complete")

	select {
	case res := <-earlier.waitCh:
		if res.err != nil {
			t.Fatalf("earlier item dispatch error: %v", res.err)
		}
	default:
		t.Fatal("earliest dispatchable item was not dispatched first")
	}
	select {
	case <-later.waitCh:
		t.Fatal("later item dispatched out of global enqueue order")
	default:
	}
	select {
	case <-blocked.waitCh:
		t.Fatal("item on a still-full provider must stay queued")
	default:
	}

	// Freeing the slot again promotes the later item; the blocked one stays put,
	// proving skipping does not consume its turn.
	m.mu.Lock()
	pending := len(m.pendingItemsLocked())
	m.mu.Unlock()
	if pending != 2 {
		t.Fatalf("expected 2 pending items after the first dispatch, got %d", pending)
	}
}

// TestGlobalPumpConcurrentReleaseCancelRefreshNoLeak drives releases, cancels,
// and config refreshes against the global pump concurrently. The invariants it
// pins are the ones a shared scheduler can silently break: no request is handed
// two leases, provider in-flight never exceeds configured capacity, and every
// counter returns to zero once the storm ends.
func TestGlobalPumpConcurrentReleaseCancelRefreshNoLeak(t *testing.T) {
	const (
		capacity   = 3
		groups     = 3
		perGroup   = 12
		nodeID     = "node-race"
		providerID = "prov-race"
	)

	newStore := func(cap int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: nodeID,
			Providers: []config.NodeProviderConf{
				{ID: providerID, Adapter: "vllm", Capacity: cap, LongContextCapacity: 1},
			},
		})
		return store
	}

	m := newModelQueueManager(nil)
	m.setStore(newStore(capacity))

	entry := &edgenode.NodeEntry{NodeID: nodeID}
	cands := []candidateNode{{entry: entry, capacity: capacity, providerID: providerID, longContextCapacity: 1}}

	var (
		mu         sync.Mutex
		seenLeases = map[uint64]int{}
		peak       int
		peakLong   int
	)
	observe := func() {
		inFlight, longInFlight := providerResourceCounts(m, nodeID, providerID)
		mu.Lock()
		if inFlight > peak {
			peak = inFlight
		}
		if longInFlight > peakLong {
			peakLong = longInFlight
		}
		mu.Unlock()
	}

	stop := make(chan struct{})
	var refreshWG sync.WaitGroup
	refreshWG.Add(1)
	go func() {
		defer refreshWG.Done()
		for {
			select {
			case <-stop:
				return
			default:
			}
			// Config refresh alternates capacity; both values are pumped, so the
			// scheduler re-evaluates every waiter on each swap.
			m.setStore(newStore(capacity))
			m.setStore(newStore(capacity + 1))
			observe()
		}
	}()

	var wg sync.WaitGroup
	for g := 0; g < groups; g++ {
		groupKey := fmt.Sprintf("g-race-%d", g)
		for i := 0; i < perGroup; i++ {
			wg.Add(1)
			long := i%3 == 0
			// A share of the callers abandon their request mid-flight, so cancel
			// races the dispatch hand-off.
			cancelEarly := i%4 == 0
			go func() {
				defer wg.Done()
				ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
				defer cancel()
				if cancelEarly {
					go func() {
						time.Sleep(time.Millisecond)
						cancel()
					}()
				}
				admitted, err := m.admit(ctx, groupKey, "", "", cands, groupPolicy{maxQueue: 64, queueTimeout: 150 * time.Millisecond}, nil, long)
				if err != nil || admitted == nil {
					return
				}
				mu.Lock()
				seenLeases[admitted.leaseID]++
				mu.Unlock()
				observe()
				time.Sleep(time.Millisecond)
				m.releaseLease(admitted.leaseID, "complete")
			}()
		}
	}

	wg.Wait()
	close(stop)
	refreshWG.Wait()

	// Every admitted request holds exactly one lease identity.
	mu.Lock()
	for id, count := range seenLeases {
		if count != 1 {
			mu.Unlock()
			t.Fatalf("lease %d handed out %d times; a request must own exactly one lease", id, count)
		}
	}
	observedPeak, observedPeakLong := peak, peakLong
	mu.Unlock()

	if observedPeak > capacity+1 {
		t.Errorf("provider in-flight peaked at %d, above the highest configured capacity %d", observedPeak, capacity+1)
	}
	if observedPeakLong > 1 {
		t.Errorf("long in-flight peaked at %d, above the configured long capacity 1", observedPeakLong)
	}

	// Final state must be fully drained: no held resource, lease, or pending item.
	inFlight, longInFlight := providerResourceCounts(m, nodeID, providerID)
	if inFlight != 0 || longInFlight != 0 {
		t.Errorf("expected zero in-flight after drain, got inflight=%d long=%d", inFlight, longInFlight)
	}
	if got := leaseCount(m); got != 0 {
		t.Errorf("expected zero live leases after drain, got %d", got)
	}
	if got := pendingItemCount(m); got != 0 {
		t.Errorf("expected zero pending items after drain, got %d", got)
	}
}