package server

import (
	"bytes"
	"context"
	"errors"
	"log/slog"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/format"
	"github.com/ollama/ollama/fs/ggml"
	"github.com/ollama/ollama/llm"
	"github.com/ollama/ollama/ml"
	"github.com/ollama/ollama/types/model"
)

func TestMain(m *testing.M) {
	os.Setenv("OLLAMA_DEBUG", "1")
	logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))
	slog.SetDefault(logger)
	os.Exit(m.Run())
}

func TestSchedInit(t *testing.T) {
	ctx, done := context.WithCancel(t.Context())
	defer done()
	s := InitScheduler(ctx)
	s.loadedMu.Lock()
	require.NotNil(t, s.loaded)
	s.loadedMu.Unlock()
}

func TestSchedLoad(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 20*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond

	modelPath, _ := createBinFile(t, ggml.KV{
		"general.architecture":          "llama",
		"llama.context_length":          uint32(32),
		"llama.embedding_length":        uint32(4096),
		"llama.block_count":             uint32(1),
		"llama.attention.head_count":    uint32(32),
		"llama.attention.head_count_kv": uint32(32),
		"tokenizer.ggml.tokens":         []string{" "},
		"tokenizer.ggml.scores":         []float32{0},
		"tokenizer.ggml.token_type":     []int32{0},
	}, []*ggml.Tensor{
		{Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
		{Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
	})

	req := &LlmRequest{
		ctx:             ctx,
		model:           &Model{ModelPath: modelPath},
		opts:            api.DefaultOptions(),
		successCh:       make(chan *runnerRef, 1),
		errCh:           make(chan error, 1),
		sessionDuration: &api.Duration{Duration: 2 * time.Second},
	}
	// Fail to load model first
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		return nil, errors.New("something failed to load model blah")
	}
	gpus := []ml.DeviceInfo{}
	systemInfo := ml.SystemInfo{}
	s.load(req, systemInfo, gpus, false)
	require.Empty(t, req.successCh)
	require.Len(t, req.errCh, 1)
	s.loadedMu.Lock()
	require.Empty(t, s.loaded)
	s.loadedMu.Unlock()
	err := <-req.errCh
	require.Contains(t, err.Error(), "this model may be incompatible")

	server := &mockLlm{vramSize: 10, vramByGPU: map[ml.DeviceID]uint64{}}
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		server.modelPath = model
		return server, nil
	}
	s.load(req, systemInfo, gpus, false)
	select {
	case err := <-req.errCh:
		require.NoError(t, err)
	case resp := <-req.successCh:
		require.Equal(t, uint64(10), resp.vramSize)
		require.Equal(t, uint(1), resp.refCount)
		s.loadedMu.Lock()
		require.Len(t, s.loaded, 1)
		s.loadedMu.Unlock()
	}

	modelPath2, _ := createBinFile(t, ggml.KV{
		"general.architecture":          "llama",
		"llama.context_length":          uint32(32),
		"llama.embedding_length":        uint32(4096),
		"llama.block_count":             uint32(1),
		"llama.attention.head_count":    uint32(32),
		"llama.attention.head_count_kv": uint32(32),
		"tokenizer.ggml.tokens":         []string{" "},
		"tokenizer.ggml.scores":         []float32{0},
		"tokenizer.ggml.token_type":     []int32{0},
	}, []*ggml.Tensor{
		{Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
		{Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
	})

	req.model.ModelPath = modelPath2
	server.waitResp = errors.New("wait failure")
	s.load(req, systemInfo, gpus, false)
	select {
	case err := <-req.errCh:
		require.Contains(t, err.Error(), "wait failure")
	case resp := <-req.successCh:
		t.Fatalf("unexpected success %v", resp)
	}
	s.loadedMu.Lock()
	runner := s.loaded[modelPath2]
	s.loadedMu.Unlock()
	require.NotNil(t, runner)
	require.Equal(t, uint(0), runner.refCount)
	time.Sleep(1 * time.Millisecond)
	require.Len(t, s.expiredCh, 1)
}

func TestSchedLoadStoresEffectiveContextLength(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()

	s := InitScheduler(ctx)
	scenario := newScenarioRequest(t, ctx, "test", 10, nil, map[ml.DeviceID]uint64{})
	scenario.req.opts.NumCtx = 262144
	scenario.req.numCtxAuto = true
	scenario.srv.contextLength = 131072
	s.newServerFn = scenario.newServer

	s.load(scenario.req, ml.SystemInfo{}, nil, false)

	select {
	case err := <-scenario.req.errCh:
		require.NoError(t, err)
	case runner := <-scenario.req.successCh:
		require.Equal(t, 131072, runner.Options.NumCtx)
	}
}

func TestSchedLoadStoresEffectiveExplicitContextLength(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()

	s := InitScheduler(ctx)
	scenario := newScenarioRequest(t, ctx, "test", 10, nil, map[ml.DeviceID]uint64{})
	scenario.req.opts.NumCtx = 262144
	scenario.srv.contextLength = 131072
	s.newServerFn = scenario.newServer

	s.load(scenario.req, ml.SystemInfo{}, nil, false)

	select {
	case err := <-scenario.req.errCh:
		require.NoError(t, err)
	case runner := <-scenario.req.successCh:
		require.Equal(t, 131072, runner.Options.NumCtx)
	}
}

func TestSchedVisionContextFloor(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()

	visionModel := &Model{
		Name: "vision-test",
		Config: model.ConfigV2{
			Capabilities: []string{string(model.CapabilityVision)},
		},
	}

	t.Run("automatic num_ctx is floored", func(t *testing.T) {
		s := InitScheduler(ctx)
		opts := api.DefaultOptions()
		opts.NumCtx = 128

		s.getRunner(ctx, visionModel, opts, nil, true, false, nil)

		req := <-s.pendingReqCh
		require.Equal(t, 2048, req.opts.NumCtx)
		require.True(t, req.numCtxAuto)
	})

	t.Run("explicit num_ctx is floored", func(t *testing.T) {
		s := InitScheduler(ctx)
		opts := api.DefaultOptions()
		opts.NumCtx = 128

		s.getRunner(ctx, visionModel, opts, nil, false, false, nil)

		req := <-s.pendingReqCh
		require.Equal(t, 2048, req.opts.NumCtx)
		require.False(t, req.numCtxAuto)
	})
}

type reqBundle struct {
	ctx     context.Context //nolint:containedctx
	ctxDone func()
	srv     *mockLlm
	req     *LlmRequest
}

func (scenario *reqBundle) newServer(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
	scenario.srv.modelPath = model
	return scenario.srv, nil
}

func newScenarioRequest(t *testing.T, ctx context.Context, modelName string, vramSize uint64, duration *api.Duration, vramByGPU map[ml.DeviceID]uint64) *reqBundle {
	return newScenarioRequestWithContext(t, ctx, modelName, vramSize, duration, vramByGPU, 32)
}

func newScenarioRequestWithContext(t *testing.T, ctx context.Context, modelName string, vramSize uint64, duration *api.Duration, vramByGPU map[ml.DeviceID]uint64, trainCtx uint32) *reqBundle {
	b := &reqBundle{}
	b.ctx, b.ctxDone = context.WithCancel(ctx)
	t.Helper()

	p, _ := createBinFile(t, ggml.KV{
		"general.architecture":          "llama",
		"llama.context_length":          trainCtx,
		"llama.embedding_length":        uint32(4096),
		"llama.block_count":             uint32(1),
		"llama.attention.head_count":    uint32(32),
		"llama.attention.head_count_kv": uint32(32),
		"tokenizer.ggml.tokens":         []string{" "},
		"tokenizer.ggml.scores":         []float32{0},
		"tokenizer.ggml.token_type":     []int32{0},
	}, []*ggml.Tensor{
		{Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
		{Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
	})

	model := &Model{Name: modelName, ModelPath: p}
	if duration == nil {
		duration = &api.Duration{Duration: 5 * time.Millisecond}
	}
	b.req = &LlmRequest{
		ctx:             b.ctx,
		model:           model,
		opts:            api.DefaultOptions(),
		sessionDuration: duration,
		successCh:       make(chan *runnerRef, 1),
		errCh:           make(chan error, 1),
	}
	b.srv = &mockLlm{vramSize: vramSize, vramByGPU: vramByGPU}
	return b
}

func getGpuFn(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
	slog.Info("test getGpuFn called", "runners", runners)
	g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
	g.TotalMemory = 24 * format.GigaByte
	g.FreeMemory = 12 * format.GigaByte
	return []ml.DeviceInfo{g}
}

func getSystemInfoFn() ml.SystemInfo {
	slog.Info("test getSystemInfoFn called")
	return ml.SystemInfo{
		TotalMemory: 32 * format.GigaByte,
		FreeMemory:  26 * format.GigaByte,
	}
}

func TestSchedRequestsSameModelSameRequest(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = getGpuFn
	s.getSystemInfoFn = getSystemInfoFn
	a := newScenarioRequest(t, ctx, "ollama-model-1", 10, &api.Duration{Duration: 5 * time.Millisecond}, nil)
	b := newScenarioRequest(t, ctx, "ollama-model-1", 11, &api.Duration{Duration: 0}, nil)
	b.req.model = a.req.model

	s.newServerFn = a.newServer
	slog.Info("a")
	s.pendingReqCh <- a.req
	require.Len(t, s.pendingReqCh, 1)
	s.Run(ctx)
	select {
	case resp := <-a.req.successCh:
		require.Equal(t, resp.llama, a.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, a.req.errCh)
	case err := <-a.req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}

	// Same runner as first request due to not needing a reload
	s.newServerFn = b.newServer
	slog.Info("b")
	s.pendingReqCh <- b.req
	select {
	case resp := <-b.req.successCh:
		require.Equal(t, resp.llama, a.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, b.req.errCh)
	case err := <-b.req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}
}

func TestSchedRequestsSimpleReloadSameModel(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 5000*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
	g.TotalMemory = 24 * format.GigaByte
	g.FreeMemory = 12 * format.GigaByte
	gMu := sync.Mutex{}
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		slog.Info("test getGpuFn called", "runners", runners)
		gMu.Lock()
		defer gMu.Unlock()
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn
	a := newScenarioRequest(t, ctx, "ollama-model-1", 10, &api.Duration{Duration: 5 * time.Millisecond}, nil)
	b := newScenarioRequest(t, ctx, "ollama-model-1", 20, &api.Duration{Duration: 5 * time.Millisecond}, nil)
	tmpModel := *a.req.model
	b.req.model = &tmpModel

	s.newServerFn = a.newServer
	slog.Info("a")
	s.pendingReqCh <- a.req
	require.Len(t, s.pendingReqCh, 1)
	s.Run(ctx)
	select {
	case resp := <-a.req.successCh:
		require.Equal(t, resp.llama, a.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, a.req.errCh)
	case err := <-a.req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}

	// Trigger a reload
	s.newServerFn = b.newServer
	b.req.model.AdapterPaths = []string{"new"}
	slog.Info("b")
	s.pendingReqCh <- b.req
	// finish first two requests, so model can reload
	time.Sleep(1 * time.Millisecond)
	a.ctxDone()
	// Report recovered VRAM usage
	time.Sleep(1 * time.Millisecond)
	gMu.Lock()
	g.FreeMemory = 24 * format.GigaByte
	gMu.Unlock()
	select {
	case resp := <-b.req.successCh:
		require.Equal(t, resp.llama, b.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, b.req.errCh)
	case err := <-b.req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}
}

func TestSchedRequestsMultipleLoadedModels(t *testing.T) {
	slog.Info("TestRequestsMultipleLoadedModels")
	ctx, done := context.WithTimeout(t.Context(), 1000*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
	g.TotalMemory = 24 * format.GigaByte
	g.FreeMemory = 12 * format.GigaByte
	gMu := sync.Mutex{}
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		slog.Info("test getGpuFn called", "runners", runners)
		gMu.Lock()
		defer gMu.Unlock()
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	// Multiple loaded models
	a := newScenarioRequest(t, ctx, "model-a-1g-gpu", 1*format.GigaByte, nil, map[ml.DeviceID]uint64{{Library: "Metal"}: 1 * format.GigaByte})
	a.req.sessionDuration = &api.Duration{Duration: 5 * time.Millisecond}
	b := newScenarioRequest(t, ctx, "model-b-10g-gpu", 10*format.GigaByte, nil, map[ml.DeviceID]uint64{{Library: "Metal"}: 10 * format.GigaByte})
	b.req.sessionDuration = &api.Duration{Duration: 5 * time.Millisecond}
	c := newScenarioRequest(t, ctx, "model-c-10g-cpu", 10*format.GigaByte, nil, nil /* No GPU load */)
	c.req.opts.NumGPU = 0                                                                                                                         // CPU load, will be allowed
	b.req.sessionDuration = &api.Duration{Duration: 10 * time.Millisecond}                                                                        // longer than b to cause the scheduler to favor unloading b over c
	d := newScenarioRequest(t, ctx, "model-d-10g-gpu", 13*format.GigaByte, nil, map[ml.DeviceID]uint64{{Library: "Metal"}: 13 * format.GigaByte}) // Needs prior unloaded

	s.newServerFn = a.newServer
	slog.Info("Loading A")
	s.pendingReqCh <- a.req
	s.Run(ctx)
	select {
	case resp := <-a.req.successCh:
		require.Equal(t, resp.llama, a.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, a.req.errCh)
	case err := <-a.req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 1)
	s.loadedMu.Unlock()

	t.Setenv("OLLAMA_MAX_LOADED_MODELS", "0")
	s.newServerFn = b.newServer
	slog.Info("Loading B")
	s.pendingReqCh <- b.req
	select {
	case resp := <-b.req.successCh:
		require.Equal(t, resp.llama, b.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, b.req.errCh)
	case err := <-b.req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 2)
	s.loadedMu.Unlock()

	// This is a CPU load with NumGPU = 0 so it should load
	s.newServerFn = c.newServer
	slog.Info("Loading C")
	s.pendingReqCh <- c.req
	select {
	case resp := <-c.req.successCh:
		require.Equal(t, resp.llama, c.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, c.req.errCh)
	case err := <-c.req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		slog.Info("FAIL: scheduler state", "s.loaded", s.loaded)
		t.Fatal("timeout")
	}
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 3)
	s.loadedMu.Unlock()

	// Try to load a model that won't fit
	s.newServerFn = d.newServer
	slog.Info("d")
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 3)
	s.loadedMu.Unlock()
	a.ctxDone() // Won't help since this one isn't big enough to make room
	time.Sleep(2 * time.Millisecond)
	s.pendingReqCh <- d.req
	// finish prior request, so new model can load
	time.Sleep(6 * time.Millisecond)
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 2)
	s.loadedMu.Unlock()
	// Mark b done so it can unload
	b.ctxDone()
	// Report recovered VRAM usage so scheduler will finish waiting and unload
	time.Sleep(1 * time.Millisecond)
	gMu.Lock()
	g.FreeMemory = 24 * format.GigaByte
	gMu.Unlock()
	select {
	case resp := <-d.req.successCh:
		require.Equal(t, resp.llama, d.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, d.req.errCh)
	case <-ctx.Done():
		t.Fatal("timeout")
	}
	// Wait for b to close
closeWait:
	for {
		select {
		case <-ctx.Done():
			t.Fatal("timeout")
		default:
			if b.srv.closeCalled {
				break closeWait
			}
			time.Sleep(1 * time.Millisecond)
		}
	}
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 2)
	s.loadedMu.Unlock()
}

func TestSchedGetRunner(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 3*time.Second)
	defer done()

	a := newScenarioRequest(t, ctx, "ollama-model-1a", 10, &api.Duration{Duration: 2 * time.Millisecond}, nil)
	b := newScenarioRequest(t, ctx, "ollama-model-1b", 10, &api.Duration{Duration: 2 * time.Millisecond}, nil)
	c := newScenarioRequest(t, ctx, "ollama-model-1c", 10, &api.Duration{Duration: 2 * time.Millisecond}, nil)
	t.Setenv("OLLAMA_MAX_QUEUE", "1")
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = getGpuFn
	s.getSystemInfoFn = getSystemInfoFn
	s.newServerFn = a.newServer
	slog.Info("a")
	successCh1a, errCh1a := s.GetRunner(a.ctx, a.req.model, a.req.opts, a.req.sessionDuration)
	require.Len(t, s.pendingReqCh, 1)
	slog.Info("b")
	successCh1b, errCh1b := s.GetRunner(b.ctx, b.req.model, b.req.opts, b.req.sessionDuration)
	require.Len(t, s.pendingReqCh, 1)
	require.Empty(t, successCh1b)
	require.Len(t, errCh1b, 1)
	err := <-errCh1b
	require.Contains(t, err.Error(), "server busy")
	s.Run(ctx)
	select {
	case resp := <-successCh1a:
		require.Equal(t, resp.llama, a.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, errCh1a)
	case err := <-errCh1a:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}
	a.ctxDone() // Set "a" model to idle so it can unload
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 1)
	s.loadedMu.Unlock()

	c.req.model.ModelPath = "bad path"
	slog.Info("c")
	successCh1c, errCh1c := s.GetRunner(c.ctx, c.req.model, c.req.opts, c.req.sessionDuration)
	// Starts in pending channel, then should be quickly processed to return an error
	time.Sleep(50 * time.Millisecond) // Long enough for the "a" model to expire and unload
	require.Empty(t, successCh1c)
	s.loadedMu.Lock()
	require.Empty(t, s.loaded)
	s.loadedMu.Unlock()
	require.Len(t, errCh1c, 1)
	err = <-errCh1c
	require.Contains(t, err.Error(), "bad path")
	b.ctxDone()
}

