ModelRouter/internal/core/core.go

// Package core wires together the Lua VM, provider registry, scheduler and
// runtime store and exposes management operations (hot reload, adapters,
// sources) for the web UI and gateway.
package core

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sort"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"llmsproxy/internal/billing"
	"llmsproxy/internal/config"
	"llmsproxy/internal/lua"
	"llmsproxy/internal/provider"
	"llmsproxy/internal/scheduler"
)

// Core owns the running configuration and adapters. mu guards every read and
// write of c.cfg (keys / sources / auto rules) from the management API while
// request paths look keys up concurrently; the AUTO chain itself is swapped
// atomically and needs no lock.
type Core struct {
	mu             sync.Mutex
	cfg            *config.Config
	vm             *lua.VM
	plugins        *lua.Plugins
	store          *config.Store
	scheduler      *scheduler.Scheduler
	registry       *provider.Registry
	autoChain      atomic.Pointer[scheduler.Chain] // chat AUTO chain
	autoImageChain atomic.Pointer[scheduler.Chain] // image-generation AUTO chain

	// keyAutoChains caches per-key AUTO chains, keyed by the gateway key.
	// Entries are rebuilt by rebuildKeyAutoChains whenever sources or key
	// config change; AutoChainFor reads them without taking c.mu so request
	// handling stays lock-free like the global chain lookup.
	keyAutoChains atomic.Pointer[map[string]*scheduler.Chain]
}

// New builds the core from a config file plus runtime overlay.
func New(cfgPath string) (*Core, error) {
	cfg, err := config.Load(cfgPath)
	if err != nil {
		return nil, err
	}
	return NewFromConfig(cfg)
}

// NewFromConfig builds the core from an already-loaded config.
func NewFromConfig(cfg *config.Config) (*Core, error) {
	c := &Core{cfg: cfg}
	c.vm = lua.NewVM(cfg.AdapterDir)
	if err := c.vm.Start(); err != nil {
		return nil, fmt.Errorf("lua vm: %w", err)
	}
	// Plugins load AFTER the VM is up (a plugin is just another Lua state) but
	// BEFORE anything can serve traffic, so a plugin's request_end hook sees
	// the very first request. A plugin that fails to load is logged and skipped
	// rather than failing startup: plugins are optional extensions, and failing
	// to boot the gateway because a third-party .lua has a typo would be the
	// wrong trade.
	c.plugins = lua.NewPlugins(c.vm, cfg.PluginDir)
	// Plugin outbound HTTP is scoped to the hosts of configured sources. Set
	// before LoadDir so a plugin that requests on its first hook already sees
	// the real allowlist rather than the deny-all placeholder.
	c.plugins.SetOutboundHosts(sourceHosts(cfg))
	// Let the plugin layer validate access specs against the keys that exist.
	// A stored allowlist naming a key that does not exist refuses every caller,
	// which looks exactly like a broken endpoint.
	lua.SetKeyNameLookup(func() []string {
		// ListKeys is the structured view (bare `gateway_keys:` entries are
		// expanded into records with names); cfg.GatewayKeys is still []string.
		keys := c.ListKeys()
		out := make([]string, 0, len(keys))
		for _, k := range keys {
			if k.Name != "" {
				out = append(out, k.Name)
			}
		}
		return out
	})
	// Operator overrides win over the plugin's own declared audience.
	for name, aud := range cfg.PluginVisibility {
		c.plugins.SetPluginVisibility(name, aud)
	}
	if err := c.plugins.SeedBundled(); err != nil {
		log.Printf("[core] seed bundled plugins: %v", err)
	}
	if err := c.plugins.LoadDir(); err != nil {
		log.Printf("[core] plugin dir: %v", err)
	}
	for _, p := range c.plugins.List() {
		if msg, bad := p["error"]; bad {
			log.Printf("[core] plugin %v failed to load: %v", p["name"], msg)
			continue
		}
		log.Printf("[core] plugin %v %v loaded (hooks=%v ui=%v)",
			p["name"], p["version"], p["hooks"], p["ui"])
	}
	c.applyBillingDSL()
	c.store = config.NewStore(cfg.RuntimeFile)
	// Share one box between the runtime store and config.yaml so a single
	// master.key seals both files. config.Load left the config holding
	// ciphertext (if it was sealed); unseal it NOW, before anything reads a
	// credential.
	//
	// ORDER IS LOAD-BEARING. These steps each consume secrets and must run
	// after the unseal:
	//   - seedKeys compares cfg.Keys[i].Key against the plaintext
	//     gateway_keys entries; with ciphertext keys the comparison never
	//     matched and every restart appended another duplicate admin key
	//     (production had four copies of the same admin key).
	//   - rebuildRegistry hands cfg.Sources[i].APIKey to the providers; with
	//     ciphertext it built every upstream client with "enc:v1:..." as its
	//     bearer token.
	// Previously the unseal happened at the END of this function, and things
	// only appeared to work because seedKeys' Save() unsealed memory as a side
	// effect. Removing the redundant saves exposed the real ordering bug.
	cfg.AttachSecretBox(c.store.SecretBox())
	hadPlaintextSecrets, err := cfg.UnsealSecrets(c.store.SecretBox())
	if err != nil {
		return nil, fmt.Errorf("unseal secrets: %w", err)
	}
	if err := c.store.Load(); err != nil {
		return nil, fmt.Errorf("runtime store: %w", err)
	}
	c.scheduler = scheduler.New(buildRetries(cfg))
	// One-time migration: lift auto rules and keys from legacy runtime.json
	// into config.yaml so all configuration lives in one place.
	c.migrateFromRuntime()
	if err := c.seedKeys(); err != nil {
		return nil, err
	}
	if err := c.seedAuto(); err != nil {
		return nil, err
	}
	if err := c.rebuildRegistry(); err != nil {
		return nil, err
	}
	// Seed built-in preset templates on first run so the WebUI shows
	// "从模板创建" for popular providers out of the box.
	if err := c.seedPresetTemplates(); err != nil {
		return nil, fmt.Errorf("seed preset templates: %w", err)
	}
	// Seal any credential still in the clear in config.yaml. Startup writes
	// nothing when the file was already sealed (hadPlaintextSecrets is decided
	// from the on-disk state, before the unseal above).
	if err := cfg.SealIfNeeded(hadPlaintextSecrets); err != nil {
		return nil, fmt.Errorf("normalize secrets: %w", err)
	}
	return c, nil
}

