// 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"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync/atomic"
	"time"

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

// Core owns the running configuration and adapters.
type Core struct {
	cfg       *config.Config
	vm        *lua.VM
	store     *config.Store
	scheduler *scheduler.Scheduler
	registry  *provider.Registry
	autoChain atomic.Pointer[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)
	}
	c.store = config.NewStore(cfg.RuntimeFile)
	if err := c.store.Load(); err != nil {
		return nil, fmt.Errorf("runtime store: %w", err)
	}
	c.scheduler = scheduler.New(buildRetries(cfg))
	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
	}
	return c, nil
}

// seedAuto migrates the legacy per-source model priority into flat AUTO
// scheduling slots (one slot per model, priority order) the first time no
// explicit auto rules exist.
func (c *Core) seedAuto() error {
	if len(c.store.AutoRules()) > 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})
	}
	return c.store.SaveAutoRules(entries)
}

// seedKeys migrates the static config gateway_keys into the runtime store as
// admin keys (once), so later UI-created keys can share the same store.
func (c *Core) seedKeys() error {
	existing := map[string]bool{}
	for _, k := range c.store.ListKeys() {
		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)
		}
		if err := c.store.SaveKey(config.GWKey{
			Key:       raw,
			Role:      "admin",
			Name:      name,
			CreatedAt: time.Now().Unix(),
			Seed:      true,
		}); err != nil {
			return err
		}
		changed = true
	}
	if changed {
		return c.store.Load()
	}
	return nil
}

func buildRetries(cfg *config.Config) int {
	return len(cfg.Sources) // allow fallback across all sources
}

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

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 (immutable after build;
// a rebuilt chain is swapped in atomically). nil before the first build or
// when no auto slots could be resolved.
func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.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 }

// TLS returns the configured cert/key file paths. Empty strings mean HTTP only.
func (c *Core) TLS() (cert, key string) {
	return c.cfg.TLSCertFile, c.cfg.TLSKeyFile
}

// PublicBaseURL returns the externally advertised base used in generated
// connection snippets, or "" to infer it from the incoming request.
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 { return c.store.ListKeys() }

// FindKey looks up a gateway key record by its secret value.
func (c *Core) FindKey(key string) (config.GWKey, bool) { return c.store.KeyByValue(key) }

// CreateKey builds a new random gateway key and persists it.
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
	models = cleanScopes(models)
	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(),
	}
	if rec.Role == "" {
		rec.Role = "user"
	}
	if err := c.store.SaveKey(rec); err != nil {
		return config.GWKey{}, err
	}
	return rec, nil
}

// UpdateKey mutates a key's name/role/model scope and persists it.
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
	rec, ok := c.store.KeyByValue(key)
	if !ok {
		return config.GWKey{}, fmt.Errorf("key not found")
	}
	if name != "" {
		rec.Name = name
	}
	if role == "admin" || role == "user" {
		rec.Role = role
	}
	rec.Models = cleanScopes(models)
	rec.Note = note
	if err := c.store.SaveKey(rec); err != nil {
		return config.GWKey{}, err
	}
	return rec, nil
}

// DeleteKey removes a key record; returns false if it did not exist.
func (c *Core) DeleteKey(key string) (bool, error) { return c.store.DeleteKey(key) }

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

// AutoRules returns the AUTO scheduling slots in priority order (slot 0 =
// highest priority).
func (c *Core) AutoRules() []config.ModelScope { return c.store.AutoRules() }

func cleanScopes(entries []config.ModelScope) []config.ModelScope {
	clean := make([]config.ModelScope, 0, len(entries))
	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, rebuilds the chain and
// clears the cooldown of every slot in it — preference scores are kept, so a
// reliably good model keeps its edge while an edited chain applies
// immediately. Providers are NOT rebuilt here (their per-model state survives
// the edit, plan 2.4 lifecycle); rebuildRegistry covers source edits.
func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
	if err := c.store.SaveAutoRules(cleanScopes(entries)); err != nil {
		return err
	}
	c.buildAutoChain()
	if ch := c.autoChain.Load(); ch != nil {
		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)
				}
			}
		}
	}
	return nil
}

// ResetHealth clears the scheduling backoff state of every provider (admin
// UI action). Unlike SaveAutoRules this does not touch the chain itself.
func (c *Core) ResetHealth() {
	for _, p := range c.registry.Providers() {
		p.ResetHealth()
	}
}

// Registry resolves model -> owning provider.
func (c *Core) ProviderForModel(model string) *provider.Provider {
	return c.registry.ProviderForModel(model)
}

// ProviderForSlot resolves a (model, source) scheduling slot to a provider;
// source "" falls back to ProviderForModel.
func (c *Core) ProviderForSlot(model, source string) *provider.Provider {
	return c.registry.ProviderForSlot(model, source)
}

// Config exposes the underlying configuration (read-only usage).
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])
			// Runtime sources (web UI edits) are persisted without timeout
			// fields; apply the same defaults the YAML path gets so a dead
			// upstream cannot hold a concurrency slot forever (P10-3).
			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 (env var overrides inline api_key) and
// decrypts an inline enc:v1: ciphertext (useful for YAML-sourced keys).
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:") {
		s.APIKey = box.MustDecrypt(s.APIKey)
	}
	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()
	return nil
}

// buildAutoChain rebuilds the AUTO chain snapshot from the persisted rules
// against the current providers. Slots whose (model, source) no longer exists
// and image-kind models are dropped; a chain with no slots makes AUTO
// requests answer "no auto slot configured".
func (c *Core) buildAutoChain() {
	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.store.AutoRules()
	sr := make([]scheduler.Rule, 0, len(rules))
	for _, e := range rules {
		r := scheduler.Rule{
			Model:  e.Model,
			Source: e.Source,
			Tier:   e.Tier,
			Quota:  e.TokenQuota,
			Period: e.Period,
			Hours:  e.Hours,
		}
		if r.Source == "" {
			// canonicalize to the owning source so summaries/audit/quota
			// windows always carry a real source name
			if p := c.registry.ProviderForSlot(e.Model, ""); p != nil {
				r.Source = p.Name()
			}
		}
		sr = append(sr, r)
	}
	c.autoChain.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"`
}

// AutoSlotStates returns per-slot health (preference, failure count,
// cooldown) for every slot of the current AUTO chain, mirroring the chain
// order so the priority-page UI can annotate its blocks.
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)
			out = append(out, AutoSlotState{
				Model:         sl.Model,
				Source:        sl.Source,
				Pref:          pref,
				FailCount:     fail,
				CooldownUntil: until,
				Cooling:       until > now,
			})
		}
	}
	return out
}

// Reload re-reads the runtime store and rebuilds sources (adapter reload is not
// strictly needed since adapters are loaded into the VM at startup; uploaded
// adapters are placed in the adapter dir and loaded by the web UI).
func (c *Core) Reload() error {
	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()
}

// UploadAdapter saves a new Lua adapter script to the adapter dir and loads it.
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)
}

// RemoveAdapter deletes an adapter script and evicts it from the VM. The file
// is removed for real (adapter dir is authoritative after first run), so the
// adapter stays gone across restarts.
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 {
	if err := normalizeSource(&src); err != nil {
		return err
	}
	if err := c.store.Upsert(src); err != nil {
		return err
	}
	return c.rebuildRegistry()
}

func (c *Core) RemoveSource(name string) error {
	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 { return c.mergedSources() }

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() {
	if c.vm != nil {
		c.vm.Stop()
	}
}