func TestSchedGetRunnerUsesDigestKeyWhenModelPathEmpty(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	s := InitScheduler(ctx)
	opts := api.DefaultOptions()
	opts.NumCtx = 4

	loadedModel := &Model{Name: "safetensors-a", Digest: "sha-a"}
	loadedRunner := &runnerRef{
		model:       loadedModel,
		modelKey:    schedulerModelKey(loadedModel),
		llama:       &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}},
		Options:     &opts,
		numParallel: 1,
	}

	s.loadedMu.Lock()
	s.loaded[loadedRunner.modelKey] = loadedRunner
	s.loadedMu.Unlock()

	reqModel := &Model{Name: "safetensors-b", Digest: "sha-b"}
	successCh, errCh := s.GetRunner(ctx, reqModel, opts, nil)

	require.Empty(t, successCh)
	require.Empty(t, errCh)
	require.Len(t, s.pendingReqCh, 1)
}

func TestSchedGetRunnerReusesSameDigestWhenModelPathEmpty(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	s := InitScheduler(ctx)
	opts := api.DefaultOptions()
	opts.NumCtx = 4

	loadedModel := &Model{Name: "safetensors-a", Digest: "sha-a"}
	loadedRunner := &runnerRef{
		model:       loadedModel,
		modelKey:    schedulerModelKey(loadedModel),
		llama:       &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}},
		Options:     &opts,
		numParallel: 1,
	}

	s.loadedMu.Lock()
	s.loaded[loadedRunner.modelKey] = loadedRunner
	s.loadedMu.Unlock()

	reqCtx, cancelReq := context.WithCancel(ctx)
	successCh, errCh := s.GetRunner(reqCtx, &Model{Name: "safetensors-a-copy", Digest: "sha-a"}, opts, nil)
	cancelReq()

	select {
	case runner := <-successCh:
		require.Equal(t, loadedRunner, runner)
	default:
		t.Fatal("expected existing runner to be reused")
	}

	require.Empty(t, errCh)
	require.Empty(t, s.pendingReqCh)
}