// migrateFromRuntime lifts auto rules and keys from legacy runtime.json into
// c.cfg — but only if they are not already present in the YAML config. This
// lets people upgrade without losing their data; once migrated, the YAML file
// is authoritative and runtime.json's auto/keys are ignored.
func (c *Core) migrateFromRuntime() {
	legacy := c.store.LoadLegacy()
	if legacy == nil {
		return
	}
	changed := false
	if len(legacy.Auto) > 0 && len(c.cfg.Auto) == 0 {
		c.cfg.Auto = legacy.Auto
		changed = true
	}
	if len(legacy.Keys) > 0 && len(c.cfg.Keys) == 0 {
		c.cfg.Keys = legacy.Keys
		changed = true
	}
	if changed {
		if err := c.cfg.Save(); err != nil {
			fmt.Printf("[core] migrate to config.yaml: %v\n", err)
		}
	}
}

// seedAuto uses the existing auto rules from config.yaml, or creates an
// initial chain from legacy per-source model priority (first run only).
func (c *Core) seedAuto() error {
	if len(c.cfg.Auto) > 0 {
		return nil
	}
	type item struct {
		model string
		prio  int
	}
	var flat []item
	for _, s := range c.mergedSources() {
		for _, m := range s.Models {
			if m.Kind == "image" {
				continue
			}
			flat = append(flat, item{m.ID, m.Priority})
		}
	}
	sort.SliceStable(flat, func(i, j int) bool {
		if flat[i].prio != flat[j].prio {
			return flat[i].prio > flat[j].prio
		}
		return flat[i].model < flat[j].model
	})
	entries := make([]config.ModelScope, 0, len(flat))
	for _, it := range flat {
		entries = append(entries, config.ModelScope{Model: it.model})
	}
	c.cfg.Auto = entries
	return c.cfg.Save()
}

