iop/apps/edge/internal/service/provider_resolution.go

package service

import (
	"fmt"
	"strings"
	"time"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
)

// resolveQueueCandidates returns candidate nodes and the group policy for the
// given request. Provider-pool requests are resolved via the model catalog;
// legacy requests are filtered by adapter/target capability.
func (s *Service) resolveQueueCandidates(req SubmitRunRequest) ([]candidateNode, groupPolicy, error) {
	store, catalog, _ := s.runtimeConfigSnapshot()
	if req.ProviderPool {
		return s.resolveProviderPoolCandidates(req, store, catalog)
	}
	if req.NodeRef != "" {
		entry, err := s.ResolveNode(req.NodeRef)
		if err != nil {
			return nil, groupPolicy{}, err
		}
		cap := defaultNodeCapacity
		if store != nil {
			if rec, ok := store.FindByID(entry.NodeID); ok {
				res := resolveAdapterForNode(rec, req.Adapter, req.Target)
				if !res.supported {
					msg := fmt.Sprintf("node %q does not support adapter %q target %q", entry.NodeID, req.Adapter, req.Target)
					if res.ambiguous {
						msg = fmt.Sprintf("node %q: adapter %q is ambiguous (multiple enabled instances); use an instance key", entry.NodeID, req.Adapter)
					}
					return nil, groupPolicy{}, fmt.Errorf("%s", msg)
				}
				cap = res.capacity
			}
		}
		policy := groupPolicyFromRequestOrStore(req, store, []*edgenode.NodeEntry{entry})
		return []candidateNode{{entry: entry, capacity: cap}}, policy, nil
	}

	all := s.registry.All()
	if len(all) == 0 {
		return nil, groupPolicy{}, fmt.Errorf("no nodes connected")
	}

	var candidates []candidateNode
	for _, entry := range all {
		cap := defaultNodeCapacity
		if store != nil {
			rec, ok := store.FindByID(entry.NodeID)
			if ok {
				res := resolveAdapterForNode(rec, req.Adapter, req.Target)
				if !res.supported {
					continue
				}
				cap = res.capacity
			}
		}
		candidates = append(candidates, candidateNode{entry: entry, capacity: cap})
	}
	if len(candidates) == 0 {
		return nil, groupPolicy{}, fmt.Errorf("no nodes support adapter %q target %q", req.Adapter, req.Target)
	}

	entries := make([]*edgenode.NodeEntry, len(candidates))
	for i, c := range candidates {
		entries[i] = c.entry
	}
	policy := groupPolicyFromRequestOrStore(req, store, entries)
	return candidates, policy, nil
}

// adapterResolution holds the resolved capacity and queue policy for a single
// node/adapter combination. ambiguous is true when a type-name lookup (e.g.,
// "ollama") matches 2+ enabled instances on that node, mirroring the Node
// router's exact-instance-key/ambiguity contract.
type adapterResolution struct {
	supported      bool
	ambiguous      bool
	capacity       int
	maxQueue       int
	queueTimeoutMS int
}

func positiveOr(v, fallback int) int {
	if v > 0 {
		return v
	}
	return fallback
}

