package mill

import (
	"errors"
	"io"
	"log/slog"
	"testing"

	"tangled.org/core/api/tangled"
	"tangled.org/core/spindle/models"

	millproto "tangled.org/core/spindle/mill/proto"
)

type testPlacementPlanner struct {
	requirements []string
	err          error
}

func (p testPlacementPlanner) Requirements(tangled.Pipeline_Workflow, tangled.Pipeline) ([]string, error) {
	return p.requirements, p.err
}

func TestMillEngineStoresCanonicalPlacementRequirements(t *testing.T) {
	m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	eng := NewEngineWithPlanner("microvm", m, testPlacementPlanner{
		requirements: []string{"runner/qemu", "image/nixos", "runner/qemu"},
	})

	wf, err := eng.InitWorkflow(tangled.Pipeline_Workflow{Name: "build"}, tangled.Pipeline{})
	if err != nil {
		t.Fatalf("InitWorkflow() error = %v", err)
	}
	got := wf.Data.(*millWorkflowState).RequiredCapabilities
	want := []string{"image/nixos", "runner/qemu"}
	if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
		t.Fatalf("required capabilities = %v, want %v", got, want)
	}
}

func TestMillEngineRejectsInvalidPlacementRequirements(t *testing.T) {
	m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	plannerErr := errors.New("invalid image")
	eng := NewEngineWithPlanner("microvm", m, testPlacementPlanner{err: plannerErr})
	if _, err := eng.InitWorkflow(tangled.Pipeline_Workflow{}, tangled.Pipeline{}); !errors.Is(err, plannerErr) {
		t.Fatalf("InitWorkflow() error = %v, want wrapped %v", err, plannerErr)
	}

	eng = NewEngineWithPlanner("microvm", m, testPlacementPlanner{requirements: []string{"image/nixos arm"}})
	if _, err := eng.InitWorkflow(tangled.Pipeline_Workflow{}, tangled.Pipeline{}); err == nil {
		t.Fatal("InitWorkflow() accepted whitespace-containing capability")
	}
}

func TestRankCandidatesIntersectsLabelsAndEngineCapabilities(t *testing.T) {
	m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	incapable := addCandidateSession(t, m, "low-load-incapable", []string{"region/eu"}, 0.05, nil)
	incapable.snapshot.Engines["dummy"].Capabilities = []string{"image/nixos-aarch64"}
	wrongRegion := addCandidateSession(t, m, "wrong-region", []string{"region/us"}, 0.10, nil)
	wrongRegion.snapshot.Engines["dummy"].Capabilities = []string{"image/nixos", "runner/qemu"}
	partial := addCandidateSession(t, m, "partial", []string{"region/eu"}, 0.20, nil)
	partial.snapshot.Engines["dummy"].Capabilities = []string{"image/nixos"}
	eligible := addCandidateSession(t, m, "eligible", []string{"region/eu"}, 0.70, nil)
	eligible.snapshot.Engines["dummy"].Capabilities = []string{"runner/qemu", "image/nixos"}

	got := m.rankCandidates("dummy", []string{"region/eu"}, []string{"image/nixos", "runner/qemu"})
	assertRankedNodes(t, got, []string{"eligible"})
}

func TestBidWithMissingEngineCapabilitySendsNoReserve(t *testing.T) {
	m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	reserveSent := false
	sess := addCandidateSession(t, m, "arm", nil, 0.05, scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetReserveSeat() != nil {
			reserveSent = true
		}
		return nil
	}))
	sess.snapshot.Engines["dummy"].Capabilities = []string{"image/nixos-aarch64"}
	wf := testWorkflow("build")
	wf.Data.(*millWorkflowState).RequiredCapabilities = []string{"image/nixos"}

	lease, err := m.bid(t.Context(), "dummy", wf)
	if err != nil {
		t.Fatalf("bid() error = %v", err)
	}
	if lease != nil {
		t.Fatalf("bid() lease = %v, want nil", lease)
	}
	if reserveSent {
		t.Fatal("bid() sent ReserveSeat to an engine missing a required capability")
	}
}

func TestRequiredCapabilitiesWithoutMillStateAreEmpty(t *testing.T) {
	if got := requiredCapabilities(&models.Workflow{}); got != nil {
		t.Fatalf("requiredCapabilities() = %v, want nil", got)
	}
}