// seedKeys ensures config.yaml has admin keys for every gateway_keys entry
// from the YAML config (once). After first save they become GWKey records
// in c.cfg.Keys and gateway_keys is no longer read for auth.
func (c *Core) seedKeys() error {
	existing := map[string]bool{}
	for _, k := range c.cfg.Keys {
		existing[k.Key] = true
	}
	changed := false
	for i, raw := range c.cfg.GatewayKeys {
		if raw == "" || existing[raw] {
			continue
		}
		name := "admin"
		if i > 0 {
			name = fmt.Sprintf("admin-%d", i+1)
		}
		c.cfg.Keys = append(c.cfg.Keys, config.GWKey{
			Key:       raw,
			Role:      "admin",
			Name:      name,
			CreatedAt: time.Now().Unix(),
			Seed:      true,
		})
		changed = true
	}
	if changed {
		return c.cfg.Save()
	}
	return nil
}

// seedPresetTemplates adds the built-in provider presets to the template
// store on first run. It only adds templates whose name is not already
// present, so a user's edits to a preset survive restarts and a deliberately
// deleted preset does not silently return on the next start (its name is
// remembered in the store's preset marker list).
func (c *Core) seedPresetTemplates() error {
	if len(config.PresetTemplates) == 0 {
		return nil
	}
	return c.store.SeedTemplates(config.PresetTemplates)
}

// saveConfig writes the current config (including auto rules and keys) back
// to config.yaml.
func (c *Core) saveConfig() error {
	return c.cfg.Save()
}

func buildRetries(cfg *config.Config) int {
	return len(cfg.Sources)
}

// VM exposes the Lua adapter runtime.
func (c *Core) VM() *lua.VM { return c.vm }

// Plugins exposes the loaded plugin set. Never nil once New* has returned, so
// the gateway can fire stages unconditionally; an unconfigured plugin
// directory yields an empty registry whose Fire is a no-op.
func (c *Core) Plugins() *lua.Plugins { return c.plugins }

func (c *Core) Scheduler() *scheduler.Scheduler { return c.scheduler }

func (c *Core) Registry() *provider.Registry { return c.registry }

// AutoChain returns the current AUTO scheduling chain.
func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.Load() }

// AutoChainFor returns the AUTO chain to use for a request authenticated with
// key. A key that declares its own chain gets that chain; every other key
// inherits the global chain, so keys created before per-key chains existed
// keep their current behaviour with no configuration change.
//
// It returns (chain, ok). ok is false only when the key has its own chain but
// it compiled to nothing usable (e.g. every slot names a model that no longer
// exists) — the caller must surface that rather than silently downgrading the
// user to the global chain, which would be a confusing, invisible switch.
func (c *Core) AutoChainFor(key string) (*scheduler.Chain, bool) {
	if m := c.keyAutoChains.Load(); m != nil {
		if ch, ok := (*m)[key]; ok {
			return ch, true
		}
	}
	return c.autoChain.Load(), true
}

// rebuildKeyAutoChains recompiles every per-key AUTO chain from the current
// config and provider registry, then swaps the cache in one shot so requests
// never observe a half-built map. Keys without their own chain are absent
// from the map and fall back to the global chain.
func (c *Core) rebuildKeyAutoChains() {
	m := make(map[string]*scheduler.Chain)
	for _, k := range c.cfg.Keys {
		if !k.HasOwnAuto() {
			continue
		}
		m[k.Key] = c.chainForRules(k.Auto, false)
	}
	c.keyAutoChains.Store(&m)
}

// SaveKeyAuto persists one key's own AUTO chain and recompiles the per-key
// chain cache. Passing an empty list clears the override so the key inherits
// the global chain again.
func (c *Core) SaveKeyAuto(key string, entries []config.ModelScope) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	idx := -1
	for i, k := range c.cfg.Keys {
		if k.Key == key {
			idx = i
			break
		}
	}
	if idx < 0 {
		return fmt.Errorf("key not found")
	}
	if err := ValidateScopeQuotas(entries); err != nil {
		return err
	}
	c.cfg.Keys[idx].Auto = cleanScopes(entries)
	if err := c.saveConfig(); err != nil {
		return err
	}
	// Rebuild the whole cache so a cleared override stops shadowing the
	// global chain for this key.
	c.rebuildKeyAutoChains()
	return nil
}

// AutoImageChain returns the persisted image-generation AUTO chain snapshot.
func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() }

func (c *Core) DefaultModel() string { return c.cfg.DefaultModel }

func (c *Core) GatewayKeys() []string { return c.cfg.GatewayKeys }

func (c *Core) Listen() string { return c.cfg.Listen }

