package service

import (
	"context"
	"sync"
	"testing"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
)

// reservationFixture admits one slot on a single-capacity provider and returns
// the manager, the reservation for that admission, and the slot key it holds.
func reservationFixture(t *testing.T, long bool, longCapacity int) (*modelQueueManager, *queueReservation, string) {
	t.Helper()
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-res",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			// The node record is the live source for capacity and long-context
			// capacity, so it must declare the same long limit the candidate does:
			// the scheduler re-resolves these fields from the record at dispatch.
			{ID: "prov-1", Capacity: 1, Adapter: "vllm", LongContextCapacity: longCapacity},
		},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-res"}
	selected := &candidateNode{
		entry:               entry,
		capacity:            1,
		providerID:          "prov-1",
		longContextCapacity: longCapacity,
	}
	cands := []candidateNode{*selected}

	m := newModelQueueManager(store)
	admitted, _, err := m.admitWithReason(context.Background(), "g-res", "", "", cands, groupPolicy{}, nil, long, false)
	if err != nil || admitted == nil {
		t.Fatalf("admit: %v", err)
	}

	return m, newQueueReservation(m, admitted), "node-res:prov-1"
}

func inflightCount(t *testing.T, m *modelQueueManager, groupKey, slot string) (int, int) {
	t.Helper()
	m.mu.Lock()
	defer m.mu.Unlock()

	var nodeID, providerID string
	if colonIdx := findLastColon(slot); colonIdx > 0 {
		nodeID = slot[:colonIdx]
		providerID = slot[colonIdx+1:]
	} else {
		nodeID = slot
	}

	if providerID != "" {
		key := providerResourceKey{nodeID: nodeID, providerID: providerID}
		if res, ok := m.resources[key]; ok {
			return res.inFlight, res.longInFlight
		}
		return 0, 0
	}

	group, ok := m.groups[groupKey]
	if !ok {
		t.Fatalf("group %q not found", groupKey)
	}
	return group.inflight[slot], group.longInflight[slot]
}

// TestQueueReservationReleasesExactlyOnce pins the release contract at each
// boundary a dispatch can fail on: the slot is freed once, and a repeated
// release cannot free a second slot.
func TestQueueReservationReleasesExactlyOnce(t *testing.T) {
	cases := []struct {
		name string
		// track reports whether the dispatch got as far as registering the run
		// inflight before failing.
		track  bool
		reason string
	}{
		{name: "prepare tunnel error before track", reason: "prepare-tunnel-error"},
		{name: "prepare run error before track", reason: "prepare-run-error"},
		{name: "build error before track", reason: "build-error"},
		{name: "unknown execution path before track", reason: "unknown-execution-path"},
		{name: "send error after track", track: true, reason: "send-error"},
		{name: "no event bus after track", track: true, reason: "no-event-bus"},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			m, reservation, slot := reservationFixture(t, false, 0)

			if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 1 {
				t.Fatalf("expected 1 reserved slot after admission, got %d", inflight)
			}

			if tc.track {
				reservation.track("run-res")
			}

			reservation.release(tc.reason)
			if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
				t.Fatalf("expected the slot to be released, got %d in flight", inflight)
			}

			// A second release must not decrement a slot that now belongs to
			// another request.
			reservation.release(tc.reason)
			if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
				t.Fatalf("expected release to be idempotent, got %d in flight", inflight)
			}
			if tc.track {
				m.mu.Lock()
				_, stillTracked := m.leaseByRun["run-res"]
				m.mu.Unlock()
				if stillTracked {
					t.Fatal("expected the run to be removed from the lease index")
				}
			}
			if got := leaseCount(m); got != 0 {
				t.Fatalf("expected the lease record to be dropped, got %d live leases", got)
			}
		})
	}
}

