package nodes

import (
	"context"
	"encoding/json"
	"errors"
	"runtime"
	"sync"
	"time"

	"github.com/nats-io/nats.go"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
	"gorm.io/gorm"

	"github.com/mudler/LocalAI/core/config"
	"github.com/mudler/LocalAI/core/gallery"
	"github.com/mudler/LocalAI/core/services/galleryop"
	"github.com/mudler/LocalAI/core/services/messaging"
	"github.com/mudler/LocalAI/core/services/testutil"
)

// scriptedMessagingClient maps a NATS subject to a canned reply payload
// (or error). Used so each fan-out request can simulate a different worker
// outcome without spinning up real NATS.
type scriptedMessagingClient struct {
	mu                         sync.Mutex
	replies                    map[string][]byte
	errs                       map[string]error
	calls                      []requestCall
	matchedReplies             map[string][]matchedReply
	publishes                  []progressPublishCall
	scheduledProgressPublishes []scheduledProgressPublish
	subscribes                 []string
}

// progressPublishCall records a single Publish invocation. The progress
// publisher tests assert on the sequence of BackendInstallProgressEvent
// values written to a per-op subject, so we capture both subject and the
// decoded event. Named to avoid clashing with the simpler `publishCall`
// already defined in unloader_test.go (which stores raw JSON bytes for
// non-progress assertions).
type progressPublishCall struct {
	Subject string
	Event   messaging.BackendInstallProgressEvent
}

// scheduledProgressPublish queues a batch of BackendInstallProgressEvent
// values to be delivered the next time Subscribe is called with the matching
// subject. This lets master-side tests assert that the adapter installs its
// handler BEFORE publishing the install request, by scripting events to be
// delivered as soon as the subscription appears.
type scheduledProgressPublish struct {
	subject string
	events  []messaging.BackendInstallProgressEvent
}

// matchedReply lets a test script a canned reply that only fires when the
// inbound request matches a predicate. Used by scriptReplyMatching to
// distinguish "install Force=true" (the fallback) from "install Force=false"
// on the same subject.
type matchedReply struct {
	pred        func(messaging.BackendInstallRequest) bool
	reply       []byte
	fallback    []byte
	fallbackErr error
}

func newScriptedMessagingClient() *scriptedMessagingClient {
	return &scriptedMessagingClient{
		replies: map[string][]byte{},
		errs:    map[string]error{},
	}
}

func (s *scriptedMessagingClient) scriptReply(subject string, reply any) {
	raw, err := json.Marshal(reply)
	Expect(err).ToNot(HaveOccurred())
	s.mu.Lock()
	defer s.mu.Unlock()
	s.replies[subject] = raw
}

func (s *scriptedMessagingClient) scriptErr(subject string, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.errs[subject] = err
}

// scriptNoResponders scripts a nats.ErrNoResponders error for `subject` so
// tests can simulate "old worker without backend.upgrade subscription"
// scenarios. Uses the real nats sentinel so errors.Is(...) works at the
// caller (the manager's NoResponders fallback path).
func (s *scriptedMessagingClient) scriptNoResponders(subject string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.errs[subject] = nats.ErrNoResponders
}

// scriptReplyMatching is like scriptReply but the canned reply only fires
// when the inbound request payload matches `pred(req)`. Lets tests
// differentiate "install with Force=true" from "install Force=false" on
// the same subject — useful for asserting the rolling-update fallback
// path actually sets Force=true on its retry.
//
// If `pred` returns false (or the unmarshal of the payload into the
// predicate's expected type fails), the subject falls through to whatever
// was scripted before (or to the unscripted default ErrNoResponders).
func (s *scriptedMessagingClient) scriptReplyMatching(subject string, pred func(messaging.BackendInstallRequest) bool, reply messaging.BackendInstallReply) {
	raw, err := json.Marshal(reply)
	Expect(err).ToNot(HaveOccurred())
	s.mu.Lock()
	defer s.mu.Unlock()
	prev := s.replies[subject] // may be nil — that's fine
	prevErr := s.errs[subject] // may be nil — that's fine
	if s.matchedReplies == nil {
		s.matchedReplies = map[string][]matchedReply{}
	}
	s.matchedReplies[subject] = append(s.matchedReplies[subject], matchedReply{
		pred:        pred,
		reply:       raw,
		fallback:    prev,
		fallbackErr: prevErr,
	})
}