func TestSchedExpireRunner(t *testing.T) {
	ctx, done := context.WithCancel(t.Context())
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond

	modelPath, _ := createBinFile(t, ggml.KV{
		"general.architecture":          "llama",
		"llama.context_length":          uint32(32),
		"llama.embedding_length":        uint32(4096),
		"llama.block_count":             uint32(1),
		"llama.attention.head_count":    uint32(32),
		"llama.attention.head_count_kv": uint32(32),
		"tokenizer.ggml.tokens":         []string{" "},
		"tokenizer.ggml.scores":         []float32{0},
		"tokenizer.ggml.token_type":     []int32{0},
	}, []*ggml.Tensor{
		{Name: "blk.0.attn.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
		{Name: "output.weight", Kind: uint32(0), Offset: uint64(0), Shape: []uint64{1, 1, 1, 1}, WriterTo: bytes.NewReader(make([]byte, 32))},
	})

	reqCtx, cancelReq := context.WithCancel(ctx)
	defer cancelReq()

	req := &LlmRequest{
		ctx:             reqCtx,
		model:           &Model{ModelPath: modelPath},
		opts:            api.DefaultOptions(),
		successCh:       make(chan *runnerRef, 1),
		errCh:           make(chan error, 1),
		sessionDuration: &api.Duration{Duration: 2 * time.Minute},
	}

	gpus := []ml.DeviceInfo{}
	systemInfo := ml.SystemInfo{}
	server := &mockLlm{vramSize: 10, vramByGPU: map[ml.DeviceID]uint64{}}
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		server.modelPath = model
		return server, nil
	}
	s.load(req, systemInfo, gpus, false)

	select {
	case err := <-req.errCh:
		if err != nil {
			t.Fatalf("expected no errors when loading, got '%s'", err.Error())
		}
	case resp := <-req.successCh:
		s.loadedMu.Lock()
		if resp.refCount != uint(1) || len(s.loaded) != 1 {
			t.Fatalf("expected a model to be loaded")
		}
		s.loadedMu.Unlock()
	}

	completedDone := make(chan struct{})
	go func() {
		defer close(completedDone)
		s.processCompleted(ctx)
	}()

	s.expireRunner(&Model{ModelPath: modelPath})
	cancelReq()

	select {
	case <-s.unloadedCh:
	case <-time.After(5 * time.Second):
		t.Fatal("expected model to be unloaded")
	}

	s.loadedMu.Lock()
	if len(s.loaded) != 0 {
		t.Fatalf("expected model to be unloaded")
	}
	s.loadedMu.Unlock()

	done()
	select {
	case <-completedDone:
	case <-time.After(5 * time.Second):
		t.Fatal("expected completed loop to stop")
	}
}

// TODO - add one scenario that triggers the bogus finished event with positive ref count
func TestSchedPrematureExpired(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 1000*time.Millisecond)
	defer done()

	// Same model, same request
	scenario1a := newScenarioRequest(t, ctx, "ollama-model-1a", 10, &api.Duration{Duration: 100 * time.Millisecond}, nil)
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = getGpuFn
	s.getSystemInfoFn = getSystemInfoFn
	s.newServerFn = scenario1a.newServer
	successCh1a, errCh1a := s.GetRunner(scenario1a.ctx, scenario1a.req.model, scenario1a.req.opts, scenario1a.req.sessionDuration)
	require.Len(t, s.pendingReqCh, 1)
	s.Run(ctx)
	select {
	case resp := <-successCh1a:
		require.Equal(t, resp.llama, scenario1a.srv)
		require.Empty(t, s.pendingReqCh)
		require.Empty(t, errCh1a)
		s.loadedMu.Lock()
		require.Len(t, s.loaded, 1)
		s.loadedMu.Unlock()
		slog.Info("sending premature expired event now")
		s.expiredCh <- resp // Shouldn't happen in real life, but make sure its safe
	case err := <-errCh1a:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}
	time.Sleep(scenario1a.req.sessionDuration.Duration)
	scenario1a.ctxDone()
	time.Sleep(20 * time.Millisecond)
	require.LessOrEqual(t, len(s.finishedReqCh), 1)
	time.Sleep(10 * time.Millisecond)
	require.Empty(t, s.finishedReqCh)
	s.loadedMu.Lock()
	require.Empty(t, s.loaded)
	s.loadedMu.Unlock()

	// also shouldn't happen in real life
	s.finishedReqCh <- scenario1a.req
	time.Sleep(5 * time.Millisecond)
}

func TestSchedUseLoadedRunner(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	req := &LlmRequest{
		ctx:             ctx,
		opts:            api.DefaultOptions(),
		successCh:       make(chan *runnerRef, 1),
		sessionDuration: &api.Duration{Duration: 2},
	}
	finished := make(chan *LlmRequest)
	llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	r1 := &runnerRef{llama: llm1, sessionDuration: 1, numParallel: 1}
	req.useLoadedRunner(r1, finished)
	require.Equal(t, uint(1), r1.refCount)
	require.Equal(t, time.Duration(2), r1.sessionDuration)
	select {
	case success := <-req.successCh:
		require.Equal(t, r1, success)
	case err := <-req.errCh:
		t.Fatal(err.Error())
	case <-ctx.Done():
		t.Fatal("timeout")
	}
	done()
	fin := <-finished
	require.Equal(t, req, fin)
}

func TestSchedUpdateFreeSpace(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()
	gpus := []ml.DeviceInfo{
		{
			DeviceID: ml.DeviceID{
				ID: "1",
			},
		},
		{
			DeviceID: ml.DeviceID{
				ID: "2",
			},
		},
	}
	gpus[0].TotalMemory = 1000
	gpus[0].FreeMemory = 900
	gpus[1].TotalMemory = 2000
	gpus[1].FreeMemory = 1900
	gpuIDs := []ml.DeviceID{
		{
			ID: "1",
		},
		{
			ID: "2",
		},
	}
	llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{{ID: "1"}: 50, {ID: "2"}: 50}}
	llm2 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{{ID: "1"}: 125, {ID: "2"}: 75}}
	r1 := &runnerRef{llama: llm1, gpus: gpuIDs, numParallel: 1}
	r2 := &runnerRef{llama: llm2, gpus: gpuIDs, numParallel: 1}

	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.loadedMu.Lock()
	s.loaded["a"] = r1
	s.loaded["b"] = r2
	s.loadedMu.Unlock()

	s.updateFreeSpace(gpus)
	require.Equal(t, uint64(1000-50-125), gpus[0].FreeMemory)
	require.Equal(t, uint64(2000-50-75), gpus[1].FreeMemory)
}