func (c *Core) TLS() (cert, key string) {
	return c.cfg.TLSCertFile, c.cfg.TLSKeyFile
}

func (c *Core) PublicBaseURL() string { return c.cfg.PublicBaseURL }

// ---- gateway key management (web UI) ----

// ListKeys returns all gateway keys (admin view).
func (c *Core) ListKeys() []config.GWKey {
	c.mu.Lock()
	defer c.mu.Unlock()
	out := make([]config.GWKey, len(c.cfg.Keys))
	copy(out, c.cfg.Keys)
	return out
}

// FindKey looks up a gateway key record by its secret value.
func (c *Core) FindKey(key string) (config.GWKey, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, k := range c.cfg.Keys {
		if k.Key == key {
			return k, true
		}
	}
	return config.GWKey{}, false
}

// CreateKey builds a new random gateway key and persists it to config.yaml.
// Quotas live on the model scope entries, so a new key's budget is whatever
// its scopes carry.
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	models = cleanScopes(models)
	if err := ValidateScopeQuotas(models); err != nil {
		return config.GWKey{}, err
	}
	key := make([]byte, 16)
	if _, err := rand.Read(key); err != nil {
		return config.GWKey{}, err
	}
	rec := config.GWKey{
		Key:       "sk-gw-" + hex.EncodeToString(key),
		Role:      role,
		Name:      name,
		Models:    models,
		Note:      note,
		CreatedAt: time.Now().Unix(),
	}
	rec.Role = config.NormalizeRole(rec.Role)
	c.cfg.Keys = append(c.cfg.Keys, rec)
	if err := c.saveConfig(); err != nil {
		return config.GWKey{}, err
	}
	return rec, nil
}

// UpdateKey mutates a key's name/role/model scope and persists it. The scope
// entries carry their own quotas, so replacing the scope replaces the budgets.
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if models != nil {
		if err := ValidateScopeQuotas(models); err != nil {
			return config.GWKey{}, err
		}
	}
	for i, k := range c.cfg.Keys {
		if k.Key == key {
			if name != "" {
				c.cfg.Keys[i].Name = name
			}
			if role == "admin" || role == "user" {
				c.cfg.Keys[i].Role = role
			}
			// models == nil means the caller did not provide a scope (leave
			// the existing one untouched); an explicit [] clears it.
			if models != nil {
				c.cfg.Keys[i].Models = cleanScopes(models)
			}
			c.cfg.Keys[i].Note = note
			if err := c.saveConfig(); err != nil {
				return config.GWKey{}, err
			}
			return c.cfg.Keys[i], nil
		}
	}
	return config.GWKey{}, fmt.Errorf("key not found")
}

// DeleteKey removes a key record; returns false if it did not exist.
func (c *Core) DeleteKey(key string) (bool, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	for i, k := range c.cfg.Keys {
		if k.Key == key {
			c.cfg.Keys = append(c.cfg.Keys[:i], c.cfg.Keys[i+1:]...)
			return true, c.saveConfig()
		}
	}
	return false, nil
}

// ---- AUTO scheduling slots (web UI) ----

// AutoRules returns the AUTO scheduling slots in priority order.
func (c *Core) AutoRules() []config.ModelScope {
	c.mu.Lock()
	defer c.mu.Unlock()
	out := make([]config.ModelScope, len(c.cfg.Auto))
	copy(out, c.cfg.Auto)
	return out
}

// AutoImageRules returns the configured image-generation AUTO chain rules.
func (c *Core) AutoImageRules() []config.ModelScope {
	c.mu.Lock()
	defer c.mu.Unlock()
	out := make([]config.ModelScope, len(c.cfg.AutoImage))
	copy(out, c.cfg.AutoImage)
	return out
}

// cleanScopes drops empty model entries and normalizes placeholder source
// names; it returns nil when no entries survive so an empty scope means
// "unrestricted" (nil) instead of a restrictive-but-empty list — a non-nil
// empty slice would 403 every model in-process yet become unrestricted again
// after a restart (config omits empty models with omitempty).
func cleanScopes(entries []config.ModelScope) []config.ModelScope {
	var clean []config.ModelScope
	for _, e := range entries {
		if e.Model == "" {
			continue
		}
		if e.Source == "undefined" || e.Source == "null" {
			e.Source = ""
		}
		clean = append(clean, e)
	}
	return clean
}