func (s *scriptedMessagingClient) Request(subject string, data []byte, timeout time.Duration) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.calls = append(s.calls, requestCall{Subject: subject, Data: data, Timeout: timeout})

	// Predicate-matched replies take precedence over flat scriptReply.
	if matchers, ok := s.matchedReplies[subject]; ok {
		var req messaging.BackendInstallRequest
		_ = json.Unmarshal(data, &req)
		for _, m := range matchers {
			if m.pred(req) {
				return m.reply, nil
			}
		}
		// No predicate matched — fall through to the recorded fallback
		// (whatever was scripted before scriptReplyMatching took over).
		if matchers[0].fallback != nil {
			return matchers[0].fallback, nil
		}
		if matchers[0].fallbackErr != nil {
			return nil, matchers[0].fallbackErr
		}
		// No fallback either — default to ErrNoResponders.
		return nil, nats.ErrNoResponders
	}

	if err, ok := s.errs[subject]; ok && err != nil {
		return nil, err
	}
	if reply, ok := s.replies[subject]; ok {
		return reply, nil
	}
	// Simulate ErrNoResponders for any unscripted subject so tests fail
	// loudly when they forget to script a node.
	return nil, &fakeNoRespondersErr{}
}

// Publish records each call so progress-publisher tests can assert on the
// stream of events written to a subject. The real messaging.Client JSON
// encodes the payload before sending, but our publisher hands a typed
// struct directly, so we handle both shapes.
func (s *scriptedMessagingClient) Publish(subject string, data any) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	switch ev := data.(type) {
	case messaging.BackendInstallProgressEvent:
		s.publishes = append(s.publishes, progressPublishCall{Subject: subject, Event: ev})
	case []byte:
		var e messaging.BackendInstallProgressEvent
		_ = json.Unmarshal(ev, &e)
		s.publishes = append(s.publishes, progressPublishCall{Subject: subject, Event: e})
	}
	return nil
}

// publishCalls returns every BackendInstallProgressEvent that was published
// to `subject`, in order. Lets tests assert on debounce behavior without
// depending on internal Publish timing.
func (s *scriptedMessagingClient) publishCalls(subject string) []messaging.BackendInstallProgressEvent {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]messaging.BackendInstallProgressEvent, 0)
	for _, c := range s.publishes {
		if c.Subject != subject {
			continue
		}
		out = append(out, c.Event)
	}
	return out
}

// scheduleProgressPublish queues a set of BackendInstallProgressEvent values
// to be delivered on the next Subscribe call matching the per-op progress
// subject. A short delay before delivery gives the subscriber time to install
// its message handler before the events arrive.
func (s *scriptedMessagingClient) scheduleProgressPublish(nodeID, opID string, events []messaging.BackendInstallProgressEvent) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.scheduledProgressPublishes = append(s.scheduledProgressPublishes, scheduledProgressPublish{
		subject: messaging.SubjectNodeBackendInstallProgress(nodeID, opID),
		events:  events,
	})
}

// subscribeCalls returns the subjects on which Subscribe was invoked.
// Used to confirm the master skipped subscription when onProgress was nil.
func (s *scriptedMessagingClient) subscribeCalls() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]string, len(s.subscribes))
	copy(out, s.subscribes)
	return out
}

func (s *scriptedMessagingClient) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) {
	s.mu.Lock()
	s.subscribes = append(s.subscribes, subject)
	matched := []scheduledProgressPublish{}
	remaining := s.scheduledProgressPublishes[:0]
	for _, sp := range s.scheduledProgressPublishes {
		if sp.subject == subject {
			matched = append(matched, sp)
		} else {
			remaining = append(remaining, sp)
		}
	}
	s.scheduledProgressPublishes = remaining
	s.mu.Unlock()

	go func() {
		time.Sleep(20 * time.Millisecond)
		for _, sp := range matched {
			for _, ev := range sp.events {
				raw, _ := json.Marshal(ev)
				handler(raw)
			}
		}
	}()

	return &fakeSubscription{}, nil
}
func (s *scriptedMessagingClient) QueueSubscribe(_ string, _ string, _ func([]byte)) (messaging.Subscription, error) {
	return &fakeSubscription{}, nil
}
func (s *scriptedMessagingClient) QueueSubscribeReply(_ string, _ string, _ func([]byte, func([]byte))) (messaging.Subscription, error) {
	return &fakeSubscription{}, nil
}
func (s *scriptedMessagingClient) SubscribeReply(_ string, _ func([]byte, func([]byte))) (messaging.Subscription, error) {
	return &fakeSubscription{}, nil
}
func (s *scriptedMessagingClient) IsConnected() bool { return true }
func (s *scriptedMessagingClient) Close()            {}

// recordingNodeCall captures a single UpdateNodeProgress invocation so
// per-node OpStatus tests can assert on the sequence of writes the
// DistributedBackendManager fans out into the sink.
type recordingNodeCall struct {
	OpID     string
	NodeID   string
	Progress galleryop.NodeProgress
}