// TestQueueReservationHandOffKeepsSlotHeld pins that a dispatched request keeps
// its slot: after hand-off the run/tunnel lifecycle owns the release, so a late
// release on the reservation must not free it.
func TestQueueReservationHandOffKeepsSlotHeld(t *testing.T) {
	m, reservation, slot := reservationFixture(t, false, 0)
	reservation.track("run-res")
	reservation.handOff()

	reservation.release("send-error")

	if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 1 {
		t.Fatalf("expected the dispatched run to keep its slot, got %d in flight", inflight)
	}

	// The lifecycle owner releases it exactly once, via the run registration.
	m.releaseRun("run-res", "complete")
	if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
		t.Fatalf("expected the terminal event to release the slot, got %d in flight", inflight)
	}
}

// TestQueueReservationLongSlotMatchesReservation pins that the long-context
// slot is released only when it was actually reserved: a long request on a
// provider that declares a long limit.
func TestQueueReservationLongSlotMatchesReservation(t *testing.T) {
	t.Run("long request on a long-capacity provider releases the long slot", func(t *testing.T) {
		m, reservation, slot := reservationFixture(t, true, 1)
		if _, longInflight := inflightCount(t, m, "g-res", slot); longInflight != 1 {
			t.Fatalf("expected 1 reserved long slot, got %d", longInflight)
		}

		reservation.release("send-error")

		inflight, longInflight := inflightCount(t, m, "g-res", slot)
		if inflight != 0 || longInflight != 0 {
			t.Fatalf("expected both slots released, got inflight=%d long=%d", inflight, longInflight)
		}
	})

	t.Run("long request without a declared long limit reserves no long slot", func(t *testing.T) {
		m, reservation, slot := reservationFixture(t, true, 0)
		if _, longInflight := inflightCount(t, m, "g-res", slot); longInflight != 0 {
			t.Fatalf("expected no long slot to be reserved, got %d", longInflight)
		}
		if leaseLong(t, m, reservation.leaseID) {
			t.Fatal("expected the lease not to claim a long slot")
		}

		reservation.release("send-error")

		if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
			t.Fatalf("expected the slot to be released, got %d in flight", inflight)
		}
	})
}

// leaseRaceFixture admits two leases on one capacity-2 provider and returns the
// manager, both lease ids, and the shared slot key. Holding a second lease is
// what makes an over-release observable: the resource counters clamp at zero, so
// a lone lease could absorb a double decrement silently, while here the extra
// decrement would steal the surviving lease's slot.
func leaseRaceFixture(t *testing.T, long bool, longCapacity int) (m *modelQueueManager, first, second uint64, slot string) {
	t.Helper()
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-race",
		Runtime: config.RuntimeConf{Concurrency: 2},
		Providers: []config.NodeProviderConf{
			{ID: "prov-race", Capacity: 2, LongContextCapacity: longCapacity, Adapter: "vllm"},
		},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-race"}
	cands := []candidateNode{{
		entry:               entry,
		capacity:            2,
		providerID:          "prov-race",
		longContextCapacity: longCapacity,
	}}

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

	ids := make([]uint64, 0, 2)
	for i := 0; i < 2; i++ {
		admitted, _, err := m.admitWithReason(context.Background(), "g-race", "", "", cands, groupPolicy{}, nil, long, false)
		if err != nil || admitted == nil {
			t.Fatalf("admit %d: %v", i, err)
		}
		if admitted.leaseID == 0 {
			t.Fatalf("admit %d returned no lease id", i)
		}
		ids = append(ids, admitted.leaseID)
	}
	if ids[0] == ids[1] {
		t.Fatalf("expected unique lease ids per admission, both were %d", ids[0])
	}
	return m, ids[0], ids[1], "node-race:prov-race"
}