// SaveAutoRules persists the AUTO scheduling slots to config.yaml, rebuilds
// the chain and clears the cooldown of every slot — preference scores are
// kept, so a reliably good model keeps its edge while an edited chain applies
// immediately.
func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.cfg.Auto = cleanScopes(entries)
	if err := c.saveConfig(); err != nil {
		return err
	}
	c.buildAutoChain()
	c.resetChainCooldowns(c.autoChain.Load())
	return nil
}

// SaveAutoImageRules persists the image-generation AUTO chain. Image models
// are kept as-is (unlike buildAutoChain, which skips them).
func (c *Core) SaveAutoImageRules(entries []config.ModelScope) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.cfg.AutoImage = cleanScopes(entries)
	if err := c.saveConfig(); err != nil {
		return err
	}
	c.buildAutoImageChain()
	c.resetChainCooldowns(c.autoImageChain.Load())
	return nil
}

func (c *Core) resetChainCooldowns(ch *scheduler.Chain) {
	if ch == nil {
		return
	}
	for _, tn := range ch.Tiers {
		for _, sl := range tn.Slots {
			if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
				p.ResetModelCooldown(sl.Model)
			}
		}
	}
}

// ResetHealth clears the scheduling backoff state of every provider.
func (c *Core) ResetHealth() {
	for _, p := range c.registry.Providers() {
		p.ResetHealth()
	}
}

func (c *Core) ProviderForModel(model string) *provider.Provider {
	return c.registry.ProviderForModel(model)
}

func (c *Core) ProviderForSlot(model, source string) *provider.Provider {
	return c.registry.ProviderForSlot(model, source)
}

func (c *Core) Config() *config.Config { return c.cfg }

// mergedSources = base YAML sources + runtime sources (runtime wins by name).
func (c *Core) mergedSources() []config.Source {
	byName := map[string]config.Source{}
	order := []string{}
	for _, s := range c.cfg.Sources {
		byName[s.Name] = s
		order = append(order, s.Name)
	}
	for _, s := range c.store.List() {
		if _, ok := byName[s.Name]; !ok {
			order = append(order, s.Name)
		}
		byName[s.Name] = s
	}
	out := make([]config.Source, 0, len(order))
	seen := map[string]bool{}
	for _, n := range order {
		if !seen[n] {
			seen[n] = true
			s := c.resolveSourceKey(byName[n])
			if s.Timeout == 0 {
				s.Timeout = config.DefaultSourceTimeout
			}
			if s.QueueTimeout == 0 {
				s.QueueTimeout = config.DefaultSourceQueueTimeout
			}
			if s.MaxConcurrent == 0 {
				s.MaxConcurrent = config.DefaultSourceConcurrency
			}
			out = append(out, s)
		}
	}
	return out
}

// resolveSourceKey applies api_key_env and decrypts enc:v1: ciphertext.
// Config values are already unsealed at startup (Config.NormalizeSecretsForRun),
// so the decrypt branch is the belt-and-braces path for a source that arrived
// already sealed through another route. A failure there leaves the ciphertext in
// place, which makes the upstream reject the key loudly rather than sending an
// empty Authorization header that might look like a config-less source.
func (c *Core) resolveSourceKey(s config.Source) config.Source {
	if s.APIKeyEnv != "" {
		if v := os.Getenv(s.APIKeyEnv); v != "" {
			s.APIKey = v
		}
		return s
	}
	if box := c.store.SecretBox(); box != nil && strings.HasPrefix(s.APIKey, "enc:v1:") {
		if v, err := box.Decrypt(s.APIKey); err == nil {
			s.APIKey = v
		} else {
			log.Printf("[core] source %s: api_key decrypt failed: %v", s.Name, err)
		}
	}
	return s
}

func (c *Core) rebuildRegistry() error {
	srcs := c.mergedSources()
	providers := make([]*provider.Provider, 0, len(srcs))
	adapterConcurrency := map[string]int{}
	for _, s := range srcs {
		providers = append(providers, provider.New(s, c.vm))
		if s.Adapter != "" {
			adapterConcurrency[s.Adapter] += s.MaxConcurrent
		}
	}
	c.vm.ConfigureConcurrency(adapterConcurrency)
	if c.registry == nil {
		c.registry = provider.NewRegistry(providers, c.cfg.DefaultModel)
	} else {
		c.registry.Replace(providers)
	}
	c.buildAutoChain()
	c.buildAutoImageChain()
	c.rebuildKeyAutoChains()
	// A source edit can add or remove an upstream host, which changes what a
	// plugin is allowed to reach. Keeping this beside the chain rebuilds means
	// one call site owns "config changed, refresh derived state".
	c.plugins.SetOutboundHosts(sourceHosts(c.cfg))
	return nil
}