// recordingProgressSink is a test-only nodeProgressSink that just records
// every call. Used by the per-node OpStatus specs below to assert the
// manager wrote the expected terminal and downloading entries.
type recordingProgressSink struct {
	mu    sync.Mutex
	calls []recordingNodeCall
}

func (r *recordingProgressSink) UpdateNodeProgress(opID, nodeID string, np galleryop.NodeProgress) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.calls = append(r.calls, recordingNodeCall{OpID: opID, NodeID: nodeID, Progress: np})
}

func (r *recordingProgressSink) callsFor(opID, nodeID string) []galleryop.NodeProgress {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := []galleryop.NodeProgress{}
	for _, c := range r.calls {
		if c.OpID == opID && c.NodeID == nodeID {
			out = append(out, c.Progress)
		}
	}
	return out
}

// fakeNoRespondersErr is the unscripted-subject default. It matches
// nats.ErrNoResponders by string only - used when a test forgets to script
// a node so the failure is loud but doesn't tickle errors.Is(...) sentinel
// paths the test wasn't deliberately exercising. Tests that DO want the
// real sentinel (e.g. to drive the manager's NoResponders fallback) call
// scriptNoResponders instead, which scripts nats.ErrNoResponders directly.
type fakeNoRespondersErr struct{}

func (e *fakeNoRespondersErr) Error() string { return "no responders" }

// stubLocalBackendManager satisfies galleryop.BackendManager for the
// distributed manager's `local` field. The DeleteBackend path expects to
// call into local first; in distributed mode the frontend rarely has
// backends installed, so returning gallery.ErrBackendNotFound from
// DeleteBackend reproduces the common production case (caller falls
// through to the NATS fan-out, which is what these tests exercise).
type stubLocalBackendManager struct{}

func (stubLocalBackendManager) InstallBackend(_ context.Context, _ *galleryop.ManagementOp[gallery.GalleryBackend, any], _ galleryop.ProgressCallback) error {
	return nil
}
func (stubLocalBackendManager) DeleteBackend(_ string) error { return gallery.ErrBackendNotFound }
func (stubLocalBackendManager) ListBackends() (gallery.SystemBackends, error) {
	return gallery.SystemBackends{}, nil
}
func (stubLocalBackendManager) UpgradeBackend(_ context.Context, _ string, _ string, _ galleryop.ProgressCallback) error {
	return nil
}
func (stubLocalBackendManager) CheckUpgrades(_ context.Context) (map[string]gallery.UpgradeInfo, error) {
	return nil, nil
}
func (stubLocalBackendManager) IsDistributed() bool { return false }

