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 (
	"context"
	"sort"
	"strings"
	"sync"
	"time"
)

// 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")
}

// ProviderForModel returns the provider owning the model id (nil if unknown).
func (r *Registry) ProviderForModel(model string) *Provider {
	r.mu.RLock()
	defer r.mu.RUnlock()
	p, ok := r.byModel[strings.ToLower(strings.TrimSpace(model))]
	if !ok {
		return nil
	}
	return p
}

// ProviderForSlot returns the provider for a (model, source) slot. When
// source is empty it behaves like ProviderForModel (owner of the model id);
// when source is set it returns only that exact source (nil if the source
// does not serve the model).
func (r *Registry) ProviderForSlot(model, source string) *Provider {
	model = strings.ToLower(strings.TrimSpace(model))
	source = strings.TrimSpace(source)
	r.mu.RLock()
	defer r.mu.RUnlock()
	if source == "" {
		p, ok := r.byModel[model]
		if !ok {
			return nil
		}
		return p
	}
	for _, p := range r.providers {
		if !strings.EqualFold(p.Name(), source) {
			continue
		}
		for _, m := range p.cfg.Models {
			if strings.EqualFold(m.ID, model) {
				return p
			}
		}
		return nil
	}
	return nil
}

// 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"`
	LiveAvailable bool     `json:"live_available"`
	LastError     string   `json:"last_error,omitempty"`
	LastChecked   int64    `json:"last_checked,omitempty"`
}

// ProbeAll runs a live reachability check for every provider (in parallel).
func (r *Registry) ProbeAll(ctx context.Context) {
	r.mu.RLock()
	providers := append([]*Provider(nil), r.providers...)
	r.mu.RUnlock()
	var wg sync.WaitGroup
	for _, p := range providers {
		wg.Add(1)
		go func(p *Provider) {
			defer wg.Done()
			probeCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
			defer cancel()
			p.Probe(probeCtx)
		}(p)
	}
	wg.Wait()
}

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