package llm

import (
	"context"
	"testing"
	"time"
)

// Background work must not start while he is waiting on a turn. llama-server
// serves one request at a time, so an extraction that starts first holds the
// slot for its whole budget.
func TestGateBackgroundWaitsForForeground(t *testing.T) {
	g := NewGate(0)
	g.poll = time.Millisecond
	done := g.Foreground()

	started := make(chan struct{})
	go func() {
		release, err := g.AcquireBackground(context.Background())
		if err != nil {
			t.Errorf("acquire: %v", err)
			return
		}
		close(started)
		release()
	}()

	select {
	case <-started:
		t.Fatal("background work started while a foreground request was in flight")
	case <-time.After(20 * time.Millisecond):
	}
	done()
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("background work never started after the foreground request finished")
	}
}

// Only one background request at a time, whatever the queue depth upstream. A
// first poll of a mailbox with 40 unseen messages must not put 40 extractions
// on the slot.
func TestGateOneBackgroundAtATime(t *testing.T) {
	g := NewGate(0)
	g.poll = time.Millisecond
	first, err := g.AcquireBackground(context.Background())
	if err != nil {
		t.Fatalf("first: %v", err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	if _, err := g.AcquireBackground(ctx); err == nil {
		t.Fatal("a second background request ran alongside the first")
	}
	first()
	second, err := g.AcquireBackground(context.Background())
	if err != nil {
		t.Fatalf("second after release: %v", err)
	}
	second()
}

// The quiet window covers the gap between the router call and the phraser call
// of one turn, so an extraction cannot slip in mid-turn.
func TestGateQuietWindow(t *testing.T) {
	now := time.Now()
	g := NewGate(time.Minute)
	g.poll = time.Millisecond
	g.now = func() time.Time { return now }
	g.Foreground()()

	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	if _, err := g.AcquireBackground(ctx); err == nil {
		t.Fatal("background work started inside the quiet window")
	}
	now = now.Add(2 * time.Minute)
	release, err := g.AcquireBackground(context.Background())
	if err != nil {
		t.Fatalf("acquire after the quiet window: %v", err)
	}
	release()
}

// Foreground never waits, whatever else is in flight.
func TestGateForegroundNeverBlocks(t *testing.T) {
	g := NewGate(time.Minute)
	release, err := g.AcquireBackground(context.Background())
	if err != nil {
		t.Fatalf("acquire: %v", err)
	}
	defer release()
	done := make(chan struct{})
	go func() { g.Foreground()(); close(done) }()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("a foreground request waited behind background work")
	}
}