var _ = Describe("DistributedBackendManager", func() {
	var (
		db       *gorm.DB
		registry *NodeRegistry
		mc       *scriptedMessagingClient
		adapter  *RemoteUnloaderAdapter
		mgr      *DistributedBackendManager
		ctx      context.Context
	)

	BeforeEach(func() {
		if runtime.GOOS == "darwin" {
			Skip("testcontainers requires Docker, not available on macOS CI")
		}
		db = testutil.SetupTestDB()
		var err error
		registry, err = NewNodeRegistry(db)
		Expect(err).ToNot(HaveOccurred())

		mc = newScriptedMessagingClient()
		adapter = NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute)
		mgr = &DistributedBackendManager{
			local:    stubLocalBackendManager{},
			adapter:  adapter,
			registry: registry,
		}
		ctx = context.Background()
	})

	// registerHealthyBackend registers an auto-approved backend node and
	// returns it after a fresh fetch (so the ID/Status is correct).
	registerHealthyBackend := func(name, address string) *BackendNode {
		node := &BackendNode{Name: name, NodeType: NodeTypeBackend, Address: address}
		Expect(registry.Register(ctx, node, true)).To(Succeed())
		fetched, err := registry.GetByName(ctx, name)
		Expect(err).ToNot(HaveOccurred())
		Expect(fetched.Status).To(Equal(StatusHealthy))
		return fetched
	}

	registerUnhealthyBackend := func(name, address string) *BackendNode {
		node := registerHealthyBackend(name, address)
		Expect(registry.MarkUnhealthy(ctx, node.ID)).To(Succeed())
		fetched, err := registry.Get(ctx, node.ID)
		Expect(err).ToNot(HaveOccurred())
		return fetched
	}

	op := func(name string) *galleryop.ManagementOp[gallery.GalleryBackend, any] {
		return &galleryop.ManagementOp[gallery.GalleryBackend, any]{
			GalleryElementName: name,
		}
	}

	Describe("InstallBackend", func() {
		Context("when every healthy node replies Success=true", func() {
			It("returns nil", func() {
				n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051")
				n2 := registerHealthyBackend("worker-b", "10.0.0.2:50051")

				mc.scriptReply(messaging.SubjectNodeBackendInstall(n1.ID),
					messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:50100"})
				mc.scriptReply(messaging.SubjectNodeBackendInstall(n2.ID),
					messaging.BackendInstallReply{Success: true, Address: "10.0.0.2:50100"})

				Expect(mgr.InstallBackend(ctx, op("vllm-development"), nil)).To(Succeed())
			})
		})

		Context("when every node replies Success=false with a distinct error", func() {
			It("returns an aggregated error mentioning each node and message", func() {
				n1 := registerHealthyBackend("dgx-casa", "10.0.0.1:50051")
				n2 := registerHealthyBackend("nvidia-thor", "10.0.0.2:50051")

				mc.scriptReply(messaging.SubjectNodeBackendInstall(n1.ID),
					messaging.BackendInstallReply{Success: false, Error: "no child with platform linux/arm64 in index quay.io/...master-cpu-vllm"})
				mc.scriptReply(messaging.SubjectNodeBackendInstall(n2.ID),
					messaging.BackendInstallReply{Success: false, Error: "disk full"})

				err := mgr.InstallBackend(ctx, op("vllm-development"), nil)
				Expect(err).To(HaveOccurred())
				msg := err.Error()
				Expect(msg).To(ContainSubstring("dgx-casa"))
				Expect(msg).To(ContainSubstring("no child with platform linux/arm64"))
				Expect(msg).To(ContainSubstring("nvidia-thor"))
				Expect(msg).To(ContainSubstring("disk full"))
			})
		})

		Context("when one node succeeds and another fails", func() {
			It("returns an error describing the failing node", func() {
				ok := registerHealthyBackend("worker-ok", "10.0.0.1:50051")
				bad := registerHealthyBackend("worker-bad", "10.0.0.2:50051")

				mc.scriptReply(messaging.SubjectNodeBackendInstall(ok.ID),
					messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:50100"})
				mc.scriptReply(messaging.SubjectNodeBackendInstall(bad.ID),
					messaging.BackendInstallReply{Success: false, Error: "out of memory"})

				err := mgr.InstallBackend(ctx, op("vllm-development"), nil)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("worker-bad"))
				Expect(err.Error()).To(ContainSubstring("out of memory"))
				Expect(err.Error()).ToNot(ContainSubstring("worker-ok"))
			})
		})

		Context("when every node is unhealthy at fan-out time", func() {
			It("returns nil — queued nodes are pending retry, not failures", func() {
				registerUnhealthyBackend("worker-a", "10.0.0.1:50051")
				registerUnhealthyBackend("worker-b", "10.0.0.2:50051")

				// No replies scripted: if the manager tried to call Request,
				// it would hit the "no responders" default and we'd see it.
				Expect(mgr.InstallBackend(ctx, op("vllm-development"), nil)).To(Succeed())
				mc.mu.Lock()
				calls := len(mc.calls)
				mc.mu.Unlock()
				Expect(calls).To(Equal(0))
			})
		})

		Context("when there are no nodes registered at all", func() {
			It("returns nil", func() {
				Expect(mgr.InstallBackend(ctx, op("vllm-development"), nil)).To(Succeed())
			})
		})

		Context("when op.TargetNodeID is set to a healthy node", func() {
			It("installs only on that node, leaving the others untouched", func() {
				target := registerHealthyBackend("worker-target", "10.0.0.1:50051")
				other := registerHealthyBackend("worker-other", "10.0.0.2:50051")

				mc.scriptReply(messaging.SubjectNodeBackendInstall(target.ID),
					messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:50100"})
				// No reply scripted for `other`: if InstallBackend fans out
				// to it, the fakeNoRespondersErr default would surface and
				// the test would fail.

				targetedOp := &galleryop.ManagementOp[gallery.GalleryBackend, any]{
					GalleryElementName: "llama-cpp",
					TargetNodeID:       target.ID,
				}
				Expect(mgr.InstallBackend(ctx, targetedOp, nil)).To(Succeed())

				mc.mu.Lock()
				defer mc.mu.Unlock()
				Expect(mc.calls).To(HaveLen(1))
				Expect(mc.calls[0].Subject).To(Equal(messaging.SubjectNodeBackendInstall(target.ID)))
				Expect(mc.calls[0].Subject).ToNot(Equal(messaging.SubjectNodeBackendInstall(other.ID)))
			})
		})

		Context("when op.TargetNodeID is set to a node that does not exist", func() {
			It("returns nil without sending any NATS request", func() {
				registerHealthyBackend("worker-a", "10.0.0.1:50051")

				ghostOp := &galleryop.ManagementOp[gallery.GalleryBackend, any]{
					GalleryElementName: "llama-cpp",
					TargetNodeID:       "this-id-does-not-exist",
				}
				Expect(mgr.InstallBackend(ctx, ghostOp, nil)).To(Succeed())

				mc.mu.Lock()
				defer mc.mu.Unlock()
				Expect(mc.calls).To(BeEmpty())
			})
		})

		Context("when InstallBackend times out on a worker", func() {
			It("returns galleryop.ErrWorkerStillInstalling and keeps the queue row with NextRetryAt pushed out", func() {
				n := registerHealthyBackend("slow", "10.0.0.1:50051")

				// Script a NATS timeout on the install subject. The adapter
				// wraps this into galleryop.ErrWorkerStillInstalling, which
				// the manager should treat as a soft failure.
				mc.scriptErr(messaging.SubjectNodeBackendInstall(n.ID), nats.ErrTimeout)

				err := mgr.InstallBackend(ctx, op("vllm"), nil)
				Expect(err).To(HaveOccurred())
				Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(),
					"expected wrapped ErrWorkerStillInstalling, got %v", err)

				rows, err := registry.ListPendingBackendOps(ctx)
				Expect(err).ToNot(HaveOccurred())
				Expect(rows).To(HaveLen(1))
				Expect(rows[0].Backend).To(Equal("vllm"))
				// The adapter is configured with a 3m install timeout in this
				// suite (NewRemoteUnloaderAdapter above). NextRetryAt should
				// be ~now+3m; a > now+2m bound is safe-but-tight enough to
				// catch the buggy short default (30s exponential backoff).
				Expect(rows[0].NextRetryAt).To(BeTemporally(">", time.Now().Add(2*time.Minute)),
					"NextRetryAt should be pushed to ~now+installTimeout, not the short default")
			})
		})

		Context("end-to-end: timeout then successful reconcile via backend.list", func() {
			It("surfaces the install in ListBackends after the worker finishes", func() {
				// Use the same node-registration helper the Task 5 test uses
				// so the test fixture is identical to the prior context.
				node := registerHealthyBackend("jetson", "10.0.0.2:50051")

				// First install attempt: NATS times out. The adapter wraps
				// this as galleryop.ErrWorkerStillInstalling and the manager
				// keeps the pending_backend_ops row alive with NextRetryAt
				// pushed out (asserted in the previous context).
				mc.scriptErr(messaging.SubjectNodeBackendInstall(node.ID), nats.ErrTimeout)

				err := mgr.InstallBackend(ctx, op("vllm"), nil)
				Expect(err).To(HaveOccurred())
				Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(),
					"expected wrapped ErrWorkerStillInstalling, got %v", err)

				rows, listErr := registry.ListPendingBackendOps(ctx)
				Expect(listErr).ToNot(HaveOccurred())
				Expect(rows).To(HaveLen(1))

				// The worker finished installing in the background. Script
				// backend.list on the same scriptedMessagingClient so the
				// manager's ListBackends fan-out reports the backend.
				mc.scriptReply(messaging.SubjectNodeBackendList(node.ID), messaging.BackendListReply{
					Backends: []messaging.NodeBackendInfo{{Name: "vllm"}},
				})

				backends, listErr := mgr.ListBackends()
				Expect(listErr).ToNot(HaveOccurred())
				Expect(backends).To(HaveKey("vllm"))
				Expect(backends["vllm"].Nodes).To(HaveLen(1))
				Expect(backends["vllm"].Nodes[0].NodeID).To(Equal(node.ID))

				// Phase 1b shipped: ListBackends proactively clears install rows
				// whose intent is now satisfied by backend.list confirmation. The
				// operator UI clears immediately instead of waiting for the next
				// reconciler tick after NextRetryAt.
				rowsAfter, _ := registry.ListPendingBackendOps(ctx)
				Expect(rowsAfter).To(BeEmpty(),
					"install row should clear once backend.list confirms presence on the target node")
			})
		})

		Context("ListBackends clears confirmed install rows", func() {
			It("deletes the pending_backend_ops install row when the backend is reported installed on its target node", func() {
				node := registerHealthyBackend("worker-a", "10.0.0.5:50051")

				// Pre-stage: simulate an admin install that timed out at the NATS
				// round-trip, leaving an install row in the queue.
				mc.scriptErr(messaging.SubjectNodeBackendInstall(node.ID), nats.ErrTimeout)
				err := mgr.InstallBackend(ctx, op("vllm"), nil)
				Expect(err).To(HaveOccurred())
				Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue())

				rows, _ := registry.ListPendingBackendOps(ctx)
				Expect(rows).To(HaveLen(1))

				// Worker finishes installing in the background. backend.list now
				// confirms presence; ListBackends should proactively clear the row.
				mc.scriptReply(messaging.SubjectNodeBackendList(node.ID), messaging.BackendListReply{
					Backends: []messaging.NodeBackendInfo{{Name: "vllm"}},
				})

				backends, listErr := mgr.ListBackends()
				Expect(listErr).ToNot(HaveOccurred())
				Expect(backends).To(HaveKey("vllm"))

				rowsAfter, _ := registry.ListPendingBackendOps(ctx)
				Expect(rowsAfter).To(BeEmpty(),
					"ListBackends should clear install rows whose intent is now satisfied by backend.list")
			})

			It("does NOT clear an upgrade row even if the backend is reported installed", func() {
				node := registerHealthyBackend("worker-b", "10.0.0.6:50051")

				Expect(registry.UpsertPendingBackendOp(ctx, node.ID, "vllm", OpBackendUpgrade, []byte("[]"))).To(Succeed())

				mc.scriptReply(messaging.SubjectNodeBackendList(node.ID), messaging.BackendListReply{
					Backends: []messaging.NodeBackendInfo{{Name: "vllm"}},
				})

				_, listErr := mgr.ListBackends()
				Expect(listErr).ToNot(HaveOccurred())

				rowsAfter, _ := registry.ListPendingBackendOps(ctx)
				Expect(rowsAfter).To(HaveLen(1), "upgrade rows must not be cleared by backend.list presence")
			})
		})

		Context("InstallBackend streams progress events to the caller's progressCb", func() {
			It("invokes progressCb once per worker-published progress event", func() {
				node := registerHealthyBackend("worker-prog", "10.0.0.7:50051")

				mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID), messaging.BackendInstallReply{Success: true, Address: "10.0.0.7:50051"})
				mc.scheduleProgressPublish(node.ID, "op-prog-1", []messaging.BackendInstallProgressEvent{
					{OpID: "op-prog-1", NodeID: node.ID, Backend: "vllm", FileName: "vllm.tar", Current: "100 MB", Total: "1 GB", Percentage: 10},
					{OpID: "op-prog-1", NodeID: node.ID, Backend: "vllm", FileName: "vllm.tar", Current: "1 GB", Total: "1 GB", Percentage: 100},
				})

				type tick struct {
					FileName, Current, Total string
					Percentage               float64
				}
				var (
					pcCalls []tick
					mu      sync.Mutex
				)
				progressCb := func(file, current, total string, pct float64) {
					mu.Lock()
					defer mu.Unlock()
					pcCalls = append(pcCalls, tick{file, current, total, pct})
				}

				opVal := op("vllm")
				opVal.ID = "op-prog-1"
				Expect(mgr.InstallBackend(ctx, opVal, progressCb)).To(Succeed())

				Eventually(func() int {
					mu.Lock()
					defer mu.Unlock()
					return len(pcCalls)
				}, "1s").Should(Equal(2))
				mu.Lock()
				defer mu.Unlock()
				// The adapter dispatches each progress event to its own goroutine
				// (see unloader.go: `go onProgress(ev)`) so two events emitted back
				// to back can land at the bridge in either order. Assert the set of
				// percentages observed contains both ticks, rather than depending
				// on goroutine scheduling for ordering.
				pcts := []float64{pcCalls[0].Percentage, pcCalls[1].Percentage}
				Expect(pcts).To(ConsistOf(10.0, 100.0))
			})
		})

		Context("InstallBackend tolerates silent (pre-Phase-2) workers", func() {
			It("completes successfully even when no progress events are ever published", func() {
				node := registerHealthyBackend("worker-silent", "10.0.0.8:50051")
				mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID), messaging.BackendInstallReply{Success: true, Address: "10.0.0.8:50051"})
				// NO scheduleProgressPublish call - silent worker.

				var ticks int
				var mu sync.Mutex
				progressCb := func(file, current, total string, pct float64) {
					mu.Lock()
					defer mu.Unlock()
					ticks++
				}

				opVal := op("vllm")
				opVal.ID = "op-silent-1"
				Expect(mgr.InstallBackend(ctx, opVal, progressCb)).To(Succeed())

				Consistently(func() int {
					mu.Lock()
					defer mu.Unlock()
					return ticks
				}, "200ms").Should(Equal(0))
			})
		})

		Context("populates per-node OpStatus entries", func() {
			var sink *recordingProgressSink

			BeforeEach(func() {
				// Reconstruct mgr with the recording sink so the new code
				// path (per-node OpStatus writes) is exercised. The default
				// mgr in the outer BeforeEach has progressSink=nil so the
				// pre-existing specs keep verifying the no-sink behavior.
				sink = &recordingProgressSink{}
				appCfg := &config.ApplicationConfig{}
				mgr = NewDistributedBackendManager(appCfg, nil, adapter, registry, sink)
				// stubLocalBackendManager mirrors the production behaviour
				// where the frontend node rarely has the backend installed
				// locally - the NATS fan-out is what these specs verify.
				mgr.local = stubLocalBackendManager{}
			})

			It("emits a success entry for each healthy node visited", func() {
				node := registerHealthyBackend("worker-ok", "10.0.0.9:50051")
				mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID),
					messaging.BackendInstallReply{Success: true, Address: "10.0.0.9:50051"})

				opVal := op("vllm")
				opVal.ID = "op-node-success"
				Expect(mgr.InstallBackend(ctx, opVal, nil)).To(Succeed())

				calls := sink.callsFor("op-node-success", node.ID)
				Expect(calls).ToNot(BeEmpty())
				Expect(calls[len(calls)-1].Status).To(Equal(galleryop.NodeStatusSuccess))
				Expect(calls[len(calls)-1].NodeName).To(Equal("worker-ok"))
			})

			It("emits a running_on_worker entry when NATS times out", func() {
				node := registerHealthyBackend("worker-slow", "10.0.0.10:50051")
				mc.scriptErr(messaging.SubjectNodeBackendInstall(node.ID), nats.ErrTimeout)

				opVal := op("vllm")
				opVal.ID = "op-node-slow"
				// Soft failure: returns wrapped ErrWorkerStillInstalling.
				_ = mgr.InstallBackend(ctx, opVal, nil)

				calls := sink.callsFor("op-node-slow", node.ID)
				Expect(calls).ToNot(BeEmpty())
				Expect(calls[len(calls)-1].Status).To(Equal(galleryop.NodeStatusRunningOnWorker))
			})

			It("emits downloading entries from progress events", func() {
				node := registerHealthyBackend("worker-dl", "10.0.0.11:50051")
				mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID),
					messaging.BackendInstallReply{Success: true})
				mc.scheduleProgressPublish(node.ID, "op-node-dl", []messaging.BackendInstallProgressEvent{
					{OpID: "op-node-dl", NodeID: node.ID, Backend: "vllm", FileName: "vllm.tar", Current: "1 GB", Total: "1 GB", Percentage: 100, Phase: messaging.PhaseDownloading},
				})

				opVal := op("vllm")
				opVal.ID = "op-node-dl"
				Expect(mgr.InstallBackend(ctx, opVal, nil)).To(Succeed())

				Eventually(func() bool {
					for _, np := range sink.callsFor("op-node-dl", node.ID) {
						if np.Status == galleryop.NodeStatusDownloading && np.Percentage == 100.0 {
							return true
						}
					}
					return false
				}, "1s").Should(BeTrue())
			})
		})
	})

	Describe("UpgradeBackend", func() {
		// scriptInstalled tells the worker(s) named in `nodeIDs` to claim
		// `backend` is installed when DistributedBackendManager.ListBackends()
		// fans out backend.list. Anything not scripted defaults to an empty
		// reply, which means "this node has no backends installed" and so
		// upgrade should skip it.
		scriptInstalled := func(backend string, nodeIDs ...string) {
			for _, id := range nodeIDs {
				mc.scriptReply(messaging.SubjectNodeBackendList(id),
					messaging.BackendListReply{Backends: []messaging.NodeBackendInfo{{Name: backend}}})
			}
		}
		scriptNoBackends := func(nodeIDs ...string) {
			for _, id := range nodeIDs {
				mc.scriptReply(messaging.SubjectNodeBackendList(id),
					messaging.BackendListReply{Backends: nil})
			}
		}

		Context("when every node fails to upgrade", func() {
			It("returns an aggregated error", func() {
				n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051")
				n2 := registerHealthyBackend("worker-b", "10.0.0.2:50051")

				scriptInstalled("vllm-development", n1.ID, n2.ID)
				mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n1.ID),
					messaging.BackendUpgradeReply{Success: false, Error: "image manifest not found"})
				mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n2.ID),
					messaging.BackendUpgradeReply{Success: false, Error: "registry unauthorized"})

				err := mgr.UpgradeBackend(ctx, "", "vllm-development", nil)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("worker-a"))
				Expect(err.Error()).To(ContainSubstring("image manifest not found"))
				Expect(err.Error()).To(ContainSubstring("worker-b"))
				Expect(err.Error()).To(ContainSubstring("registry unauthorized"))
			})
		})

		Context("when every node succeeds", func() {
			It("returns nil", func() {
				n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051")
				scriptInstalled("vllm-development", n1.ID)
				mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n1.ID),
					messaging.BackendUpgradeReply{Success: true})
				Expect(mgr.UpgradeBackend(ctx, "", "vllm-development", nil)).To(Succeed())
			})
		})

		// Smart fan-out: only nodes that actually report the backend installed
		// receive the upgrade NATS request. Reproduces the bug where the
		// "Upgrade All" UI button asked a darwin/arm64 worker to upgrade a
		// linux-only backend it never had, producing a "no child with platform
		// darwin/arm64 in index" error and a stuck pending_backend_ops row.
		Context("when only one of two healthy nodes has the backend installed", func() {
			It("upgrades only on that node and skips the other entirely", func() {
				has := registerHealthyBackend("linux-amd64-worker", "10.0.0.1:50051")
				lacks := registerHealthyBackend("mac-mini-m4", "10.0.0.2:50051")

				scriptInstalled("cpu-insightface-development", has.ID)
				scriptNoBackends(lacks.ID)
				mc.scriptReply(messaging.SubjectNodeBackendUpgrade(has.ID),
					messaging.BackendUpgradeReply{Success: true})
				// Deliberately don't script SubjectNodeBackendUpgrade for `lacks`:
				// if the manager attempts it, the scripted-client default returns
				// fakeNoRespondersErr and the assertion below fails loudly.

				Expect(mgr.UpgradeBackend(ctx, "", "cpu-insightface-development", nil)).To(Succeed())

				mc.mu.Lock()
				defer mc.mu.Unlock()
				for _, call := range mc.calls {
					Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendUpgrade(lacks.ID)),
						"upgrade leaked to %s which does not have the backend installed", lacks.Name)
				}
			})
		})

		Context("when no node has the backend installed", func() {
			It("returns a clear error and never attempts an install request", func() {
				n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051")
				scriptNoBackends(n1.ID)

				err := mgr.UpgradeBackend(ctx, "", "vllm-development", nil)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("not installed on any node"))

				mc.mu.Lock()
				defer mc.mu.Unlock()
				for _, call := range mc.calls {
					Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendUpgrade(n1.ID)))
					Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendInstall(n1.ID)))
				}
			})
		})

		// Rolling-update fallback: pre-2026-05-08 workers don't subscribe to
		// backend.upgrade, so the manager catches nats.ErrNoResponders and
		// re-fires the legacy backend.install Force=true on the same node.
		// Drop these specs once the fallback path itself is removed (see
		// managers_distributed.go UpgradeBackend godoc for the deprecation).
		Context("rolling-update fallback", func() {
			It("falls back to backend.install Force=true when upgrade returns ErrNoResponders", func() {
				n := registerHealthyBackend("worker-old", "10.0.0.1:50051")
				scriptInstalled("vllm-development", n.ID)

				// Old worker: no subscriber on backend.upgrade.
				mc.scriptNoResponders(messaging.SubjectNodeBackendUpgrade(n.ID))
				// Fallback re-fires legacy backend.install with Force=true.
				mc.scriptReplyMatching(messaging.SubjectNodeBackendInstall(n.ID),
					func(req messaging.BackendInstallRequest) bool { return req.Force },
					messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:50100"})

				Expect(mgr.UpgradeBackend(ctx, "", "vllm-development", nil)).To(Succeed())
			})

			It("returns the upgrade error when it is not ErrNoResponders", func() {
				n := registerHealthyBackend("worker-bad", "10.0.0.1:50051")
				scriptInstalled("vllm-development", n.ID)

				mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n.ID),
					messaging.BackendUpgradeReply{Success: false, Error: "disk full"})

				err := mgr.UpgradeBackend(ctx, "", "vllm-development", nil)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("disk full"))
			})
		})
	})

	Describe("DeleteBackend", func() {
		Context("when every node fails to delete", func() {
			It("returns an aggregated error", func() {
				n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051")
				n2 := registerHealthyBackend("worker-b", "10.0.0.2:50051")

				mc.scriptReply(messaging.SubjectNodeBackendDelete(n1.ID),
					messaging.BackendDeleteReply{Success: false, Error: "backend not installed"})
				mc.scriptReply(messaging.SubjectNodeBackendDelete(n2.ID),
					messaging.BackendDeleteReply{Success: false, Error: "permission denied"})

				err := mgr.DeleteBackend("vllm-development")
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(ContainSubstring("worker-a"))
				Expect(err.Error()).To(ContainSubstring("backend not installed"))
				Expect(err.Error()).To(ContainSubstring("worker-b"))
				Expect(err.Error()).To(ContainSubstring("permission denied"))
			})
		})

		Context("when every node succeeds", func() {
			It("returns nil", func() {
				n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051")
				mc.scriptReply(messaging.SubjectNodeBackendDelete(n1.ID),
					messaging.BackendDeleteReply{Success: true})
				Expect(mgr.DeleteBackend("vllm-development")).To(Succeed())
			})
		})
	})
})
