package mill

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"strings"
	"testing"
	"time"

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

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

type scriptedEncoder func(*millproto.Message) error

func (e scriptedEncoder) Encode(msg *millproto.Message) error { return e(msg) }

func testWorkflow(name string) *models.Workflow {
	return &models.Workflow{
		Name:        name,
		Environment: map[string]string{},
		Steps:       []models.Step{remoteStep{}},
		Data: &millWorkflowState{
			TargetEngine: "dummy",
			RawWorkflow:  tangled.Pipeline_Workflow{Name: name},
			RawPipeline:  tangled.Pipeline{},
		},
	}
}

func testWorkflowWithRunsOn(name string, runsOn []string) *models.Workflow {
	wf := testWorkflow(name)
	wf.Data.(*millWorkflowState).RawWorkflow.RunsOn = runsOn
	return wf
}

func addCandidateSession(t *testing.T, m *Mill, nodeID string, labels []string, load float64, enc messageEncoder) *millSession {
	t.Helper()
	if enc == nil {
		enc = scriptedEncoder(func(*millproto.Message) error { return nil })
	}
	sess := newSession(nodeID, enc, slog.New(slog.NewTextHandler(io.Discard, nil)))
	sess.labels = labels
	sess.snapshot = &millv1.NodeSnapshot{
		NodeId: nodeID,
		Engines: map[string]*millv1.EngineAvailability{
			"dummy": {Available: load < 1.0, Load: map[string]float64{"slots": load}},
		},
	}
	m.mu.Lock()
	m.sessions[nodeID] = sess
	m.mu.Unlock()
	return sess
}

func assertRankedNodes(t *testing.T, got []*millSession, want []string) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("rankCandidates() returned %d candidates, want %d: got %v want %v", len(got), len(want), sessionIDs(got), want)
	}
	for i := range want {
		if got[i].nodeID != want[i] {
			t.Fatalf("rankCandidates()[%d] = %q, want %q; full order got %v want %v", i, got[i].nodeID, want[i], sessionIDs(got), want)
		}
	}
}

func sessionIDs(sessions []*millSession) []string {
	out := make([]string, len(sessions))
	for i, sess := range sessions {
		out[i] = sess.nodeID
	}
	return out
}

func sameStringMultiset(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	counts := make(map[string]int, len(a))
	for _, s := range a {
		counts[s]++
	}
	for _, s := range b {
		if counts[s] == 0 {
			return false
		}
		counts[s]--
	}
	return true
}

type reserveReply struct {
	accepted    bool
	rejectClass millv1.RejectClass
	reason      string
}

func addReplyingCandidateSession(t *testing.T, m *Mill, nodeID string, labels []string, load float64, asked chan<- string, reply reserveReply) *millSession {
	t.Helper()
	var sess *millSession
	sess = addCandidateSession(t, m, nodeID, labels, load, scriptedEncoder(func(msg *millproto.Message) error {
		rs := msg.GetReserveSeat()
		if rs == nil {
			return nil
		}
		if asked != nil {
			asked <- nodeID
		}
		sess.deliver(rs.GetLeaseId(), &millproto.Message{ReserveResult: &millv1.ReserveResult{
			LeaseId:      rs.GetLeaseId(),
			Accepted:     reply.accepted,
			RejectReason: reply.reason,
			RejectClass:  reply.rejectClass,
		}})
		return nil
	}))
	return sess
}

func drainAsked(ch <-chan string) []string {
	var out []string
	for {
		select {
		case nodeID := <-ch:
			out = append(out, nodeID)
		default:
			return out
		}
	}
}