func TestSchedFindRunnerToUnload(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	r1 := &runnerRef{refCount: 1, sessionDuration: 1, numParallel: 1}
	r2 := &runnerRef{sessionDuration: 2, numParallel: 1}

	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.loadedMu.Lock()
	s.loaded["a"] = r1
	s.loaded["b"] = r2
	s.loadedMu.Unlock()

	resp := s.findRunnerToUnload()
	require.Equal(t, r2, resp)
	r2.refCount = 1
	resp = s.findRunnerToUnload()
	require.Equal(t, r1, resp)
}

func TestSchedNeedsReload(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	do := api.DefaultOptions()
	runner := &runnerRef{
		model: &Model{
			AdapterPaths:   []string{"adapter1"},
			ProjectorPaths: []string{"projector1"},
		},
		Options:     &do,
		llama:       llm,
		numParallel: 1,
	}
	req := &LlmRequest{
		model: &Model{
			AdapterPaths:   []string{"adapter2"},
			ProjectorPaths: []string{"projector2"},
		},
		opts: api.DefaultOptions(),
	}
	resp := runner.needsReload(ctx, req)
	require.True(t, resp)
	req.model.AdapterPaths = runner.model.AdapterPaths
	resp = runner.needsReload(ctx, req)
	require.True(t, resp)
	req.model.ProjectorPaths = runner.model.ProjectorPaths
	runner.loading = true
	req.opts.NumBatch = 1234
	resp = runner.needsReload(ctx, req)
	require.True(t, resp)
	req.opts.NumBatch = runner.Options.NumBatch
	llm.pingResp = errors.New("foo")
	resp = runner.needsReload(ctx, req)
	require.True(t, resp)
	llm.pingResp = nil
	resp = runner.needsReload(ctx, req)
	require.False(t, resp)
	req.opts.NumGPU = 99
	resp = runner.needsReload(ctx, req)
	require.True(t, resp)
	req.opts.NumGPU = -1
	resp = runner.needsReload(ctx, req)
	require.False(t, resp)
	req.contextShift = true
	resp = runner.needsReload(ctx, req)
	require.True(t, resp)
}

func TestResolveContextShift(t *testing.T) {
	trueValue := true
	falseValue := false

	tests := []struct {
		name  string
		shift *bool
		model *Model
		want  bool
	}{
		{name: "unset defaults to shift", want: true},
		{name: "unset deepseek2 disables shift", model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false},
		{name: "unset deepseek2 family disables shift", model: &Model{Config: model.ConfigV2{ModelFamilies: []string{"llama", "deepseek2"}}}, want: false},
		{name: "explicit false disables shift", shift: &falseValue, want: false},
		{name: "explicit false disables shift for deepseek2", shift: &falseValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: false},
		{name: "explicit true enables shift", shift: &trueValue, want: true},
		{name: "explicit true enables shift for deepseek2", shift: &trueValue, model: &Model{Config: model.ConfigV2{ModelFamily: "deepseek2"}}, want: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.want, resolveContextShift(tt.shift, tt.model))
		})
	}
}

func TestSchedNeedsReloadIgnoresAutomaticNumCtxClamp(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	opts := api.DefaultOptions()
	opts.NumCtx = 131072
	model := &Model{}
	runner := &runnerRef{
		model:       model,
		Options:     &opts,
		llama:       llm,
		numParallel: 1,
		numCtxAuto:  true,
	}
	req := &LlmRequest{
		model:      model,
		opts:       api.DefaultOptions(),
		numCtxAuto: true,
	}
	req.opts.NumCtx = 262144

	require.False(t, runner.needsReload(ctx, req))

	req.numCtxAuto = false
	require.True(t, runner.needsReload(ctx, req))
}

func TestSchedNeedsReloadUsesEffectiveAutomaticContextShift(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	opts := api.DefaultOptions()
	opts.NumCtx = 128
	model := &Model{ModelPath: "model.gguf"}
	runner := &runnerRef{
		model:        model,
		Options:      &opts,
		llama:        llm,
		numParallel:  1,
		numCtxAuto:   true,
		contextShift: true,
	}
	req := &LlmRequest{
		model:      model,
		opts:       api.DefaultOptions(),
		numCtxAuto: true,
	}
	req.opts.NumCtx = 262144

	require.False(t, runner.needsReload(ctx, req))

	req.numCtxAuto = false
	require.True(t, runner.needsReload(ctx, req))
}

func TestSchedNeedsReloadUsesEffectiveExplicitContext(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	opts := api.DefaultOptions()
	opts.NumCtx = 2048
	model := &Model{ModelPath: "model.gguf"}
	runner := &runnerRef{
		model:        model,
		Options:      &opts,
		llama:        llm,
		numParallel:  1,
		contextShift: true,
		trainContext: 2048,
	}
	req := &LlmRequest{
		model: model,
		opts:  api.DefaultOptions(),
	}
	req.opts.NumCtx = 262144

	require.False(t, runner.needsReload(ctx, req))

	req.opts.NumCtx = 1024
	require.True(t, runner.needsReload(ctx, req))
}

func TestSchedNeedsReloadIgnoresAutomaticNumBatchDerivation(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	opts := api.DefaultOptions()
	opts.NumBatch = 1024
	model := &Model{}
	runner := &runnerRef{
		model:        model,
		Options:      &opts,
		llama:        llm,
		numParallel:  1,
		numBatchAuto: true,
	}
	req := &LlmRequest{
		model:        model,
		opts:         api.DefaultOptions(),
		numBatchAuto: true,
	}
	req.opts.NumBatch = 512

	require.False(t, runner.needsReload(ctx, req))

	req.numBatchAuto = false
	require.True(t, runner.needsReload(ctx, req))
}

func TestSchedNeedsReloadIgnoresAutomaticUseMMapDefault(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	llm := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	useMmap := false
	opts := api.DefaultOptions()
	opts.UseMMap = &useMmap
	model := &Model{}
	runner := &runnerRef{
		model:       model,
		Options:     &opts,
		llama:       llm,
		numParallel: 1,
		useMMapAuto: true,
	}
	req := &LlmRequest{
		model: model,
		opts:  api.DefaultOptions(),
	}

	require.False(t, runner.needsReload(ctx, req))

	explicitUseMmap := true
	req.opts.UseMMap = &explicitUseMmap
	require.True(t, runner.needsReload(ctx, req))

	req.opts.UseMMap = &useMmap
	require.False(t, runner.needsReload(ctx, req))

	runner.useMMapAuto = false
	req.opts.UseMMap = nil
	require.True(t, runner.needsReload(ctx, req))
}

func TestAutomaticGenerationBatch(t *testing.T) {
	tests := []struct {
		name         string
		effectiveCtx int
		predicted    uint64
		available    uint64
		flash        ml.FlashAttentionType
		gpus         []ml.DeviceInfo
		want         int
	}{
		{
			name:         "small context keeps default",
			effectiveCtx: 4096,
			flash:        ml.FlashAttentionAuto,
			want:         512,
		},
		{
			name:         "medium context uses 1024 with unknown memory",
			effectiveCtx: 32768,
			flash:        ml.FlashAttentionAuto,
			want:         1024,
		},
		{
			name:         "large context uses 2048 with headroom",
			effectiveCtx: 131072,
			predicted:    8 * format.GibiByte,
			available:    14 * format.GibiByte,
			flash:        ml.FlashAttentionAuto,
			want:         2048,
		},
		{
			name:         "large context steps down to 1024 without 2048 headroom",
			effectiveCtx: 131072,
			predicted:    9 * format.GibiByte,
			available:    14 * format.GibiByte,
			flash:        ml.FlashAttentionAuto,
			want:         1024,
		},
		{
			name:         "large context steps down to 1024 for headroom",
			effectiveCtx: 131072,
			predicted:    8 * format.GibiByte,
			available:    11 * format.GibiByte,
			flash:        ml.FlashAttentionAuto,
			want:         1024,
		},
		{
			name:         "medium context steps down to 512 for headroom",
			effectiveCtx: 32768,
			predicted:    8500 * format.MebiByte,
			available:    11 * format.GibiByte,
			flash:        ml.FlashAttentionAuto,
			want:         512,
		},
		{
			name:         "flash attention disabled suppresses promotion",
			effectiveCtx: 131072,
			predicted:    8 * format.GibiByte,
			available:    14 * format.GibiByte,
			flash:        ml.FlashAttentionDisabled,
			gpus:         []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 14 * format.GibiByte}},
			want:         512,
		},
		{
			name:         "constrained CUDA without flash attention uses smaller batch",
			effectiveCtx: 131072,
			predicted:    3 * format.GibiByte,
			available:    6 * format.GibiByte,
			flash:        ml.FlashAttentionDisabled,
			gpus:         []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}, FreeMemory: 6 * format.GibiByte}},
			want:         256,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.want, automaticGenerationBatch(tt.effectiveCtx, tt.predicted, tt.available, tt.flash, tt.gpus))
		})
	}
}