// resolveAdapterForNode determines whether a node can handle adapterType/target
// and computes the per-node capacity and queue policy fields.
//
// Resolution order:
//  1. Exact instance Name match across OllamaInstances / VllmInstances /
//     OpenAICompatInstances (instance-key route, matching Node router priority).
//  2. Type-name route (e.g. "ollama"): supported only when exactly 1 enabled
//     instance of that type exists. 2+ enabled instances → ambiguous (fail,
//     same semantics as Node router ambiguity error). 0 instances → fail-open
//     for legacy/unconfigured nodes.
//  3. "cli": capability gated by CLI.Enabled and profile name in target.
//  4. Default (unknown adapter type): fail-open.
func resolveAdapterForNode(rec *edgenode.NodeRecord, adapterType, target string) adapterResolution {
	if rec == nil {
		return adapterResolution{supported: true, capacity: defaultNodeCapacity}
	}
	concurrencyFallback := positiveOr(rec.Runtime.Concurrency, defaultNodeCapacity)

	// Exact instance Name match (highest priority).
	for _, inst := range rec.Adapters.OllamaInstances {
		if inst.Name == adapterType {
			return adapterResolution{
				supported:      inst.Enabled,
				capacity:       positiveOr(inst.Capacity, concurrencyFallback),
				maxQueue:       inst.MaxQueue,
				queueTimeoutMS: inst.QueueTimeoutMS,
			}
		}
	}
	for _, inst := range rec.Adapters.VllmInstances {
		if inst.Name == adapterType {
			return adapterResolution{
				supported:      inst.Enabled,
				capacity:       positiveOr(inst.Capacity, concurrencyFallback),
				maxQueue:       inst.MaxQueue,
				queueTimeoutMS: inst.QueueTimeoutMS,
			}
		}
	}
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if inst.Name == adapterType {
			return adapterResolution{
				supported:      inst.Enabled,
				capacity:       positiveOr(inst.Capacity, concurrencyFallback),
				maxQueue:       inst.MaxQueue,
				queueTimeoutMS: inst.QueueTimeoutMS,
			}
		}
	}

	// Type-name route.
	switch adapterType {
	case "ollama":
		return resolveTypeRoute(rec, ollamaEnabledInstances(rec), concurrencyFallback)
	case "vllm":
		return resolveTypeRoute(rec, vllmEnabledInstances(rec), concurrencyFallback)
	case "openai_compat":
		return resolveTypeRoute(rec, openAICompatEnabledInstances(rec), concurrencyFallback)
	case "cli":
		// CLI has no named multi-instance model; capability is gated by profile.
		if !rec.Adapters.CLI.Enabled && len(rec.Adapters.CLI.Profiles) == 0 {
			// No CLI config at all → fail-open for legacy/unconfigured nodes.
			return adapterResolution{supported: true, capacity: concurrencyFallback}
		}
		if !rec.Adapters.CLI.Enabled {
			return adapterResolution{supported: false}
		}
		if target == "" {
			return adapterResolution{supported: len(rec.Adapters.CLI.Profiles) > 0, capacity: concurrencyFallback}
		}
		_, ok := rec.Adapters.CLI.Profiles[target]
		return adapterResolution{supported: ok, capacity: concurrencyFallback}
	default:
		return adapterResolution{supported: true, capacity: concurrencyFallback}
	}
}

// instanceFields holds the capacity/policy fields extracted from a single adapter instance.
type instanceFields struct {
	capacity, maxQueue, queueTimeoutMS int
}

func ollamaEnabledInstances(rec *edgenode.NodeRecord) []instanceFields {
	var out []instanceFields
	for _, inst := range rec.Adapters.OllamaInstances {
		if inst.Enabled {
			out = append(out, instanceFields{inst.Capacity, inst.MaxQueue, inst.QueueTimeoutMS})
		}
	}
	return out
}

func vllmEnabledInstances(rec *edgenode.NodeRecord) []instanceFields {
	var out []instanceFields
	for _, inst := range rec.Adapters.VllmInstances {
		if inst.Enabled {
			out = append(out, instanceFields{inst.Capacity, inst.MaxQueue, inst.QueueTimeoutMS})
		}
	}
	return out
}

func openAICompatEnabledInstances(rec *edgenode.NodeRecord) []instanceFields {
	var out []instanceFields
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if inst.Enabled {
			out = append(out, instanceFields{inst.Capacity, inst.MaxQueue, inst.QueueTimeoutMS})
		}
	}
	return out
}

// resolveTypeRoute applies Node-router-compatible type-name resolution:
//   - 0 enabled instances → fail-open (legacy / unconfigured)
//   - 1 enabled instance → use its capacity and policy
//   - 2+ enabled instances → ambiguous (reject)
func resolveTypeRoute(rec *edgenode.NodeRecord, enabled []instanceFields, concurrencyFallback int) adapterResolution {
	switch len(enabled) {
	case 0:
		return adapterResolution{supported: true, capacity: concurrencyFallback}
	case 1:
		f := enabled[0]
		return adapterResolution{
			supported:      true,
			capacity:       positiveOr(f.capacity, concurrencyFallback),
			maxQueue:       f.maxQueue,
			queueTimeoutMS: f.queueTimeoutMS,
		}
	default:
		return adapterResolution{supported: false, ambiguous: true}
	}
}

