package service

import (
	"testing"
	"time"
)

// 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)
}

type concurrentAdmitResult struct {
	candidate *candidateNode
	err       error
}

// trackTestLease registers a tracked lease for a slot the test reserved by hand,
// standing in for the admission → trackLease sequence a real dispatch performs.
// Use it only when the test manipulates group counters directly instead of going
// through admit.
func trackTestLease(m *modelQueueManager, groupKey, runID, nodeID, providerID string, long bool) uint64 {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.leaseSeq++
	id := m.leaseSeq
	m.leases[id] = &providerLease{
		id:         id,
		groupKey:   groupKey,
		nodeID:     nodeID,
		providerID: providerID,
		long:       long,
		state:      leaseStateTracked,
		runID:      runID,
	}
	m.leaseByRun[runID] = id
	return id
}

// leaseLong reports whether the lease reserved a long-context slot.
func leaseLong(t *testing.T, m *modelQueueManager, leaseID uint64) bool {
	t.Helper()
	m.mu.Lock()
	defer m.mu.Unlock()
	lease, ok := m.leases[leaseID]
	if !ok {
		t.Fatalf("lease %d not found", leaseID)
	}
	return lease.long
}

// leaseCount reports how many leases the manager still holds.
func leaseCount(m *modelQueueManager) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.leases)
}

// providerResourceCounts reports the live normal and long-context in-flight
// counts a provider resource holds. Race fixtures use it to sample the peak
// occupancy the scheduler ever produced and to assert the final drained state.
func providerResourceCounts(m *modelQueueManager, nodeID, providerID string) (inFlight, longInFlight int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	res, ok := m.resources[providerResourceKey{nodeID: nodeID, providerID: providerID}]
	if !ok {
		return 0, 0
	}
	return res.inFlight, res.longInFlight
}

// pendingItemCount reports how many items are queued across every model group,
// which is the scheduler-wide view the global pump operates on.
func pendingItemCount(m *modelQueueManager) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.pendingItemsLocked())
}