// sourceHosts extracts the host of every configured source's base_url — the
// allowlist for plugin outbound HTTP. Unparseable URLs are skipped rather than
// failing startup: a malformed base_url is already rejected when the provider is
// built, and it must not take plugins down with it.
func sourceHosts(cfg *config.Config) []string {
	if cfg == nil {
		return nil
	}
	out := make([]string, 0, len(cfg.Sources))
	for _, s := range cfg.Sources {
		if s.BaseURL == "" {
			continue
		}
		u, err := url.Parse(s.BaseURL)
		if err != nil || u.Hostname() == "" {
			continue
		}
		out = append(out, u.Hostname())
	}
	return out
}

// buildAutoChain rebuilds the AUTO chain snapshot from config.yaml rules
// against the current providers. Slot model ids are normalized to the exact
// configured spelling (case-insensitive match), otherwise ModelFor would fall
// back to the source's best chat model and cooldown/quota bookkeeping would
// key on a name that never matches.
func (c *Core) buildAutoChain() {
	c.autoChain.Store(c.chainForRules(c.cfg.Auto, false))
}

// chainForRules compiles a chat AUTO chain from a rule list. kindFilter is
// false for the chat chain (image models are skipped) and true for the image
// chain (only image models participate). Per-key chains always compile as
// chat chains: a user's key configures the models *they* talk to.
func (c *Core) chainForRules(rules []config.ModelScope, imageOnly bool) *scheduler.Chain {
	prov := func(model, source string) scheduler.Provider {
		p := c.registry.ProviderForSlot(model, source)
		if p == nil {
			return nil
		}
		m := p.ModelByID(model)
		if imageOnly && (m == nil || m.Kind != "image") {
			return nil
		}
		if !imageOnly && m != nil && m.Kind == "image" {
			return nil
		}
		return p
	}
	sr := make([]scheduler.Rule, 0, len(rules))
	for _, e := range rules {
		model, source := e.Model, e.Source
		if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil {
			if exact := p.ModelIDFold(e.Model); exact != "" {
				model = exact
			}
			if m := p.ModelByID(model); m != nil && m.Kind == "image" {
				if !imageOnly {
					continue // image-kind models never join the chat AUTO chain
				}
			} else if imageOnly {
				continue
			}
			if source == "" {
				source = p.Name()
			}
		}
		sr = append(sr, scheduler.Rule{
			Model:  model,
			Source: source,
			Tier:   e.Tier,
			Quota:  e.TokenQuota,
			Period: e.Period,
			Hours:  e.Hours,
		})
	}
	return scheduler.BuildChain(sr, prov)
}

// buildAutoImageChain rebuilds the image-generation AUTO chain snapshot.
// Unlike buildAutoChain, only image-kind models participate: chat models in
// the rules are skipped so a stale chat slot can't receive image traffic.
func (c *Core) buildAutoImageChain() {
	prov := func(model, source string) scheduler.Provider {
		p := c.registry.ProviderForSlot(model, source)
		if p == nil {
			return nil
		}
		if m := p.ModelByID(model); m != nil && m.Kind != "image" {
			return nil
		}
		return p
	}
	rules := c.cfg.AutoImage
	sr := make([]scheduler.Rule, 0, len(rules))
	for _, e := range rules {
		model, source := e.Model, e.Source
		if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil {
			if exact := p.ModelIDFold(e.Model); exact != "" {
				model = exact
			}
			if m := p.ModelByID(model); m != nil && m.Kind != "image" {
				continue // chat-kind models never join the image AUTO chain
			}
			if source == "" {
				source = p.Name()
			}
		}
		sr = append(sr, scheduler.Rule{
			Model:  model,
			Source: source,
			Tier:   e.Tier,
			Quota:  e.TokenQuota,
			Period: e.Period,
			Hours:  e.Hours,
		})
	}
	c.autoImageChain.Store(scheduler.BuildChain(sr, prov))
}