func TestCommitRetriesAfterSessionCloseBeforeCommitted(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{BidTimeout: 25 * time.Millisecond, ReconnectGrace: time.Second})
	wf := testWorkflow("build")
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "dummy")
	lease.wid = wid
	wf.Data.(*millWorkflowState).Lease = lease

	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	var sess1 *millSession
	firstCommit := make(chan struct{})
	sess1 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetCommitLease() != nil {
			close(firstCommit)
			m.detachSession(sess1)
		}
		return nil
	}), l)
	m.attachSession(sess1)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- m.commitAndWait(ctx, wf, nil) }()

	select {
	case <-firstCommit:
	case <-ctx.Done():
		t.Fatal("first commit was not sent")
	}

	var sess2 *millSession
	sess2 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetCommitLease() == nil {
			return nil
		}
		leaseID := msg.GetCommitLease().GetLeaseId()
		sess2.deliver(leaseID, &millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}})
		m.onAttemptResult(sess2, &millv1.AttemptResult{
			Offset:         1,
			LeaseId:        leaseID,
			TerminalStatus: string(models.StatusKindSuccess),
		})
		return nil
	}), l)
	m.attachSession(sess2)
	m.sessionReady(sess2)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("commitAndWait() error = %v, want success after reconnect", err)
		}
	case <-ctx.Done():
		t.Fatal("commitAndWait() did not finish after reconnect")
	}
}

func TestDestroyRunningLeaseDoesNotDropCancelledTerminal(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "dummy")
	lease.wid = wid
	lease.setState(leaseRunning)

	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(wid)
	if lease.getState() == leaseDone {
		t.Fatal("destroy sealed the lease before the terminal result")
	}

	lease.deliverTerminal(&millv1.AttemptResult{
		LeaseId:        lease.id,
		TerminalStatus: string(models.StatusKindCancelled),
	})
	res := <-lease.terminal
	if err := terminalError(res.GetTerminalStatus()); !errors.Is(err, engine.ErrCancelled) {
		t.Fatalf("terminalError() = %v, want ErrCancelled", err)
	}
}

func TestPlaceBlocksWhenNoCapacity(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{})
	wf := testWorkflow("build")
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}

	// no executors at all: place must block until ctx expires (user sees pending).
	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
	defer cancel()

	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded", err)
	}
}

func TestRankCandidatesFiltersRequiredLabelsWithANDSemantics(t *testing.T) {
	m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	addCandidateSession(t, m, "linux-high", []string{"linux"}, 0.0, nil)
	addCandidateSession(t, m, "linux-arm", []string{"linux", "arm64"}, 0.25, nil)
	addCandidateSession(t, m, "unlabeled", nil, 0.5, nil)
	addCandidateSession(t, m, "linux-arm-gpu", []string{"linux", "arm64", "gpu"}, 0.75, nil)
	addCandidateSession(t, m, "linux-arm-full", []string{"linux", "arm64"}, 1.0, nil)

	tests := []struct {
		name           string
		requiredLabels []string
		want           []string
	}{
		{
			name: "no required labels keeps old capacity ranking",
			want: []string{"linux-high", "linux-arm", "unlabeled", "linux-arm-gpu"},
		},
		{
			name:           "single required label includes every candidate carrying it",
			requiredLabels: []string{"linux"},
			want:           []string{"linux-high", "linux-arm", "linux-arm-gpu"},
		},
		{
			name:           "all required labels must be present",
			requiredLabels: []string{"linux", "arm64"},
			want:           []string{"linux-arm", "linux-arm-gpu"},
		},
		{
			name:           "one missing required label excludes the candidate",
			requiredLabels: []string{"linux", "arm64", "gpu"},
			want:           []string{"linux-arm-gpu"},
		},
		{
			name:           "unknown required label leaves no candidate",
			requiredLabels: []string{"linux", "arm64", "metal"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assertRankedNodes(t, m.rankCandidates("dummy", tt.requiredLabels, nil), tt.want)
		})
	}
}

func TestPlaceWithMissingRequiredLabelsStaysPendingWithoutReserve(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{BidTimeout: 10 * time.Millisecond})
	reserveSent := make(chan struct{}, 1)
	addCandidateSession(t, m, "linux-only", []string{"linux"}, 0.75, scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetReserveSeat() != nil {
			select {
			case reserveSent <- struct{}{}:
			default:
			}
		}
		return nil
	}))
	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}

	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Millisecond)
	defer cancel()
	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded while job remains pending", err)
	}
	select {
	case <-reserveSent:
		t.Fatal("place() sent ReserveSeat to executor missing a required label")
	default:
	}
}

