AIRouter/internal/config/config.go

package config

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"

	"gopkg.in/yaml.v3"
)

// Config is the root configuration structure.
type Config struct {
	Server         ServerConfig           `yaml:"server"`
	Auth           AuthConfig             `yaml:"auth"`
	Scheduler      SchedulerConfig        `yaml:"scheduler"`
	Timeouts       TimeoutConfig          `yaml:"timeouts"`
	Context        ContextConfig          `yaml:"context"`
	Models         map[string]ModelConfig `yaml:"models"`
	Routing        RoutingConfig          `yaml:"routing"`
	CircuitBreaker CircuitBreakerConfig   `yaml:"circuit_breaker"`
	Backpressure   BackpressureConfig     `yaml:"backpressure"`
	Observability  ObservabilityConfig    `yaml:"observability"`
	Storage        StorageConfig          `yaml:"storage"`
}

type ServerConfig struct {
	Host               string   `yaml:"host"`
	Port               int      `yaml:"port"`
	AdminPort          int      `yaml:"admin_port"`
	MaxRequestBodyMB   int      `yaml:"max_request_body_mb"`
	CORSAllowedOrigins []string `yaml:"cors_allowed_origins"`
}

type AuthConfig struct {
	Enabled            bool     `yaml:"enabled"`
	Methods            []string `yaml:"methods"`
	JWTIssuer          string   `yaml:"jwt_issuer"`
	JWTSecretEnv       string   `yaml:"jwt_secret_env"`
	RateLimitPerMinute int      `yaml:"rate_limit_per_minute"`
	RateLimitBurst     int      `yaml:"rate_limit_burst"`
	UsageWindowMinutes int      `yaml:"usage_window_minutes"`
}

type SchedulerConfig struct {
	MaxRunningTasks       int    `yaml:"max_running_tasks"`
	MaxQueuedTasks        int    `yaml:"max_queued_tasks"`
	Fairness              string `yaml:"fairness"`
	PriorityAgingSeconds  int    `yaml:"priority_aging_seconds"`
	ReservedRealtimeSlots int    `yaml:"reserved_realtime_slots"`
}

type TimeoutConfig struct {
	DefaultConnectMs     int `yaml:"default_connect_ms"`
	DefaultQueueMs       int `yaml:"default_queue_ms"`
	DefaultFirstTokenMs  int `yaml:"default_first_token_ms"`
	DefaultInferenceMs   int `yaml:"default_inference_ms"`
	DefaultIdleMs        int `yaml:"default_idle_ms"`
	DefaultTotalMs       int `yaml:"default_total_ms"`
	CancelGracePeriodMs  int `yaml:"cancel_grace_period_ms"`
	ShutdownDrainSeconds int `yaml:"shutdown_drain_seconds"` // 优雅关机时等待运行中任务的超时时间
}

type ContextConfig struct {
	SafetyMarginRatio       float64 `yaml:"safety_margin_ratio"`
	DefaultPolicy           string  `yaml:"default_policy"`
	MaxSessionMessages      int     `yaml:"max_session_messages"`
	SessionIdleTTLMinutes   int     `yaml:"session_idle_ttl_minutes"`
	EnablePromptPersistence bool    `yaml:"enable_prompt_persistence"`
	EnableLLMSummary        bool    `yaml:"enable_llm_summary"`
	SummaryMaxTokens        int     `yaml:"summary_max_tokens"`
	SummaryTimeoutSeconds   int     `yaml:"summary_timeout_seconds"`
}

type ModelConfig struct {
	Provider          string `yaml:"provider"`
	ActualModel       string `yaml:"actual_model"`
	Endpoint          string `yaml:"endpoint"`
	ContextWindow     int    `yaml:"context_window"`
	MaxOutputTokens   int    `yaml:"max_output_tokens"`
	MaxConcurrency    int    `yaml:"max_concurrency"`
	Residency         string `yaml:"residency"`
	CancelSupported   bool   `yaml:"cancel_supported"`
	IdleUnloadSeconds int    `yaml:"idle_unload_seconds"`
}