// AutoSlotState is the UI-facing health snapshot of one AUTO chain slot.
type AutoSlotState struct {
	Model         string `json:"model"`
	Source        string `json:"source"`
	Pref          int64  `json:"pref"`
	FailCount     int64  `json:"fail_count"`
	CooldownUntil int64  `json:"cooldown_until"`
	Cooling       bool   `json:"cooling"`
	// CooldownFrom is when the current cooldown window opened and ProbeAfter is
	// its midpoint: from ProbeAfter on, a single probe request is allowed
	// through so a recovered upstream does not have to sit out the rest of the
	// window. Probing reports whether such a probe is in flight right now.
	CooldownFrom int64 `json:"cooldown_from"`
	ProbeAfter   int64 `json:"probe_after"`
	Probing      bool  `json:"probing"`
	// Probeable marks a cooling slot that is past its midpoint, i.e. the next
	// AUTO request may use it as a probe.
	Probeable bool `json:"probeable"`
}

// AutoSlotStates returns per-slot health for every slot of the current chain.
func (c *Core) AutoSlotStates() []AutoSlotState {
	ch := c.autoChain.Load()
	if ch == nil {
		return nil
	}
	now := time.Now().Unix()
	var out []AutoSlotState
	for _, tn := range ch.Tiers {
		for _, sl := range tn.Slots {
			pp, ok := sl.Prov.(*provider.Provider)
			if !ok {
				continue
			}
			pref, fail, until := pp.ModelHealthInfo(sl.Model)
			from, probeAfter, probing := pp.ModelProbeInfo(sl.Model)
			out = append(out, AutoSlotState{
				Model:         sl.Model,
				Source:        sl.Source,
				Pref:          pref,
				FailCount:     fail,
				CooldownUntil: until,
				Cooling:       until > now,
				CooldownFrom:  from,
				ProbeAfter:    probeAfter,
				Probing:       probing,
				Probeable:     until > now && probeAfter > 0 && now >= probeAfter,
			})
		}
	}
	return out
}

// Reload re-reads the runtime store and rebuilds sources.
func (c *Core) Reload() error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if err := c.store.Load(); err != nil {
		return err
	}
	return c.rebuildRegistry()
}

// ---- adapter management (web UI) ----

func (c *Core) ListAdapters() []lua.APIAdapter {
	return c.vm.ListAdapters()
}

func (c *Core) UploadAdapter(name, code string) error {
	if name == "" {
		return fmt.Errorf("adapter name required")
	}
	if err := os.MkdirAll(c.cfg.AdapterDir, 0755); err != nil {
		return err
	}
	path := filepath.Join(c.cfg.AdapterDir, name+".lua")
	if err := os.WriteFile(path, []byte(code), 0644); err != nil {
		return err
	}
	return c.vm.LoadAdapter(path)
}

func (c *Core) RemoveAdapter(name string) error {
	path := filepath.Join(c.cfg.AdapterDir, name+".lua")
	_ = os.Remove(path)
	c.vm.RemoveAdapter(name)
	return nil
}

// ---- source management (web UI) ----

func (c *Core) AddSource(src config.Source) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if err := normalizeSource(&src); err != nil {
		return err
	}
	if err := config.UpsertSourceInYAML(c.cfg.Path, src.Name, src, c.cfg.SecretBox()); err != nil {
		return err
	}
	// Update in-memory Sources so mergedSources() finds the entry.
	replaced := false
	for i, s := range c.cfg.Sources {
		if s.Name == src.Name {
			c.cfg.Sources[i] = src
			replaced = true
			break
		}
	}
	if !replaced {
		c.cfg.Sources = append(c.cfg.Sources, src)
	}
	return c.rebuildRegistry()
}

func (c *Core) RemoveSource(name string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, s := range c.cfg.Sources {
		if s.Name == name {
			if err := config.RemoveSourceFromYAML(c.cfg.Path, name); err != nil {
				return err
			}
			c.cfg.Sources = c.removeCfgSource(name)
			break
		}
	}
	if _, err := c.store.Remove(name); err != nil {
		return err
	}
	return c.rebuildRegistry()
}

func (c *Core) removeCfgSource(name string) []config.Source {
	out := c.cfg.Sources[:0]
	for _, s := range c.cfg.Sources {
		if s.Name != name {
			out = append(out, s)
		}
	}
	return out
}