func TestPlaceKeepsAskingBelowTopKAfterIncompatibleRejects(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{TopK: 2, BidTimeout: time.Second})
	asked := make(chan string, 4)
	addReplyingCandidateSession(t, m, "wrong-label", []string{"linux"}, 0.0, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "wrong-label should not be asked",
	})
	addReplyingCandidateSession(t, m, "incompatible-a", []string{"linux", "arm64"}, 0.0, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "no runner",
	})
	addReplyingCandidateSession(t, m, "incompatible-b", []string{"linux", "arm64"}, 0.25, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "bad image",
	})
	addReplyingCandidateSession(t, m, "compatible-below-window", []string{"linux", "arm64"}, 0.5, asked, reserveReply{
		accepted: true,
	})

	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	slot, err := m.place(ctx, "dummy", wid, wf)
	if err != nil {
		t.Fatalf("place() error = %v, want compatible lower-ranked executor", err)
	}
	defer slot.Release()

	ms, ok := slot.(*millSlot)
	if !ok {
		t.Fatalf("place() slot type = %T, want *millSlot", slot)
	}
	if ms.lease.nodeID != "compatible-below-window" {
		t.Fatalf("place() chose node %q, want compatible-below-window", ms.lease.nodeID)
	}
	if got, want := drainAsked(asked), []string{"incompatible-a", "incompatible-b", "compatible-below-window"}; !sameStringMultiset(got, want) {
		t.Fatalf("ReserveSeat asked nodes = %v, want %v", got, want)
	}
}

func TestPlaceKeepsTransientOnlyRejectsPending(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{TopK: 1, BidTimeout: 10 * time.Millisecond})
	asked := make(chan string, 2)
	addReplyingCandidateSession(t, m, "busy-a", []string{"linux", "arm64"}, 0.5, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_TRANSIENT,
		reason:      "draining",
	})
	addReplyingCandidateSession(t, m, "busy-b", []string{"linux", "arm64"}, 0.75, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_TRANSIENT,
		reason:      "no slot",
	})

	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Millisecond)
	defer cancel()
	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded while transient rejects leave job pending", err)
	}
	if got, want := drainAsked(asked), []string{"busy-a", "busy-b"}; !sameStringMultiset(got, want) {
		t.Fatalf("ReserveSeat asked nodes = %v, want %v", got, want)
	}
}

func TestPlaceReportsOnlyEligiblePermanentIncompatibleRejects(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{TopK: 1, BidTimeout: time.Second})
	asked := make(chan string, 3)
	addReplyingCandidateSession(t, m, "wrong-label", []string{"linux"}, 0.25, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "wrong-label should not be asked",
	})
	addReplyingCandidateSession(t, m, "incompatible-a", []string{"linux", "arm64"}, 0.5, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "no qemu-system-aarch64",
	})
	addReplyingCandidateSession(t, m, "incompatible-b", []string{"linux", "arm64"}, 0.75, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "image arch unsupported",
	})

	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	_, err := m.place(ctx, "dummy", wid, wf)
	if err == nil {
		t.Fatal("place() error = nil, want permanent incompatible failure")
	}
	errText := err.Error()
	for _, want := range []string{"no compatible executor for dummy", "no qemu-system-aarch64", "image arch unsupported"} {
		if !strings.Contains(errText, want) {
			t.Fatalf("place() error = %q, want it to contain %q", errText, want)
		}
	}
	if strings.Contains(errText, "wrong-label should not be asked") {
		t.Fatalf("place() error = %q, included a missing-label candidate as incompatible", errText)
	}
	if got, want := drainAsked(asked), []string{"incompatible-a", "incompatible-b"}; !sameStringMultiset(got, want) {
		t.Fatalf("ReserveSeat asked nodes = %v, want %v", got, want)
	}
}

func TestMaxPendingRejects(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{MaxPending: 1})

	m.mu.Lock()
	m.pending = 1
	m.mu.Unlock()

	wf2 := testWorkflow("b")
	_, err := m.place(context.Background(), "dummy", models.WorkflowId{Name: "b"}, wf2)
	if err == nil {
		t.Fatal("place() past maxPending should error")
	}
}