func TestSchedUnloadAllRunners(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 100*time.Millisecond)
	defer done()

	llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	llm2 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.unloadAllRunners()

	r1 := &runnerRef{llama: llm1, numParallel: 1}
	r2 := &runnerRef{llama: llm2, numParallel: 1}

	s.loadedMu.Lock()
	s.loaded["a"] = r1
	s.loaded["b"] = r2
	s.loadedMu.Unlock()
	s.unloadAllRunners()

	require.True(t, llm1.closeCalled)
	require.True(t, llm2.closeCalled)
}

func TestSchedUnload(t *testing.T) {
	llm1 := &mockLlm{vramByGPU: map[ml.DeviceID]uint64{}}
	r1 := &runnerRef{llama: llm1, numParallel: 1}
	r2 := &runnerRef{model: &Model{AdapterPaths: []string{"A"}}, numParallel: 1}
	r1.unload()
	require.True(t, llm1.closeCalled)
	r2.unload()
	require.Nil(t, r2.model)
}

func TestSchedAlreadyCanceled(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	dctx, done2 := context.WithCancel(ctx)
	done2()
	scenario1a := newScenarioRequest(t, dctx, "ollama-model-1", 10, &api.Duration{Duration: 0}, nil)
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	slog.Info("scenario1a")
	s.pendingReqCh <- scenario1a.req
	require.Len(t, s.pendingReqCh, 1)
	s.Run(ctx)
	time.Sleep(5 * time.Millisecond)
	require.Empty(t, s.pendingReqCh)
	require.Empty(t, scenario1a.req.errCh)
	require.Empty(t, scenario1a.req.successCh)
}

// hasLoadedRunner is a test helper that checks if any runner is loaded.
func hasLoadedRunner(s *Scheduler) bool {
	s.loadedMu.Lock()
	defer s.loadedMu.Unlock()
	for _, r := range s.loaded {
		if r.llama != nil {
			return true
		}
	}
	return false
}

func TestSchedLlamaServerEvictsWhenVRAMInsufficient(t *testing.T) {
	// When a llama-server model is predicted to exceed available VRAM,
	// the scheduler should signal eviction before spawning
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	// GPU with very little free memory — the model won't fit
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 0 // no free VRAM — forces eviction
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	// Pre-load a regular model
	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	// Create a request — the model file + KV cache will exceed 100 MiB
	scenario := newScenarioRequest(t, ctx, "llama-server-model", 1*format.GigaByte, nil, nil)

	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		return &mockLlm{modelPath: model}, nil
	}

	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.True(t, needEvict, "expected eviction when predicted VRAM exceeds free memory")
}

func TestSchedLlamaServerExplicitPartialNumGPUSkipsFullFitEviction(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 0
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	scenario := newScenarioRequest(t, ctx, "partial-llama-server-model", 1*format.GigaByte, nil, nil)
	scenario.req.opts.NumGPU = 1
	scenario.srv.vramSize = 0

	called := false
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		called = true
		require.Equal(t, 1, opts.NumGPU)
		return scenario.srv, nil
	}

	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.False(t, needEvict, "explicit partial offload should not trigger full-fit eviction")
	require.True(t, called, "scheduler should try the explicitly partial load")
}

func TestSchedLlamaServerFitsAlongside(t *testing.T) {
	// When a llama-server model is predicted to fit in remaining VRAM,
	// it should load without evicting existing models
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	// GPU with plenty of free memory
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 20 * format.GigaByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	// Pre-load a regular model
	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	// The test model GGUF is tiny (~64 bytes) — should easily fit in 20 GiB
	scenario := newScenarioRequest(t, ctx, "small-llama-server", 1*format.GigaByte, nil, nil)

	s.newServerFn = scenario.newServer

	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	// Should NOT evict — model fits alongside existing
	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.False(t, needEvict, "expected no eviction when model fits in available VRAM")
}

func TestSchedLlamaServerPredictionUsesTotalParallelContext(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	t.Setenv("OLLAMA_NUM_PARALLEL", "2")

	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 900 * format.MebiByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	scenario := newScenarioRequestWithContext(t, ctx, "parallel-context-model", 1*format.GigaByte, nil, nil, 65536)
	scenario.req.opts.NumCtx = 32768

	called := false
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		called = true
		return scenario.srv, nil
	}

	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.True(t, needEvict, "expected eviction when total parallel context exceeds available memory")
	require.False(t, called, "preflight prediction should reject before spawning llama-server")
}

func TestAvailableMemoryForLoadUsesWorstSharedMemoryMeasurement(t *testing.T) {
	tests := []struct {
		name              string
		systemFree        uint64
		gpus              []ml.DeviceInfo
		wantAvailable     uint64
		wantGPUFree       uint64
		wantSystemLimited bool
	}{
		{
			name:       "integrated metal uses lower system free",
			systemFree: 80 * format.GigaByte,
			gpus: []ml.DeviceInfo{{
				DeviceID:   ml.DeviceID{Library: "Metal"},
				Integrated: true,
				FreeMemory: 300 * format.GigaByte,
			}},
			wantAvailable:     80 * format.GigaByte,
			wantGPUFree:       300 * format.GigaByte,
			wantSystemLimited: true,
		},
		{
			name:       "integrated gpu uses lower system free",
			systemFree: 6 * format.GigaByte,
			gpus: []ml.DeviceInfo{{
				DeviceID:   ml.DeviceID{Library: "Vulkan"},
				Integrated: true,
				FreeMemory: 12 * format.GigaByte,
			}},
			wantAvailable:     6 * format.GigaByte,
			wantGPUFree:       12 * format.GigaByte,
			wantSystemLimited: true,
		},
		{
			name:       "discrete metal ignores lower system free",
			systemFree: 6 * format.GigaByte,
			gpus: []ml.DeviceInfo{{
				DeviceID:   ml.DeviceID{Library: "Metal"},
				FreeMemory: 12 * format.GigaByte,
			}},
			wantAvailable: 12 * format.GigaByte,
			wantGPUFree:   12 * format.GigaByte,
		},
		{
			name:       "discrete gpu ignores lower system free",
			systemFree: 6 * format.GigaByte,
			gpus: []ml.DeviceInfo{{
				DeviceID:   ml.DeviceID{Library: "CUDA"},
				FreeMemory: 12 * format.GigaByte,
			}},
			wantAvailable: 12 * format.GigaByte,
			wantGPUFree:   12 * format.GigaByte,
		},
		{
			name:       "mixed gpus only clamp integrated contribution",
			systemFree: 6 * format.GigaByte,
			gpus: []ml.DeviceInfo{
				{
					DeviceID:   ml.DeviceID{Library: "CUDA"},
					FreeMemory: 12 * format.GigaByte,
				},
				{
					DeviceID:   ml.DeviceID{Library: "Vulkan"},
					Integrated: true,
					FreeMemory: 10 * format.GigaByte,
				},
			},
			wantAvailable:     18 * format.GigaByte,
			wantGPUFree:       22 * format.GigaByte,
			wantSystemLimited: true,
		},
		{
			name:       "shared gpu keeps lower adjusted gpu baseline",
			systemFree: 20 * format.GigaByte,
			gpus: []ml.DeviceInfo{{
				DeviceID:   ml.DeviceID{Library: "Metal"},
				Integrated: true,
				FreeMemory: 12 * format.GigaByte,
			}},
			wantAvailable: 12 * format.GigaByte,
			wantGPUFree:   12 * format.GigaByte,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			available, gpuFree, systemLimited := availableMemoryForLoad(ml.SystemInfo{FreeMemory: tt.systemFree}, tt.gpus)
			require.Equal(t, tt.wantAvailable, available)
			require.Equal(t, tt.wantGPUFree, gpuFree)
			require.Equal(t, tt.wantSystemLimited, systemLimited)
		})
	}
}