type RoutingConfig struct {
	SensitiveDataLocalOnly      bool     `yaml:"sensitive_data_local_only"`
	AllowCloudFallbackByDefault bool     `yaml:"allow_cloud_fallback_by_default"`
	OverloadStrategy            []string `yaml:"overload_strategy"`
}

type CircuitBreakerConfig struct {
	ErrorRateThreshold  float64 `yaml:"error_rate_threshold"`
	MinRequests         int     `yaml:"min_requests"`
	WindowSeconds       int     `yaml:"window_seconds"`
	OpenDurationSeconds int     `yaml:"open_duration_seconds"`
	HalfOpenMaxRequests int     `yaml:"half_open_max_requests"`
}

type BackpressureConfig struct {
	Level1Threshold float64 `yaml:"level1_threshold"`
	Level2Threshold float64 `yaml:"level2_threshold"`
	Level3Threshold float64 `yaml:"level3_threshold"`
}

type ObservabilityConfig struct {
	MetricsEnabled     bool   `yaml:"metrics_enabled"`
	MetricsPath        string `yaml:"metrics_path"`
	TracingEnabled     bool   `yaml:"tracing_enabled"`
	PromptLogging      string `yaml:"prompt_logging"`
	AuditRetentionDays int    `yaml:"audit_retention_days"`
	LogLevel           string `yaml:"log_level"`
}

type StorageConfig struct {
	SessionDB string      `yaml:"session_db"`
	TaskState string      `yaml:"task_state"`
	Redis     RedisConfig `yaml:"redis"`
}

type RedisConfig struct {
	Enabled  bool   `yaml:"enabled"`
	Endpoint string `yaml:"endpoint"`
}

var (
	currentConfig *Config
	configMu      sync.RWMutex
)

// Load reads the config from the given YAML file path and applies env overrides.
func Load(path string) (*Config, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read config file: %w", err)
	}

	cfg := &Config{}
	if err := yaml.Unmarshal(data, cfg); err != nil {
		return nil, fmt.Errorf("parse config yaml: %w", err)
	}

	applyDefaults(cfg)
	if err := validate(cfg); err != nil {
		return nil, fmt.Errorf("config validation: %w", err)
	}

	applyEnvOverrides(cfg)

	configMu.Lock()
	currentConfig = cfg
	configMu.Unlock()

	return cfg, nil
}