// TestProviderLeaseConcurrentReleaseExactlyOnce pins the core lease invariant:
// however many terminal causes race on one lease — send failure, run terminal,
// tunnel terminal — the slot is returned exactly once and the other live lease
// keeps what it holds.
func TestProviderLeaseConcurrentReleaseExactlyOnce(t *testing.T) {
	cases := []struct {
		name         string
		long         bool
		longCapacity int
		wantLong     int
	}{
		{name: "normal slots"},
		{name: "long slots", long: true, longCapacity: 2, wantLong: 1},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			m, first, second, slot := leaseRaceFixture(t, tc.long, tc.longCapacity)
			m.trackLease(first, "run-race-1")

			const racers = 32
			var wg sync.WaitGroup
			start := make(chan struct{})
			for i := 0; i < racers; i++ {
				wg.Add(1)
				go func(i int) {
					defer wg.Done()
					<-start
					switch i % 4 {
					case 0:
						m.releaseLease(first, "send-error")
					case 1:
						m.releaseRun("run-race-1", "complete")
					case 2:
						m.releaseRun("run-race-1", "cancelled")
					default:
						m.releaseLease(first, "tunnel-END")
					}
				}(i)
			}
			close(start)
			wg.Wait()

			inflight, longInflight := inflightCount(t, m, "g-race", slot)
			if inflight != 1 {
				t.Fatalf("expected the surviving lease to keep 1 slot, got %d in flight", inflight)
			}
			if longInflight != tc.wantLong {
				t.Fatalf("expected %d long slots held, got %d", tc.wantLong, longInflight)
			}
			if got := leaseCount(m); got != 1 {
				t.Fatalf("expected only the surviving lease to remain, got %d leases", got)
			}

			m.mu.Lock()
			_, stillIndexed := m.leaseByRun["run-race-1"]
			_, survivorLive := m.leases[second]
			m.mu.Unlock()
			if stillIndexed {
				t.Error("released lease is still in the run index")
			}
			if !survivorLive {
				t.Error("the untouched lease was released by the racing calls")
			}
		})
	}
}

// TestProviderLeaseTrackAndSendFailureRace covers the window the old bool guard
// could not: a terminal cause arriving while the dispatch is still binding the
// run id. Whichever side wins, the slot is freed once and no run id is left
// pointing at a lease that no longer exists.
func TestProviderLeaseTrackAndSendFailureRace(t *testing.T) {
	for i := 0; i < 16; i++ {
		m, first, second, slot := leaseRaceFixture(t, false, 0)

		var wg sync.WaitGroup
		start := make(chan struct{})
		wg.Add(2)
		go func() {
			defer wg.Done()
			<-start
			m.trackLease(first, "run-track-race")
		}()
		go func() {
			defer wg.Done()
			<-start
			m.releaseLease(first, "send-error")
		}()
		close(start)
		wg.Wait()

		// The terminal path that lost the race still runs afterwards; it must
		// not free a second slot.
		m.releaseRun("run-track-race", "complete")
		m.releaseLease(first, "tunnel-closed")

		if inflight, _ := inflightCount(t, m, "g-race", slot); inflight != 1 {
			t.Fatalf("iteration %d: expected the surviving lease to keep 1 slot, got %d", i, inflight)
		}
		m.mu.Lock()
		_, stillIndexed := m.leaseByRun["run-track-race"]
		_, survivorLive := m.leases[second]
		m.mu.Unlock()
		if stillIndexed {
			t.Fatalf("iteration %d: run id still indexed after the lease was released", i)
		}
		if !survivorLive {
			t.Fatalf("iteration %d: the untouched lease was released", i)
		}
	}
}