func (c *Core) Sources() []config.Source {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.mergedSources()
}

// Templates returns the saved source templates.
func (c *Core) Templates() []config.SourceTemplate {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.store.ListTemplates()
}

// SaveTemplate upserts a source template.
func (c *Core) SaveTemplate(t config.SourceTemplate) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if t.Name == "" {
		return fmt.Errorf("template requires a name")
	}
	return c.store.UpsertTemplate(t)
}

// RemoveTemplate deletes a source template by name.
func (c *Core) RemoveTemplate(name string) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, err := c.store.RemoveTemplate(name); err != nil {
		return err
	}
	return nil
}

func normalizeSource(s *config.Source) error {
	if s.Name == "" || s.BaseURL == "" {
		return fmt.Errorf("source requires name and base_url")
	}
	if len(s.Models) == 0 {
		return fmt.Errorf("source requires at least one model")
	}
	if s.Adapter == "" {
		s.Adapter = "openai"
	}
	if s.MaxConcurrent == 0 {
		s.MaxConcurrent = 8
	}
	return nil
}

// Close releases resources.
func (c *Core) Close() {
	// Plugin state must be flushed BEFORE the VM stops. The saver's final write
	// reads each plugin's Lua tables; once vm.Stop() has closed those states the
	// read finds nothing and the last interval of accumulation is lost — which
	// is the exact failure this persistence was added to prevent.
	if c.plugins != nil {
		c.plugins.Close()
	}
	if c.vm != nil {
		c.vm.Stop()
	}
}

// applyBillingDSL compiles the configured billing profile and injects it into
// the billing plugin.
//
// Configured pricing replaces the hand-written JSON blob that used to be PUT
// through the state API. That path was used twice in production and both times
// it "succeeded" while being wrong: once without the required `prices` wrapper
// (silently wiping the accumulated totals), once with a peak-window shape the
// plugin does not read (peak traffic billed at off-peak rates, no error). A
// typed config with load-time validation turns both mistakes into startup
// errors naming the file and field.
//
// A DSL with no profile, or one that fails to compile, leaves the plugin's
// prices untouched: billing keeps running on whatever it had, and the problem
// is logged rather than becoming a startup failure — the gateway must forward
// even when its cost accounting is misconfigured.
func (c *Core) applyBillingDSL() {
	ps := c.plugins
	if ps == nil || ps.Count() == 0 {
		return
	}
	// The billing plugin may simply not be installed; a DSL without it is a
	// no-op, not an error.
	var have bool
	for _, row := range ps.List() {
		if row["name"] == "billing" && row["loaded"] == true {
			have = true
			break
		}
	}
	if !have {
		return
	}
	dsl := c.cfg.BillingDSL
	if dsl == nil || len(dsl.Profiles) == 0 {
		return
	}
	if err := dsl.Validate(); err != nil {
		log.Printf("[billing] config invalid, profile not applied: %v", err)
		return
	}
	profile := dsl.Resolve(dsl.Active)
	if profile == nil {
		log.Printf("[billing] no billing profile resolved (active=%q)", dsl.Active)
		return
	}
	prices, err := billing.Compile(profile, c.cfg.Sources)
	if err != nil {
		log.Printf("[billing] profile %q not applied: %v", profile.ID, err)
		return
	}
	if err := ps.SetState("billing", map[string]interface{}{"prices": prices}); err != nil {
		log.Printf("[billing] applying profile %q: %v", profile.ID, err)
		return
	}
	log.Printf("[billing] profile %q applied (%d rules, %d sources)", profile.ID, len(profile.Rules), len(c.cfg.Sources))
}

// ValidateScopeQuotas checks every scope entry's caps before they are stored.
// A typo in a period must be rejected at write time rather than silently
// becoming a never-resetting budget — the opposite of what was typed.
func ValidateScopeQuotas(entries []config.ModelScope) error {
	for _, e := range entries {
		if err := (config.KeyQuota{
			TokenQuota: e.TokenQuota,
			ReqQuota:   e.ReqQuota,
			Period:     e.Period,
			Hours:      e.Hours,
		}).Validate(); err != nil {
			return fmt.Errorf("model %q: %w", e.Model, err)
		}
	}
	return nil
}