func applyDefaults(cfg *Config) {
	if cfg.Server.Host == "" {
		cfg.Server.Host = "0.0.0.0"
	}
	if cfg.Server.Port == 0 {
		cfg.Server.Port = 39000
	}
	if cfg.Server.AdminPort == 0 {
		cfg.Server.AdminPort = 8081
	}
	if cfg.Server.MaxRequestBodyMB == 0 {
		cfg.Server.MaxRequestBodyMB = 20
	}
	if cfg.Scheduler.MaxRunningTasks == 0 {
		cfg.Scheduler.MaxRunningTasks = 8
	}
	if cfg.Scheduler.MaxQueuedTasks == 0 {
		cfg.Scheduler.MaxQueuedTasks = 500
	}
	if cfg.Scheduler.Fairness == "" {
		cfg.Scheduler.Fairness = "weighted_fair_queue"
	}
	if cfg.Scheduler.PriorityAgingSeconds == 0 {
		cfg.Scheduler.PriorityAgingSeconds = 30
	}
	if cfg.Timeouts.DefaultConnectMs == 0 {
		cfg.Timeouts.DefaultConnectMs = 5000
	}
	if cfg.Timeouts.DefaultQueueMs == 0 {
		cfg.Timeouts.DefaultQueueMs = 5000
	}
	if cfg.Timeouts.DefaultFirstTokenMs == 0 {
		cfg.Timeouts.DefaultFirstTokenMs = 10000
	}
	if cfg.Timeouts.DefaultInferenceMs == 0 {
		cfg.Timeouts.DefaultInferenceMs = 60000
	}
	if cfg.Timeouts.DefaultIdleMs == 0 {
		cfg.Timeouts.DefaultIdleMs = 15000
	}
	if cfg.Timeouts.DefaultTotalMs == 0 {
		cfg.Timeouts.DefaultTotalMs = 90000
	}
	if cfg.Timeouts.CancelGracePeriodMs == 0 {
		cfg.Timeouts.CancelGracePeriodMs = 3000
	}
	if cfg.Timeouts.ShutdownDrainSeconds == 0 {
		cfg.Timeouts.ShutdownDrainSeconds = 10
	}
	if cfg.Context.SafetyMarginRatio == 0 {
		cfg.Context.SafetyMarginRatio = 0.08
	}
	if cfg.Context.DefaultPolicy == "" {
		cfg.Context.DefaultPolicy = "summary_and_recent"
	}
	if cfg.Context.MaxSessionMessages == 0 {
		cfg.Context.MaxSessionMessages = 200
	}
	if cfg.Context.SessionIdleTTLMinutes == 0 {
		cfg.Context.SessionIdleTTLMinutes = 60
	}
	if cfg.Context.SummaryMaxTokens == 0 {
		cfg.Context.SummaryMaxTokens = 256
	}
	if cfg.Context.SummaryTimeoutSeconds == 0 {
		cfg.Context.SummaryTimeoutSeconds = 15
	}
	if cfg.CircuitBreaker.ErrorRateThreshold == 0 {
		cfg.CircuitBreaker.ErrorRateThreshold = 0.1
	}
	if cfg.CircuitBreaker.MinRequests == 0 {
		cfg.CircuitBreaker.MinRequests = 10
	}
	if cfg.CircuitBreaker.WindowSeconds == 0 {
		cfg.CircuitBreaker.WindowSeconds = 60
	}
	if cfg.CircuitBreaker.OpenDurationSeconds == 0 {
		cfg.CircuitBreaker.OpenDurationSeconds = 30
	}
	if cfg.CircuitBreaker.HalfOpenMaxRequests == 0 {
		cfg.CircuitBreaker.HalfOpenMaxRequests = 1
	}
	if cfg.Backpressure.Level1Threshold == 0 {
		cfg.Backpressure.Level1Threshold = 0.70
	}
	if cfg.Backpressure.Level2Threshold == 0 {
		cfg.Backpressure.Level2Threshold = 0.85
	}
	if cfg.Backpressure.Level3Threshold == 0 {
		cfg.Backpressure.Level3Threshold = 0.95
	}
	if cfg.Observability.MetricsPath == "" {
		cfg.Observability.MetricsPath = "/metrics"
	}
	if cfg.Observability.PromptLogging == "" {
		cfg.Observability.PromptLogging = "metadata_only"
	}
	if cfg.Observability.LogLevel == "" {
		cfg.Observability.LogLevel = "info"
	}
	if cfg.Observability.AuditRetentionDays == 0 {
		cfg.Observability.AuditRetentionDays = 180
	}
	if cfg.Auth.RateLimitPerMinute <= 0 {
		cfg.Auth.RateLimitPerMinute = 60
	}
	if cfg.Auth.RateLimitBurst <= 0 {
		cfg.Auth.RateLimitBurst = 10
	}
	if cfg.Auth.UsageWindowMinutes <= 0 {
		cfg.Auth.UsageWindowMinutes = 60
	}
	if cfg.Storage.SessionDB == "" {
		cfg.Storage.SessionDB = "sqlite:///var/lib/edgeai/sessions.db"
	}
	if cfg.Storage.TaskState == "" {
		cfg.Storage.TaskState = "sqlite:///var/lib/edgeai/tasks.db"
	}
}

