gen2d/backend/internal/config/config.go

package config

import (
	"log"

	"github.com/joho/godotenv"
	"github.com/spf13/viper"
)

// Config 应用全局配置。
type Config struct {
	Server   ServerConfig   `mapstructure:"server"`
	Database DatabaseConfig `mapstructure:"database"`
	JWT      JWTConfig      `mapstructure:"jwt"`
	LLM      LLMConfig      `mapstructure:"llm"`
	ImageGen ImageGenConfig `mapstructure:"image_gen"`
}

// ServerConfig HTTP 服务配置。
type ServerConfig struct {
	Port        int    `mapstructure:"port"`
	Mode        string `mapstructure:"mode"`
	MaxFileSize int64  `mapstructure:"max_file_size"`
}

// DatabaseConfig 数据库配置。
type DatabaseConfig struct {
	DSN string `mapstructure:"dsn"`
}

// JWTConfig JWT 签名配置。
type JWTConfig struct {
	Secret string `mapstructure:"secret"`
	Expire int64  `mapstructure:"expire"`
}

// LLMConfig 大语言模型配置。
type LLMConfig struct {
	BaseURL     string  `mapstructure:"base_url"`
	APIKey      string  `mapstructure:"api_key"`
	Model       string  `mapstructure:"model"`
	Temperature float64 `mapstructure:"temperature"`
	MaxTokens   int     `mapstructure:"max_tokens"`
}

// ImageGenConfig 文生图模型配置。
type ImageGenConfig struct {
	BaseURL   string  `mapstructure:"base_url"`
	APIKey    string  `mapstructure:"api_key"`
	Model     string  `mapstructure:"model"`
	Width     int     `mapstructure:"width"`
	Height    int     `mapstructure:"height"`
	NumImages int     `mapstructure:"num_images"`
	Steps     int     `mapstructure:"steps"`
	CFGScale  float64 `mapstructure:"cfg_scale"`
}

// Load 从 YAML 配置文件和环境变量加载配置。
// 优先级：环境变量 > YAML 文件 > 默认值。
func Load() *Config {
	if err := godotenv.Load(); err != nil {
		log.Println("config: no .env file found, using system env or defaults")
	}

	v := viper.New()
	v.SetConfigName("config")
	v.SetConfigType("yaml")
	v.AddConfigPath("internal/config")

	setDefaults(v)

	if err := v.ReadInConfig(); err != nil {
		log.Printf("config: could not read config file: %v, using env + defaults", err)
	} else {
		log.Printf("config: using config file: %s", v.ConfigFileUsed())
	}

	bindEnvVars(v)

	v.AllowEmptyEnv(true)

	var cfg Config
	if err := v.Unmarshal(&cfg); err != nil {
		log.Fatalf("config: unmarshal failed: %v", err)
	}

	return &cfg
}

func setDefaults(v *viper.Viper) {
	v.SetDefault("server.port", 8080)
	v.SetDefault("server.mode", "debug")
	v.SetDefault("server.max_file_size", int64(10<<20))

	v.SetDefault("database.dsn", "data/gen2d.db")

	v.SetDefault("jwt.secret", "gen2d-dev-secret")
	v.SetDefault("jwt.expire", int64(7200))

	v.SetDefault("llm.base_url", "https://api.openai.com/v1")
	v.SetDefault("llm.api_key", "")
	v.SetDefault("llm.model", "gpt-4o")
	v.SetDefault("llm.temperature", 0.7)
	v.SetDefault("llm.max_tokens", 2048)

	v.SetDefault("image_gen.base_url", "https://api.stability.ai/v1")
	v.SetDefault("image_gen.api_key", "")
	v.SetDefault("image_gen.model", "stable-diffusion-xl")
	v.SetDefault("image_gen.width", 1024)
	v.SetDefault("image_gen.height", 1024)
	v.SetDefault("image_gen.num_images", 1)
	v.SetDefault("image_gen.steps", 30)
	v.SetDefault("image_gen.cfg_scale", 7.0)
}

func bindEnvVars(v *viper.Viper) {
	v.BindEnv("server.port", "GEN2D_PORT")
	v.BindEnv("server.mode", "GEN2D_MODE")
	v.BindEnv("server.max_file_size", "GEN2D_MAX_FILE_SIZE")
	v.BindEnv("database.dsn", "GEN2D_DSN")
	v.BindEnv("jwt.secret", "GEN2D_JWT_SECRET")
	v.BindEnv("jwt.expire", "GEN2D_JWT_EXPIRE")

	v.BindEnv("llm.base_url", "GEN2D_LLM_BASE_URL")
	v.BindEnv("llm.api_key", "GEN2D_LLM_API_KEY")
	v.BindEnv("llm.model", "GEN2D_LLM_MODEL")
	v.BindEnv("llm.temperature", "GEN2D_LLM_TEMPERATURE")
	v.BindEnv("llm.max_tokens", "GEN2D_LLM_MAX_TOKENS")

	v.BindEnv("image_gen.base_url", "GEN2D_IMAGE_BASE_URL")
	v.BindEnv("image_gen.api_key", "GEN2D_IMAGE_API_KEY")
	v.BindEnv("image_gen.model", "GEN2D_IMAGE_MODEL")
	v.BindEnv("image_gen.width", "GEN2D_IMAGE_WIDTH")
	v.BindEnv("image_gen.height", "GEN2D_IMAGE_HEIGHT")
	v.BindEnv("image_gen.num_images", "GEN2D_IMAGE_NUM_IMAGES")
	v.BindEnv("image_gen.steps", "GEN2D_IMAGE_STEPS")
	v.BindEnv("image_gen.cfg_scale", "GEN2D_IMAGE_CFG_SCALE")
}