// leaseDisconnectFixture admits a primary lease on one node and a control lease
// on another, returning the manager, both lease ids, the primary slot key, and
// the control slot key. Two separate (node, provider) resources let us observe
// that a node disconnect only touches the disconnected node's counters.
func leaseDisconnectFixture(t *testing.T, long bool, longCapacity int) (m *modelQueueManager, primaryLeaseID, controlLeaseID uint64, primarySlot, controlSlot string) {
	t.Helper()
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-disconn-primary",
		Runtime: config.RuntimeConf{Concurrency: 2},
		Providers: []config.NodeProviderConf{
			{ID: "prov-disconn", Capacity: 2, LongContextCapacity: longCapacity, Adapter: "vllm"},
		},
	})
	store.Add(&edgenode.NodeRecord{
		ID:      "node-disconn-control",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			{ID: "prov-disconn", Capacity: 1, Adapter: "vllm"},
		},
	})

	primaryEntry := &edgenode.NodeEntry{NodeID: "node-disconn-primary"}
	controlEntry := &edgenode.NodeEntry{NodeID: "node-disconn-control"}
	primaryCands := []candidateNode{{
		entry:               primaryEntry,
		capacity:            2,
		providerID:          "prov-disconn",
		longContextCapacity: longCapacity,
	}}
	controlCands := []candidateNode{{
		entry:               controlEntry,
		capacity:            1,
		providerID:          "prov-disconn",
		longContextCapacity: 0,
	}}

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

	// Admit primary lease.
	admitted, _, err := m.admitWithReason(context.Background(), "g-disconn", "", "", primaryCands, groupPolicy{}, nil, long, false)
	if err != nil || admitted == nil {
		t.Fatalf("admit primary: %v", err)
	}
	primaryLeaseID = admitted.leaseID
	primarySlot = primaryEntry.NodeID + ":" + "prov-disconn"

	// Admit control lease on a different node.
	admitted2, _, err := m.admitWithReason(context.Background(), "g-disconn-ctrl", "", "", controlCands, groupPolicy{}, nil, false, false)
	if err != nil || admitted2 == nil {
		t.Fatalf("admit control: %v", err)
	}
	controlLeaseID = admitted2.leaseID
	controlSlot = controlEntry.NodeID + ":" + "prov-disconn"

	return m, primaryLeaseID, controlLeaseID, primarySlot, controlSlot
}

