ModelRouter/internal/provider/registry.go

// Package provider also provides the top-level registry that owns all sources,
// routes model requests (explicit or AUTO), and supports hot reload.
package provider

import (
	"sort"
	"strings"
	"sync"
)

// Registry holds all configured providers and routes model requests.
type Registry struct {
	mu        sync.RWMutex
	providers []*Provider
	byModel   map[string]*Provider // modelID -> provider
	defaultM  string               // default model id ("" means AUTO)
}

func NewRegistry(providers []*Provider, defaultModel string) *Registry {
	r := &Registry{byModel: map[string]*Provider{}, defaultM: defaultModel}
	r.set(providers)
	return r
}

// Replace atomically swaps the provider set (hot reload).
func (r *Registry) Replace(providers []*Provider) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.set(providers)
}

func (r *Registry) set(providers []*Provider) {
	r.providers = providers
	r.byModel = map[string]*Provider{}
	for _, p := range providers {
		for _, m := range p.Models() {
			r.byModel[strings.ToLower(m)] = p
		}
	}
}

func (r *Registry) Providers() []*Provider {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]*Provider, len(r.providers))
	copy(out, r.providers)
	return out
}

func (r *Registry) Count() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.providers)
}

// ModelList returns all exposed model ids (chat + image).
func (r *Registry) ModelList() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	seen := map[string]bool{}
	for _, p := range r.providers {
		for _, m := range p.Models() {
			if !seen[m] {
				seen[m] = true
			}
		}
	}
	out := make([]string, 0, len(seen))
	for m := range seen {
		out = append(out, m)
	}
	sort.Strings(out)
	return out
}

// Resolve returns the ordered candidate providers to try for a request,
// honoring explicit model selection or AUTO (priority order, healthy first).
//
// model "" or "AUTO" -> all sources sorted by (priority desc, healthy first).
// Otherwise the owning provider, if healthy; else its source anyway.
func (r *Registry) Resolve(model string) []*Provider {
	r.mu.RLock()
	defer r.mu.RUnlock()

	model = strings.TrimSpace(model)
	if model == "" || strings.EqualFold(model, "AUTO") {
		// priority chain across all models
		type cand struct {
			prov     *Provider
			priority int
		}
		var cands []cand
		seen := map[string]bool{}
		for _, p := range r.providers {
			prio := -1
			for _, m := range p.cfg.Models {
				if m.Priority > prio {
					prio = m.Priority
				}
			}
			if prio < 0 {
				prio = 0
			}
			cands = append(cands, cand{p, prio})
			seen[p.Name()] = true
		}
		sort.SliceStable(cands, func(i, j int) bool {
			if cands[i].priority != cands[j].priority {
				return cands[i].priority > cands[j].priority
			}
			// healthy preferred at same priority
			return cands[i].prov.Available() && !cands[j].prov.Available()
		})
		out := make([]*Provider, 0, len(cands))
		for _, c := range cands {
			out = append(out, c.prov)
		}
		return out
	}

	// explicit model
	if p, ok := r.byModel[strings.ToLower(model)]; ok {
		// switch to the owning source but pin the model via request
		return []*Provider{p}
	}
	// unknown model -> fall back to default/AUTO chain
	return r.AUTOChain()
}

// AUTOChain returns the priority-sorted providers for AUTO.
func (r *Registry) AUTOChain() []*Provider {
	return r.Resolve("AUTO")
}

// Default returns the highest-priority available provider.
func (r *Registry) Default() *Provider {
	chain := r.AUTOChain()
	if len(chain) == 0 {
		return nil
	}
	return chain[0]
}

// ModelStatus is a web-UI friendly snapshot per source.
type SourceStatus struct {
	Name          string   `json:"name"`
	Adapter       string   `json:"adapter"`
	BaseURL       string   `json:"base_url"`
	Models        []string `json:"models"`
	Available     bool     `json:"available"`
	Healthy       bool     `json:"healthy"`
	MaxConcurrent int      `json:"max_concurrent"`
}

func (r *Registry) Status() []SourceStatus {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]SourceStatus, 0, len(r.providers))
	for _, p := range r.providers {
		s := SourceStatus{
			Name:          p.Name(),
			Adapter:       p.Adapter(),
			BaseURL:       p.Config().BaseURL,
			Models:        p.Models(),
			Available:     p.Available(),
			Healthy:       p.Available(),
			MaxConcurrent: p.MaxConcurrent(),
		}
		out = append(out, s)
	}
	return out
}