func TestSelectLlamaServerPlacement(t *testing.T) {
	systemInfo := ml.SystemInfo{FreeMemory: 14 * format.GigaByte}

	tests := []struct {
		name             string
		gpus             []ml.DeviceInfo
		predictedVRAM    uint64
		opts             api.Options
		schedSpread      string
		wantLibrary      string
		wantMainGPU      *int
		wantSelectedGPUs int
		wantGPUID        string
	}{
		{
			name:          "compacts onto largest same-backend GPU",
			predictedVRAM: 8 * format.GigaByte,
			gpus: []ml.DeviceInfo{
				{DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, Name: "small", FreeMemory: 10 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "1", Library: "CUDA"}, Name: "large", FreeMemory: 20 * format.GigaByte},
			},
			opts:             api.DefaultOptions(),
			wantLibrary:      "CUDA",
			wantMainGPU:      testIntPtr(0),
			wantSelectedGPUs: 1,
			wantGPUID:        "1",
		},
		{
			name:          "explicit main gpu selects matching backend group",
			predictedVRAM: 8 * format.GigaByte,
			gpus: []ml.DeviceInfo{
				{DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, FreeMemory: 10 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "0", Library: "ROCm"}, FreeMemory: 20 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "1", Library: "ROCm"}, FreeMemory: 24 * format.GigaByte},
			},
			opts: api.Options{
				Runner: api.Runner{MainGPU: testIntPtr(1), NumGPU: -1},
			},
			wantLibrary:      "ROCm",
			wantMainGPU:      testIntPtr(0),
			wantSelectedGPUs: 1,
			wantGPUID:        "1",
		},
		{
			name:          "integrated GPU is capped by system free memory",
			predictedVRAM: 12 * format.GigaByte,
			gpus: []ml.DeviceInfo{
				{DeviceID: ml.DeviceID{ID: "0", Library: "Metal"}, Integrated: true, FreeMemory: 32 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "1", Library: "Metal"}, FreeMemory: 16 * format.GigaByte},
			},
			opts:             api.DefaultOptions(),
			wantLibrary:      "Metal",
			wantMainGPU:      testIntPtr(0),
			wantSelectedGPUs: 1,
			wantGPUID:        "1",
		},
		{
			name:          "prefers discrete GPU over integrated GPU with more available memory",
			predictedVRAM: 8 * format.GigaByte,
			gpus: []ml.DeviceInfo{
				{DeviceID: ml.DeviceID{ID: "0", Library: "Vulkan"}, Name: "integrated", Integrated: true, FreeMemory: 32 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "1", Library: "Vulkan"}, Name: "discrete", FreeMemory: 10 * format.GigaByte},
			},
			opts:             api.DefaultOptions(),
			wantLibrary:      "Vulkan",
			wantMainGPU:      testIntPtr(0),
			wantSelectedGPUs: 1,
			wantGPUID:        "1",
		},
		{
			name:          "spread disables automatic compaction",
			predictedVRAM: 8 * format.GigaByte,
			schedSpread:   "1",
			gpus: []ml.DeviceInfo{
				{DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, FreeMemory: 10 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "1", Library: "CUDA"}, FreeMemory: 20 * format.GigaByte},
			},
			opts:             api.DefaultOptions(),
			wantLibrary:      "CUDA",
			wantSelectedGPUs: 2,
		},
		{
			name:          "no single fit chooses best backend group for llama-server split",
			predictedVRAM: 30 * format.GigaByte,
			gpus: []ml.DeviceInfo{
				{DeviceID: ml.DeviceID{ID: "0", Library: "CUDA"}, FreeMemory: 10 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "1", Library: "CUDA"}, FreeMemory: 18 * format.GigaByte},
				{DeviceID: ml.DeviceID{ID: "0", Library: "ROCm"}, FreeMemory: 12 * format.GigaByte},
			},
			opts:             api.DefaultOptions(),
			wantLibrary:      "CUDA",
			wantSelectedGPUs: 2,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Setenv("OLLAMA_SCHED_SPREAD", tt.schedSpread)

			selected, launchOpts := selectLlamaServerPlacement(systemInfo, tt.gpus, tt.predictedVRAM, tt.opts)
			require.Len(t, selected, tt.wantSelectedGPUs)
			require.Equal(t, tt.wantLibrary, selected[0].Library)
			if tt.wantGPUID != "" {
				require.Equal(t, tt.wantGPUID, selected[0].ID)
			}
			if tt.wantMainGPU == nil {
				require.Nil(t, launchOpts.MainGPU)
			} else {
				require.NotNil(t, launchOpts.MainGPU)
				require.Equal(t, *tt.wantMainGPU, *launchOpts.MainGPU)
			}
		})
	}
}

func testIntPtr(v int) *int {
	return &v
}

func TestDisableMmapDefaultReason(t *testing.T) {
	useMmap := true
	tests := []struct {
		name          string
		goos          string
		opts          api.Options
		gpus          []ml.DeviceInfo
		blockCount    uint64
		predictedVRAM uint64
		availableVRAM uint64
		want          string
	}{
		{
			name: "explicit use_mmap true wins",
			goos: "windows",
			opts: api.Options{Runner: api.Runner{NumGPU: -1, UseMMap: &useMmap}},
			gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}},
		},
		{
			name: "cpu-only request disables mmap",
			goos: "linux",
			opts: api.Options{Runner: api.Runner{NumGPU: 0}},
			gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}},
			want: "cpu",
		},
		{
			name: "no GPU devices disables mmap",
			goos: "linux",
			opts: api.Options{Runner: api.Runner{NumGPU: -1}},
			want: "cpu",
		},
		{
			name: "windows cuda disables mmap",
			goos: "windows",
			opts: api.Options{Runner: api.Runner{NumGPU: -1}},
			gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}},
			want: "windows_cuda",
		},
		{
			name:       "metal partial offload disables mmap",
			goos:       "darwin",
			opts:       api.Options{Runner: api.Runner{NumGPU: 10}},
			gpus:       []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}},
			blockCount: 20,
			want:       "metal_partial_offload",
		},
		{
			name:       "metal full offload keeps default",
			goos:       "darwin",
			opts:       api.Options{Runner: api.Runner{NumGPU: 21}},
			gpus:       []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}},
			blockCount: 20,
		},
		{
			name:          "metal auto partial offload disables mmap",
			goos:          "darwin",
			opts:          api.Options{Runner: api.Runner{NumGPU: -1}},
			gpus:          []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}},
			predictedVRAM: 30 * format.GigaByte,
			availableVRAM: 20 * format.GigaByte,
			want:          "metal_partial_offload",
		},
		{
			name:          "metal auto full offload keeps default",
			goos:          "darwin",
			opts:          api.Options{Runner: api.Runner{NumGPU: -1}},
			gpus:          []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "Metal"}}},
			predictedVRAM: 10 * format.GigaByte,
			availableVRAM: 20 * format.GigaByte,
		},
		{
			name: "linux cuda keeps default",
			goos: "linux",
			opts: api.Options{Runner: api.Runner{NumGPU: -1}},
			gpus: []ml.DeviceInfo{{DeviceID: ml.DeviceID{Library: "CUDA"}}},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.want, disableMmapDefaultReason(tt.goos, tt.opts, tt.gpus, tt.blockCount, tt.predictedVRAM, tt.availableVRAM))
		})
	}
}

func TestDisableMmapForHostPressure(t *testing.T) {
	gpus := []ml.DeviceInfo{{
		DeviceID:    ml.DeviceID{Library: "CUDA"},
		TotalMemory: 100 * format.GigaByte,
		FreeMemory:  80 * format.GigaByte,
	}}
	systemInfo := ml.SystemInfo{
		TotalMemory: 100 * format.GigaByte,
		FreeMemory:  50 * format.GigaByte,
	}

	require.True(t, disableMmapForHostPressure(
		"linux",
		api.Options{},
		systemInfo,
		gpus,
		20*format.GigaByte,
		25*format.GigaByte,
		30*format.GigaByte,
		80*format.GigaByte,
	))

	useMmap := true
	require.False(t, disableMmapForHostPressure(
		"linux",
		api.Options{Runner: api.Runner{UseMMap: &useMmap}},
		systemInfo,
		gpus,
		20*format.GigaByte,
		25*format.GigaByte,
		30*format.GigaByte,
		80*format.GigaByte,
	), "explicit use_mmap=true should win")

	require.False(t, disableMmapForHostPressure(
		"darwin",
		api.Options{},
		systemInfo,
		gpus,
		20*format.GigaByte,
		25*format.GigaByte,
		30*format.GigaByte,
		80*format.GigaByte,
	), "only the Linux pressure heuristic is restored")

	igpu := append([]ml.DeviceInfo(nil), gpus...)
	igpu[0].Integrated = true
	require.False(t, disableMmapForHostPressure(
		"linux",
		api.Options{},
		systemInfo,
		igpu,
		20*format.GigaByte,
		25*format.GigaByte,
		30*format.GigaByte,
		80*format.GigaByte,
	), "shared-memory GPU loads should keep the normal mmap path")

	require.False(t, disableMmapForHostPressure(
		"linux",
		api.Options{},
		systemInfo,
		gpus,
		20*format.GigaByte,
		25*format.GigaByte,
		70*format.GigaByte,
		80*format.GigaByte,
	), "when VRAM is tight, no-mmap could make partial CPU offload worse")
}