// TestProviderLeaseNodeDisconnectRaceExactlyOnce pins that when the current-owner
// node disconnects, the race between releaseNode and any other terminal cause
// (releaseLease, releaseRun) frees the disconnected slot exactly once, and that
// a later re-admission on the same (node, provider) is isolated from stale
// releases of the old lease identity. Both normal and long-context cases are
// verified.
func TestProviderLeaseNodeDisconnectRaceExactlyOnce(t *testing.T) {
	cases := []struct {
		name         string
		long         bool
		longCapacity int
		wantLong     int
	}{
		{name: "normal slots"},
		{name: "long slots", long: true, longCapacity: 2, wantLong: 1},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			m, primaryLeaseID, controlLeaseID, primarySlot, controlSlot := leaseDisconnectFixture(t, tc.long, tc.longCapacity)

			// Track the primary lease so releaseRun has a run id to resolve.
			primaryRunID := "run-disconn-race"
			m.trackLease(primaryLeaseID, primaryRunID)

			// Barrier: race releaseNode against releaseLease and releaseRun.
			const racers = 32
			var wg sync.WaitGroup
			start := make(chan struct{})
			for i := 0; i < racers; i++ {
				wg.Add(1)
				go func(i int) {
					defer wg.Done()
					<-start
					switch i % 3 {
					case 0:
						m.releaseNode("node-disconn-primary", "transport-closed")
					case 1:
						m.releaseRun(primaryRunID, "complete")
					default:
						m.releaseLease(primaryLeaseID, "tunnel-END")
					}
				}(i)
			}
			close(start)
			wg.Wait()

			// After the race, the disconnected primary resource must be zeroed.
			inflight, longInflight := inflightCount(t, m, "g-disconn", primarySlot)
			if inflight != 0 {
				t.Fatalf("expected disconnected primary inflight=0, got %d", inflight)
			}
			if longInflight != 0 {
				t.Fatalf("expected disconnected primary longInflight=0, got %d", longInflight)
			}

			// The control lease on the other node must be untouched.
			ctrlInflight, ctrlLongInflight := inflightCount(t, m, "g-disconn-ctrl", controlSlot)
			if ctrlInflight != 1 {
				t.Fatalf("expected control inflight=1, got %d", ctrlInflight)
			}
			if ctrlLongInflight != 0 {
				t.Fatalf("expected control longInflight=0, got %d", ctrlLongInflight)
			}

			// The primary lease must be gone from the lease map.
			m.mu.Lock()
			_, primaryStillLive := m.leases[primaryLeaseID]
			_, controlStillLive := m.leases[controlLeaseID]
			_, stillIndexed := m.leaseByRun[primaryRunID]
			m.mu.Unlock()
			if primaryStillLive {
				t.Error("primary lease still in lease map after disconnect race")
			}
			if !controlStillLive {
				t.Error("control lease was released by the racing calls")
			}
			if stillIndexed {
				t.Error("primary run id still indexed after disconnect race")
			}

			// Re-admit on the same (node, provider) that just disconnected.
			primaryEntry := &edgenode.NodeEntry{NodeID: "node-disconn-primary"}
			primaryCands := []candidateNode{{
				entry:               primaryEntry,
				capacity:            2,
				providerID:          "prov-disconn",
				longContextCapacity: tc.longCapacity,
			}}
			newAdmitted, _, err := m.admitWithReason(context.Background(), "g-disconn-re", "", "", primaryCands, groupPolicy{}, nil, tc.long, false)
			if err != nil || newAdmitted == nil {
				t.Fatalf("re-admit: %v", err)
			}
			newLeaseID := newAdmitted.leaseID

			// New lease must hold exactly the expected counters.
			newInflight, newLongInflight := inflightCount(t, m, "g-disconn-re", primarySlot)
			if newInflight != 1 {
				t.Fatalf("expected re-admitted inflight=1, got %d", newInflight)
			}
			if newLongInflight != tc.wantLong {
				t.Fatalf("expected re-admitted longInflight=%d, got %d", tc.wantLong, newLongInflight)
			}

			// Now issue late releases of the OLD lease identity. They must be
			// no-ops: the new lease counters must not decrement.
			m.releaseLease(primaryLeaseID, "late-send-error")
			m.releaseRun(primaryRunID, "late-complete")

			finalInflight, finalLongInflight := inflightCount(t, m, "g-disconn-re", primarySlot)
			if finalInflight != 1 {
				t.Fatalf("expected late old-identity release to be a no-op, inflight=%d (want 1)", finalInflight)
			}
			if finalLongInflight != tc.wantLong {
				t.Fatalf("expected late old-identity release to be a no-op, longInflight=%d (want %d)", finalLongInflight, tc.wantLong)
			}

			// The new lease must still be live.
			m.mu.Lock()
			_, newStillLive := m.leases[newLeaseID]
			m.mu.Unlock()
			if !newStillLive {
				t.Error("new lease was released by stale old-identity calls")
			}
		})
	}
}

func TestQueueReservationReleaseWithCapacityDisabled(t *testing.T) {
	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID:      "node-res-disabled",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			{ID: "prov-1", Capacity: 1, LongContextCapacity: 1, Adapter: "vllm"},
		},
	}
	store.Add(rec)

	entry := &edgenode.NodeEntry{NodeID: "node-res-disabled"}
	selected := &candidateNode{
		entry:               entry,
		capacity:            1,
		providerID:          "prov-1",
		longContextCapacity: 1,
	}
	cands := []candidateNode{*selected}

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

	admitted, _, err := m.admitWithReason(context.Background(), "g-res", "", "", cands, groupPolicy{}, nil, true, false)
	if err != nil || admitted == nil {
		t.Fatalf("admit: %v", err)
	}

	reservation := newQueueReservation(m, admitted)
	if got := m.longInFlightForProvider("node-res-disabled", "prov-1"); got != 1 {
		t.Fatalf("expected longInFlight=1, got %d", got)
	}

	// Disable capacity via refresh
	rec.Providers[0].Capacity = 0
	rec.Providers[0].LongContextCapacity = 0
	m.setStore(store)

	// Release
	reservation.release("send-error")

	if got := m.longInFlightForProvider("node-res-disabled", "prov-1"); got != 0 {
		t.Errorf("longInFlight after capacity disabled release = %d, want 0", got)
	}
}