func groupPolicyFromRequestOrStore(req SubmitRunRequest, store *edgenode.NodeStore, entries []*edgenode.NodeEntry) groupPolicy {
	if req.MaxQueue > 0 || req.QueueTimeoutMS > 0 {
		maxQueue := req.MaxQueue
		if maxQueue <= 0 {
			maxQueue = defaultGroupMaxQueue
		}
		return groupPolicy{
			maxQueue:        maxQueue,
			queueTimeout:    time.Duration(req.QueueTimeoutMS) * time.Millisecond,
			queueTimeoutSet: true,
		}
	}
	return groupPolicyFromStore(store, entries, req.Adapter, req.Target)
}

// groupPolicyFromStore derives queue policy from the first resolved candidate node.
func groupPolicyFromStore(store *edgenode.NodeStore, entries []*edgenode.NodeEntry, adapterType, target string) groupPolicy {
	if store != nil {
		for _, e := range entries {
			rec, ok := store.FindByID(e.NodeID)
			if !ok {
				continue
			}
			res := resolveAdapterForNode(rec, adapterType, target)
			if !res.supported {
				continue
			}
			if res.maxQueue > 0 || res.queueTimeoutMS > 0 {
				p := groupPolicy{
					maxQueue:        positiveOr(res.maxQueue, defaultGroupMaxQueue),
					queueTimeout:    time.Duration(res.queueTimeoutMS) * time.Millisecond,
					queueTimeoutSet: true,
				}
				return p
			}
		}
	}
	return groupPolicy{maxQueue: defaultGroupMaxQueue, queueTimeout: defaultQueueTimeout, queueTimeoutSet: true}
}

// providerCanServe checks whether the provider advertises the served model in
// its own models list (defensive SDD compliance).
func providerCanServe(prov config.NodeProviderConf, servedModel string) bool {
	for _, m := range prov.Models {
		if m == servedModel {
			return true
		}
	}
	return false
}

// providerAdapterKey returns the dispatch adapter key for a provider.
// For legacy/compat providers the explicit Adapter field is used; for
// provider-first providers (Adapter is empty) the provider ID is used.
func providerAdapterKey(prov config.NodeProviderConf) string {
	if k := strings.TrimSpace(prov.Adapter); k != "" {
		return k
	}
	return prov.ID
}

// applyProviderDispatchFields copies the dispatch inputs a candidate derives from
// provider config onto c. Both the initial provider-pool resolution and the
// scheduler's re-resolution of an already-queued candidate go through it, so a
// request that waited across a config refresh is dispatched under exactly the
// same adapter/priority/execution-path rules as one admitted immediately. The
// candidate's identity fields (node entry, provider id, served target) are the
// caller's request and are deliberately left untouched.
func applyProviderDispatchFields(c *candidateNode, prov config.NodeProviderConf) {
	c.capacity = prov.Capacity
	c.longContextCapacity = prov.LongContextCapacity
	c.priority = prov.Priority
	c.providerType = prov.Type
	c.adapter = providerAdapterKey(prov)
	c.executionPath = classifyProviderExecutionPath(prov.Type)
}

// isProviderAvailable checks provider health status. Only "available" (and
// optionally "healthy" as an alias) are considered dispatchable.
func isProviderAvailable(health string) bool {
	h := strings.ToLower(strings.TrimSpace(health))
	return h == "available" || h == "healthy"
}

// isProviderAdapterInstanceValid is a defensive check in provider-pool candidate
// resolution. It returns false only when the adapter name resolves to a disabled
// exact instance, an ambiguous type route (2+ enabled instances of that type),
// a type route with zero enabled instances, or an unknown/missing adapter key.
// Only enabled exact instances and single-enabled type routes are considered valid.
func isProviderAdapterInstanceValid(rec *edgenode.NodeRecord, adapter string) bool {
	if rec == nil {
		return true
	}
	for _, inst := range rec.Adapters.OllamaInstances {
		if inst.Name == adapter {
			return inst.Enabled
		}
	}
	for _, inst := range rec.Adapters.VllmInstances {
		if inst.Name == adapter {
			return inst.Enabled
		}
	}
	for _, inst := range rec.Adapters.OpenAICompatInstances {
		if inst.Name == adapter {
			return inst.Enabled
		}
	}
	if adapter == "cli" {
		return rec.Adapters.CLI.Enabled
	}
	switch adapter {
	case "ollama":
		// Match edgevalidate.buildAdapterIndex: legacy ollama.Enabled counts as one enabled instance/key.
		count := 0
		if rec.Adapters.Ollama.Enabled {
			count++
		}
		count += len(ollamaEnabledInstances(rec))
		return count == 1
	case "vllm":
		count := 0
		if rec.Adapters.Vllm.Enabled {
			count++
		}
		count += len(vllmEnabledInstances(rec))
		return count == 1
	case "openai_compat":
		count := 0
		if rec.Adapters.OpenAICompat.Enabled {
			count++
		}
		count += len(openAICompatEnabledInstances(rec))
		return count == 1
	}
	return false // unknown/custom adapter key: excluded from provider-pool candidates
}