// TestSchedLoadCrashTriggersEvictAllAndRetry verifies that a post-spawn
// Load() OOM while other models are resident signals evict-all-and-retry
// on the first attempt, but fails fast on the second attempt.
func TestSchedLoadCrashTriggersEvictAllAndRetry(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 20 * format.GigaByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	// Pre-load a different model so evict-all has something to evict.
	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "/fake/existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	// newServerFn returns a mockLlm that crashes in Load()
	loadCrash := errors.New("cudaMalloc failed: out of memory")
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		return &mockLlm{modelPath: model, loadErr: loadCrash}, nil
	}

	scenario := newScenarioRequest(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil)
	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	// First attempt: should signal evict-all by returning true and NOT send
	// the error to errCh (so the caller will retry).
	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.True(t, needEvict, "first load OOM should signal eviction")
	require.True(t, scenario.req.oomRetryAttempted, "oomRetryAttempted should be set")
	select {
	case err := <-scenario.req.errCh:
		t.Fatalf("errCh should be empty on first crash, got %v", err)
	default:
	}

	// Second attempt (simulating the retry after processPending evicted all
	// other runners): same crash, but this time oomRetryAttempted is set so
	// load() should fail fast and report the error.
	needEvict = s.load(scenario.req, systemInfo, gpus, true)
	require.False(t, needEvict, "second load OOM should not ask for another eviction")
	select {
	case err := <-scenario.req.errCh:
		require.ErrorIs(t, err, loadCrash)
	case <-time.After(100 * time.Millisecond):
		t.Fatal("expected error on errCh after second crash")
	}
}

func TestSchedLoadOOMReducesAutomaticContextBeforeRetry(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 20 * format.GigaByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "/fake/existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	loadCrash := errors.New("cudaMalloc failed: out of memory")
	var seenNumCtx []int
	var seenNumBatch []int
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		seenNumCtx = append(seenNumCtx, opts.NumCtx)
		seenNumBatch = append(seenNumBatch, opts.NumBatch)
		return &mockLlm{modelPath: model, loadErr: loadCrash}, nil
	}

	scenario := newScenarioRequestWithContext(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil, 131072)
	scenario.req.opts.NumCtx = 262144
	scenario.req.numCtxAuto = true
	scenario.req.numBatchAuto = true
	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.True(t, needEvict, "first automatic-context load OOM should signal eviction and retry")
	require.True(t, scenario.req.oomRetryAttempted)
	require.Equal(t, 32768, scenario.req.opts.NumCtx)
	require.Equal(t, 1024, scenario.req.opts.NumBatch)
	select {
	case err := <-scenario.req.errCh:
		t.Fatalf("errCh should be empty on first crash, got %v", err)
	default:
	}

	needEvict = s.load(scenario.req, systemInfo, gpus, true)
	require.False(t, needEvict, "second load OOM should not ask for another eviction")
	require.Equal(t, []int{262144, 32768}, seenNumCtx)
	require.Equal(t, []int{2048, 1024}, seenNumBatch)
	select {
	case err := <-scenario.req.errCh:
		require.ErrorIs(t, err, loadCrash)
	case <-time.After(100 * time.Millisecond):
		t.Fatal("expected error on errCh after second crash")
	}
}

func TestSchedLoadOOMKeepsExplicitContextBeforeRetry(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 20 * format.GigaByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "/fake/existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	loadCrash := errors.New("cudaMalloc failed: out of memory")
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		return &mockLlm{modelPath: model, loadErr: loadCrash}, nil
	}

	scenario := newScenarioRequestWithContext(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil, 131072)
	scenario.req.opts.NumCtx = 262144
	scenario.req.numCtxAuto = false
	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.True(t, needEvict, "explicit-context load OOM should still evict and retry once")
	require.True(t, scenario.req.oomRetryAttempted)
	require.Equal(t, 262144, scenario.req.opts.NumCtx)
	select {
	case err := <-scenario.req.errCh:
		t.Fatalf("errCh should be empty on first crash, got %v", err)
	default:
	}
}

func TestSchedFirstLoadOOMReducesAutomaticContextAndRetries(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), time.Second)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 20 * format.GigaByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	loadCrash := errors.New("cudaMalloc failed: out of memory")
	var seenNumCtx []int
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		seenNumCtx = append(seenNumCtx, opts.NumCtx)
		if len(seenNumCtx) == 1 {
			return &mockLlm{modelPath: model, loadErr: loadCrash}, nil
		}
		return &mockLlm{modelPath: model, vramSize: 1 * format.GigaByte, contextLength: opts.NumCtx}, nil
	}

	scenario := newScenarioRequestWithContext(t, ctx, "first-load-crashing-model", 1*format.GigaByte, nil, nil, 131072)
	scenario.req.opts.NumCtx = 262144
	scenario.req.numCtxAuto = true

	s.pendingReqCh <- scenario.req
	s.Run(ctx)

	select {
	case runner := <-scenario.req.successCh:
		require.Equal(t, 32768, runner.Options.NumCtx)
		require.Equal(t, []int{262144, 32768}, seenNumCtx)
	case err := <-scenario.req.errCh:
		t.Fatalf("expected retry success, got error %v", err)
	case <-ctx.Done():
		t.Fatal("timed out waiting for first-load retry")
	}
}

// TestSchedLoadCrashNoOtherModelsFailsFast verifies that a Load() crash with
// no other resident models reports the error immediately (no retry).
func TestSchedLoadCrashNoOtherModelsFailsFast(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 20 * format.GigaByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	loadCrash := errors.New("simulated llama-server OOM crash")
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		return &mockLlm{modelPath: model, loadErr: loadCrash}, nil
	}

	scenario := newScenarioRequest(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil)
	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.False(t, needEvict, "crash with no other runners should not ask for eviction")
	require.False(t, scenario.req.oomRetryAttempted, "oomRetryAttempted must stay false")
	select {
	case err := <-scenario.req.errCh:
		require.ErrorIs(t, err, loadCrash)
	case <-time.After(100 * time.Millisecond):
		t.Fatal("expected error on errCh immediately")
	}
}

func TestSchedLoadNonOOMWithOtherModelsFailsFast(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()
	s := InitScheduler(ctx)
	s.waitForRecovery = 10 * time.Millisecond
	s.getGpuFn = func(ctx context.Context, runners []ml.FilteredRunnerDiscovery) []ml.DeviceInfo {
		g := ml.DeviceInfo{DeviceID: ml.DeviceID{Library: "Metal"}}
		g.TotalMemory = 24 * format.GigaByte
		g.FreeMemory = 20 * format.GigaByte
		return []ml.DeviceInfo{g}
	}
	s.getSystemInfoFn = getSystemInfoFn

	s.loadedMu.Lock()
	s.loaded["existing-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "/fake/existing"},
		modelKey: "existing-model",
	}
	s.loadedMu.Unlock()

	loadCrash := errors.New("server parse failed")
	s.newServerFn = func(systemInfo ml.SystemInfo, gpus []ml.DeviceInfo, model string, f *ggml.GGML, adapters []string, projectors []string, opts api.Options, numParallel int, _ llm.LlamaServerConfig) (llm.LlamaServer, error) {
		return &mockLlm{modelPath: model, loadErr: loadCrash}, nil
	}

	scenario := newScenarioRequest(t, ctx, "crashing-model", 1*format.GigaByte, nil, nil)
	systemInfo := getSystemInfoFn()
	gpus := s.getGpuFn(ctx, nil)

	needEvict := s.load(scenario.req, systemInfo, gpus, true)
	require.False(t, needEvict, "non-OOM load crash should not ask for eviction")
	require.False(t, scenario.req.oomRetryAttempted, "oomRetryAttempted must stay false")
	select {
	case err := <-scenario.req.errCh:
		require.ErrorIs(t, err, loadCrash)
	case <-time.After(100 * time.Millisecond):
		t.Fatal("expected error on errCh immediately")
	}
}