func validate(cfg *Config) error {
	if cfg.Server.Port <= 0 || cfg.Server.Port > 65535 {
		return fmt.Errorf("server.port must be in [1, 65535]")
	}
	if cfg.Server.AdminPort < 0 || cfg.Server.AdminPort > 65535 {
		return fmt.Errorf("server.admin_port must be in [0, 65535]")
	}
	if cfg.Server.AdminPort > 0 && cfg.Server.AdminPort == cfg.Server.Port {
		return fmt.Errorf("server.admin_port must differ from server.port")
	}
	if cfg.Server.MaxRequestBodyMB <= 0 {
		return fmt.Errorf("server.max_request_body_mb must be positive")
	}
	if cfg.Scheduler.MaxRunningTasks <= 0 {
		return fmt.Errorf("scheduler.max_running_tasks must be positive")
	}
	if cfg.Scheduler.MaxQueuedTasks <= 0 {
		return fmt.Errorf("scheduler.max_queued_tasks must be positive")
	}
	if cfg.Context.SafetyMarginRatio < 0 || cfg.Context.SafetyMarginRatio >= 1 {
		return fmt.Errorf("context.safety_margin_ratio must be in [0, 1)")
	}
	if cfg.Context.MaxSessionMessages <= 0 {
		return fmt.Errorf("context.max_session_messages must be positive")
	}
	if cfg.Backpressure.Level1Threshold >= cfg.Backpressure.Level2Threshold {
		return fmt.Errorf("backpressure level1 threshold must be less than level2")
	}
	if cfg.Backpressure.Level2Threshold >= cfg.Backpressure.Level3Threshold {
		return fmt.Errorf("backpressure level2 threshold must be less than level3")
	}
	if cfg.Auth.RateLimitPerMinute <= 0 {
		return fmt.Errorf("auth.rate_limit_per_minute must be positive")
	}
	if cfg.Auth.RateLimitBurst <= 0 {
		return fmt.Errorf("auth.rate_limit_burst must be positive")
	}
	if len(cfg.Models) == 0 {
		return fmt.Errorf("models must not be empty")
	}
	for name, mc := range cfg.Models {
		if mc.Provider == "" {
			return fmt.Errorf("models.%s.provider must not be empty", name)
		}
		if mc.Endpoint == "" {
			return fmt.Errorf("models.%s.endpoint must not be empty", name)
		}
		if mc.ActualModel == "" {
			return fmt.Errorf("models.%s.actual_model must not be empty", name)
		}
		if mc.ContextWindow < 0 {
			return fmt.Errorf("models.%s.context_window must not be negative", name)
		}
		if mc.MaxOutputTokens < 0 {
			return fmt.Errorf("models.%s.max_output_tokens must not be negative", name)
		}
	}
	return nil
}

func applyEnvOverrides(cfg *Config) {
	if v := os.Getenv("EDGEAI_SERVER_PORT"); v != "" {
		if port, err := strconv.Atoi(v); err == nil {
			cfg.Server.Port = port
		}
	}
	if v := os.Getenv("EDGEAI_ADMIN_PORT"); v != "" {
		if port, err := strconv.Atoi(v); err == nil {
			cfg.Server.AdminPort = port
		}
	}
	if v := os.Getenv("EDGEAI_LOG_LEVEL"); v != "" {
		cfg.Observability.LogLevel = v
	}
	if v := os.Getenv("EDGEAI_DB_PATH"); v != "" {
		cfg.Storage.SessionDB = "sqlite://" + v + "/sessions.db"
		cfg.Storage.TaskState = "sqlite://" + v + "/tasks.db"
	}
	if v := os.Getenv("EDGEAI_CONFIG_PATH"); v != "" {
		// already handled by Load path
		_ = v
	}
}

// Get returns the current config (thread-safe).
func Get() *Config {
	configMu.RLock()
	defer configMu.RUnlock()
	return currentConfig
}

// Update replaces the current config (thread-safe).
func Update(cfg *Config) {
	configMu.Lock()
	currentConfig = cfg
	configMu.Unlock()
}

// ConfigPath returns the config file path from env or default.
func ConfigPath() string {
	path := os.Getenv("EDGEAI_CONFIG_PATH")
	if path == "" {
		return "configs/config.yaml"
	}
	return path
}

// PriorityName returns the string name for a priority level.
func PriorityName(p int) string {
	names := []string{"P0", "P1", "P2", "P3", "P4"}
	if p >= 0 && p < len(names) {
		return names[p]
	}
	return "P2"
}

// ParsePriority parses a priority string like "P0" to an int.
func ParsePriority(s string) int {
	s = strings.ToUpper(s)
	for i, name := range []string{"P0", "P1", "P2", "P3", "P4"} {
		if s == name {
			return i
		}
	}
	return 2 // default P2
}