// classifyProviderExecutionPath classifies a provider's execution path based on
// its type. OpenAI-compatible aliases (openai_compat, openai_api, vllm, vllm-mlx,
// lemonade, sglang, seulgivibe_claude, seulgivibe_openai) are routed to the
// tunnel/passthrough path. Ollama, CLI, and unknown/native types use the
// normalized path. This mirrors the SDD requirement that OpenAI-compatible
// callers go through passthrough while Ollama/CLI/native use normalized.
func classifyProviderExecutionPath(providerType string) providerExecutionPath {
	switch strings.ToLower(strings.TrimSpace(providerType)) {
	case "openai_compat", "openai_api", "vllm", "vllm-mlx", "lemonade", "sglang",
		"seulgivibe_claude", "seulgivibe_openai":
		return providerExecutionPathTunnel
	default:
		// ollama, cli, and any unknown/native type → normalized.
		return providerExecutionPathNormalized
	}
}

// resolveProviderPoolCandidates builds candidates for a provider-pool request.
// It scans connected nodes for providers referenced in the model catalog entry
// that matches req.ModelGroupKey, then assembles per-provider candidateNodes
// carrying the concrete served model name for target rewrite after admission.
// Filters: served-model membership, dispatch adapter presence, and available health.
func (s *Service) resolveProviderPoolCandidates(req SubmitRunRequest, store *edgenode.NodeStore, catalog []config.ModelCatalogEntry) ([]candidateNode, groupPolicy, error) {
	var catalogEntry *config.ModelCatalogEntry
	for i := range catalog {
		if catalog[i].ID == req.ModelGroupKey {
			catalogEntry = &catalog[i]
			break
		}
	}
	if catalogEntry == nil {
		return nil, groupPolicy{}, fmt.Errorf("provider pool model %q not found in catalog", req.ModelGroupKey)
	}

	all := s.registry.All()
	if len(all) == 0 {
		return nil, groupPolicy{}, fmt.Errorf("no nodes connected")
	}

	candidates := make([]candidateNode, 0)
	for _, entry := range all {
		if store == nil {
			continue
		}
		rec, ok := store.FindByID(entry.NodeID)
		if !ok {
			continue
		}
		for _, prov := range rec.Providers {
			servedModel, inCatalog := catalogEntry.Providers[prov.ID]
			if !inCatalog {
				continue
			}
			// Exclude disabled providers from dispatch.
			if !config.ProviderEnabled(prov) {
				continue
			}
			// Defensive SDD compliance: served target must be in provider's own models list.
			if !providerCanServe(prov, servedModel) {
				continue
			}
			// Derive dispatch adapter key: explicit adapter wins; provider-first uses provider ID.
			adapterKey := providerAdapterKey(prov)
			if strings.TrimSpace(prov.Adapter) != "" {
				// Legacy/compat: adapter must resolve to an enabled instance on this node.
				if !isProviderAdapterInstanceValid(rec, adapterKey) {
					continue
				}
			}
			// Only available/healthy providers are dispatchable.
			if !isProviderAvailable(prov.Health) {
				continue
			}
			// SDD compliance: capacity 0 or unknown providers are excluded from
			// dispatch candidates. Do NOT fall back to runtime/default concurrency.
			if prov.Capacity <= 0 {
				continue
			}
			candidate := candidateNode{
				entry:        entry,
				providerID:   prov.ID,
				servedTarget: servedModel,
			}
			applyProviderDispatchFields(&candidate, prov)
			candidates = append(candidates, candidate)
		}
	}

	if len(candidates) == 0 {
		return nil, groupPolicy{}, fmt.Errorf("no connected nodes support provider pool model %q", req.ModelGroupKey)
	}
	// Policy is always zero for provider-pool resolution: the atomic runtime
	// snapshot owns the canonical root policy and callers use it directly.
	return candidates, groupPolicy{}, nil
}