func TestSchedRuntimeOOMExpiresLoadedRunners(t *testing.T) {
	ctx, done := context.WithCancel(t.Context())
	defer done()
	s := InitScheduler(ctx)

	currentModel := &Model{ModelPath: "/tmp/current.gguf"}
	current := &runnerRef{
		model:           currentModel,
		modelKey:        schedulerModelKey(currentModel),
		sessionDuration: time.Hour,
		llama:           &mockLlm{modelPath: "/tmp/current.gguf"},
	}
	otherModel := &Model{ModelPath: "/tmp/other.gguf"}
	other := &runnerRef{
		model:           otherModel,
		modelKey:        schedulerModelKey(otherModel),
		sessionDuration: time.Hour,
		llama:           &mockLlm{modelPath: "/tmp/other.gguf"},
	}

	s.loadedMu.Lock()
	s.loaded[current.modelKey] = current
	s.loaded[other.modelKey] = other
	s.loadedMu.Unlock()

	s.expireRunnersForRuntimeOOM(currentModel, errors.New("cudaMalloc failed: out of memory"))

	require.Equal(t, time.Duration(0), current.sessionDuration)
	require.Equal(t, time.Duration(0), other.sessionDuration)
	require.Len(t, s.expiredCh, 2)
}

func TestSchedLlamaServerEvictsExistingOnPending(t *testing.T) {
	// When a llama-server runner is already loaded and a new model is requested,
	// the scheduler should evict the llama-server runner
	ctx, done := context.WithCancel(t.Context())
	defer done()
	s := InitScheduler(ctx)

	// Load a llama-server runner
	s.loadedMu.Lock()
	s.loaded["llama-model"] = &runnerRef{
		llama:    &mockLlm{modelPath: "/tmp/model.gguf"},
		modelKey: "llama-model",
	}
	s.loadedMu.Unlock()

	require.True(t, hasLoadedRunner(s))

	// The findRunnerToUnload should find and return the llama-server runner
	runner := s.findRunnerToUnload()
	require.NotNil(t, runner)
}

type mockLlm struct {
	modelPath         string
	pingResp          error
	waitResp          error
	completionResp    error
	embeddingResp     []float32
	embeddingRespErr  error
	tokenizeResp      []int
	tokenizeRespErr   error
	detokenizeResp    string
	detonekizeRespErr error
	closeResp         error
	closeCalled       bool
	vramSize          uint64
	totalSize         uint64
	contextLength     int
	vramByGPU         map[ml.DeviceID]uint64

	// loadErr, if non-nil, is returned from Load() to simulate a post-spawn
	// load failure (e.g. llama-server crashing due to under-predicted VRAM).
	loadErr error
}

func (s *mockLlm) ModelPath() string {
	return s.modelPath
}

func (s *mockLlm) Load(ctx context.Context, sytemInfo ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error) {
	if s.loadErr != nil {
		return nil, s.loadErr
	}
	if requireFull {
		if len(gpus) == 0 {
			slog.Info("mockLlm.Load CPU based load")
			return nil, nil
		}
		for _, g := range gpus {
			if g.FreeMemory >= s.vramSize {
				return []ml.DeviceID{g.DeviceID}, nil
			}
		}

		return nil, llm.ErrLoadRequiredFull
	}
	gpuIDs := make([]ml.DeviceID, len(gpus))
	for i := range gpus {
		gpuIDs[i] = gpus[i].DeviceID
	}
	return gpuIDs, nil
}
func (s *mockLlm) Ping(ctx context.Context) error             { return s.pingResp }
func (s *mockLlm) WaitUntilRunning(ctx context.Context) error { return s.waitResp }
func (s *mockLlm) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
	return s.completionResp
}

func (s *mockLlm) Chat(ctx context.Context, req llm.ChatRequest, fn func(llm.ChatResponse)) error {
	return errors.New("not implemented")
}

func (s *mockLlm) ApplyChatTemplate(ctx context.Context, req llm.ChatRequest) (string, error) {
	return "", errors.New("not implemented")
}

func (s *mockLlm) Embedding(ctx context.Context, input string) ([]float32, int, error) {
	return s.embeddingResp, 0, s.embeddingRespErr
}

func (s *mockLlm) Tokenize(ctx context.Context, content string) ([]int, error) {
	return s.tokenizeResp, s.tokenizeRespErr
}

func (s *mockLlm) Detokenize(ctx context.Context, tokens []int) (string, error) {
	return s.detokenizeResp, s.detonekizeRespErr
}

func (s *mockLlm) Close() error {
	s.closeCalled = true
	return s.closeResp
}
func (s *mockLlm) MemorySize() (uint64, uint64)                       { return s.totalSize, s.vramSize }
func (s *mockLlm) VRAMByGPU(id ml.DeviceID) uint64                    { return s.vramByGPU[id] }
func (s *mockLlm) Pid() int                                           { return -1 }
func (s *mockLlm) GetPort() int                                       { return -1 }
func (s *mockLlm) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo { return nil }
func (s *mockLlm) HasExited() bool                                    { return false }
func (s *mockLlm) GetActiveDeviceIDs() []ml.DeviceID                  { return nil }
func (s *mockLlm) ContextLength() int                                 { return s.contextLength }

// TestImageGenRunnerCanBeEvicted verifies that an image generation model
// loaded in the scheduler can be evicted when idle.
func TestImageGenRunnerCanBeEvicted(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()

	s := InitScheduler(ctx)
	s.getGpuFn = getGpuFn
	s.getSystemInfoFn = getSystemInfoFn

	// Simulate an image gen runner already loaded
	imageGenRunner := &runnerRef{
		model:           &Model{Name: "z-image", ModelPath: "/fake/image/model"},
		modelPath:       "/fake/image/model",
		llama:           &mockLlm{vramSize: 21 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{}},
		sessionDuration: 5 * time.Millisecond,
		refCount:        0, // idle
	}

	s.loadedMu.Lock()
	s.loaded["/fake/image/model"] = imageGenRunner
	s.loadedMu.Unlock()

	// Verify the image gen runner is loaded
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 1)
	s.loadedMu.Unlock()

	// findRunnerToUnload should find the idle image gen runner
	runner := s.findRunnerToUnload()
	require.NotNil(t, runner)
	require.Equal(t, "/fake/image/model", runner.modelPath)
}

// TestImageGenSchedulerCoexistence verifies that image generation models
// can coexist with language models in the scheduler and VRAM is tracked correctly.
func TestImageGenSchedulerCoexistence(t *testing.T) {
	ctx, done := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer done()

	s := InitScheduler(ctx)
	s.getGpuFn = getGpuFn
	s.getSystemInfoFn = getSystemInfoFn

	// Load both an imagegen runner and a language model runner
	imageGenRunner := &runnerRef{
		model:           &Model{Name: "flux", ModelPath: "/fake/flux/model"},
		modelPath:       "/fake/flux/model",
		llama:           &mockLlm{vramSize: 8 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{{Library: "Metal"}: 8 * format.GigaByte}},
		sessionDuration: 10 * time.Millisecond,
		numParallel:     1,
		refCount:        0,
	}

	langModelRunner := &runnerRef{
		model:           &Model{Name: "llama3", ModelPath: "/fake/llama3/model"},
		modelPath:       "/fake/llama3/model",
		llama:           &mockLlm{vramSize: 4 * format.GigaByte, vramByGPU: map[ml.DeviceID]uint64{{Library: "Metal"}: 4 * format.GigaByte}},
		sessionDuration: 10 * time.Millisecond,
		numParallel:     1,
		refCount:        0,
	}

	s.loadedMu.Lock()
	s.loaded["/fake/flux/model"] = imageGenRunner
	s.loaded["/fake/llama3/model"] = langModelRunner
	s.loadedMu.Unlock()

	// Verify both are loaded
	s.loadedMu.Lock()
	require.Len(t, s.loaded, 2)
	require.NotNil(t, s.loaded["/fake/flux/model"])
	require.NotNil(t, s.loaded["/fake/llama3/model"])
	s.loadedMu.Unlock()

	// Verify updateFreeSpace accounts for both
	gpus := []ml.DeviceInfo{
		{
			DeviceID:    ml.DeviceID{Library: "Metal"},
			TotalMemory: 24 * format.GigaByte,
			FreeMemory:  24 * format.GigaByte,
		},
	}
	s.updateFreeSpace(gpus)

	// Free memory should be reduced by both models
	expectedFree := uint64(24*format.GigaByte) - uint64(8*format.GigaByte) - uint64(4*format.GigaByte)
	require.Equal(t, expectedFree, gpus[0].FreeMemory)
}
