package mill

import (
	"context"
	"slices"
	"strings"
	"sync"
	"testing"
	"time"

	"tangled.org/core/spindle/db"
	"tangled.org/core/spindle/models"
	"tangled.org/core/spindle/observability"

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

func workflowOwnedBy(wf *models.Workflow, owner string) *models.Workflow {
	state := wf.Data.(*millWorkflowState)
	md := *state.RawPipeline.TriggerMetadata
	repo := *md.Repo
	repo.Did = owner
	md.Repo = &repo
	state.RawPipeline.TriggerMetadata = &md
	return wf
}

// a node that turns bids away until the test frees seats on it
type seatedNode struct {
	mu   sync.Mutex
	free int
}

func (n *seatedNode) open(m *Mill, seats int) {
	n.mu.Lock()
	n.free += seats
	n.mu.Unlock()
	m.notifyChange()
}

func addSeatedNode(t *testing.T, m *Mill) *seatedNode {
	t.Helper()
	node := &seatedNode{}
	var sess *millSession
	sess = addCandidateSession(t, m, "node-a", nil, 0, scriptedEncoder(func(msg *millproto.Message) error {
		rs := msg.GetReserveSeat()
		if rs == nil {
			return nil
		}
		node.mu.Lock()
		accepted := node.free > 0
		if accepted {
			node.free--
		}
		node.mu.Unlock()
		sess.deliver(rs.GetLeaseId(), &millproto.Message{ReserveResult: &millv1.ReserveResult{
			LeaseId:  rs.GetLeaseId(),
			Accepted: accepted,
		}})
		return nil
	}))
	return node
}

func runningFor(m *Mill, owner string, n int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for range n {
		lease := newLease(m.nextLeaseIDLocked(), "node-a", "", "dummy")
		lease.ownerDID = owner
		m.leases[lease.id] = lease
	}
}

func waitForTriedWaiters(t *testing.T, m *Mill, want int) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		m.mu.Lock()
		tried := 0
		for w := range m.waiters {
			if w.tried {
				tried++
			}
		}
		m.mu.Unlock()
		if tried == want {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("waiters never all bid, want %d", want)
}

func placeInBackground(t *testing.T, m *Mill, wf *models.Workflow) <-chan string {
	t.Helper()
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	placed := make(chan string, 1)
	go func() {
		wid := models.WorkflowId{PipelineId: models.PipelineId("p"), Name: wf.Name}
		if _, err := m.place(ctx, "dummy", wid, wf); err != nil {
			placed <- "error: " + err.Error()
			return
		}
		placed <- wf.Name
	}()
	return placed
}

func nextPlaced(t *testing.T, placed ...<-chan string) string {
	t.Helper()
	cases := make(chan string, len(placed))
	for _, ch := range placed {
		go func() {
			if name, ok := <-ch; ok {
				cases <- name
			}
		}()
	}
	select {
	case name := <-cases:
		return name
	case <-time.After(2 * time.Second):
		t.Fatal("no workflow was placed")
		return ""
	}
}

func TestFreedSeatGoesToTheOwnerWithFewerRunning(t *testing.T) {
	m, _ := quotaMill(t, &scriptedStore{})
	node := addSeatedNode(t, m)
	runningFor(m, "did:plc:busy", 2)

	busy := placeInBackground(t, m, workflowOwnedBy(testWorkflow("busy"), "did:plc:busy"))
	waitForTriedWaiters(t, m, 1)
	quiet := placeInBackground(t, m, workflowOwnedBy(testWorkflow("quiet"), "did:plc:quiet"))
	waitForTriedWaiters(t, m, 2)

	node.open(m, 1)
	if got := nextPlaced(t, busy, quiet); got != "quiet" {
		t.Fatalf("first freed seat went to %q, want the owner with nothing running", got)
	}
}

func TestWaiterThatCannotFitDoesNotBlockTheRest(t *testing.T) {
	m, bdb := quotaMill(t, &scriptedStore{})
	registerTestExecutor(t, bdb, "gpu-1", "hash-gpu-1", []string{"gpu"})
	node := addSeatedNode(t, m)
	runningFor(m, "did:plc:busy", 2)

	// ranks first but the only node with the label it needs is offline
	gpu := placeInBackground(t, m, workflowOwnedBy(testWorkflowWithRunsOn("gpu", []string{"gpu"}), "did:plc:quiet"))
	busy := placeInBackground(t, m, workflowOwnedBy(testWorkflow("busy"), "did:plc:busy"))
	waitForTriedWaiters(t, m, 2)

	node.open(m, 1)
	if got := nextPlaced(t, gpu, busy); got != "busy" {
		t.Fatalf("freed seat went to %q, want the waiter that fits", got)
	}
}

func TestWorkflowCountsNameTheRepoAndSplitWaiting(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	if err := bdb.AddRepo(db.Repo{Knot: "knot", Owner: "did:plc:alice", Rkey: "core", RepoDid: "did:plc:core", Name: "core"}); err != nil {
		t.Fatal(err)
	}

	m.mu.Lock()
	for range 2 {
		lease := newLease(m.nextLeaseIDLocked(), "node-a", "", "dummy")
		lease.ownerDID, lease.repoDID = "did:plc:alice", "did:plc:core"
		m.leases[lease.id] = lease
	}
	m.mu.Unlock()
	m.joinPlacement("did:plc:alice", "did:plc:core", "dummy")
	m.setAdmitting(m.joinPlacement("did:plc:alice", "did:plc:core", "dummy"), true)
	m.joinPlacement("did:plc:bob", "did:plc:fork", "dummy")

	got := m.workflowCounts()
	slices.SortFunc(got, func(a, b observability.MillWorkflows) int {
		return strings.Compare(a.State+a.RepoDID, b.State+b.RepoDID)
	})
	want := []observability.MillWorkflows{
		{State: "running", OwnerDID: "did:plc:alice", RepoDID: "did:plc:core", Repo: "core", Count: 2},
		{State: "waiting_quota", OwnerDID: "did:plc:alice", RepoDID: "did:plc:core", Repo: "core", Count: 1},
		{State: "waiting_seat", OwnerDID: "did:plc:alice", RepoDID: "did:plc:core", Repo: "core", Count: 1},
		{State: "waiting_seat", OwnerDID: "did:plc:bob", RepoDID: "did:plc:fork", Repo: "did:plc:fork", Count: 1},
	}
	if !slices.Equal(got, want) {
		t.Fatalf("workflow counts = %+v, want %+v", got, want)
	}
}