nex/backend/cmd/server/main.go

package main

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"

	"nex/backend/internal/config"
	"nex/backend/internal/conversion"
	"nex/backend/internal/conversion/anthropic"
	"nex/backend/internal/conversion/openai"
	"nex/backend/internal/database"
	"nex/backend/internal/handler"
	"nex/backend/internal/handler/middleware"
	"nex/backend/internal/provider"
	"nex/backend/internal/repository"
	"nex/backend/internal/service"
	pkgLogger "nex/backend/pkg/logger"
)

func main() {
	minimalLogger := pkgLogger.NewMinimal()

	cfg, err := config.LoadConfig()
	if err != nil {
		minimalLogger.Fatal("加载配置失败", zap.Error(err))
	}

	zapLogger, err := pkgLogger.Upgrade(minimalLogger, pkgLogger.Config{
		Level:      cfg.Log.Level,
		Path:       cfg.Log.Path,
		MaxSize:    cfg.Log.MaxSize,
		MaxBackups: cfg.Log.MaxBackups,
		MaxAge:     cfg.Log.MaxAge,
		Compress:   cfg.Log.Compress,
	})
	if err != nil {
		minimalLogger.Fatal("初始化日志失败", zap.Error(err))
	}
	defer zapLogger.Sync()

	cfg.PrintSummary(zapLogger)

	db, err := database.Init(&cfg.Database, zapLogger)
	if err != nil {
		zapLogger.Fatal("初始化数据库失败", zap.Error(err))
	}
	defer database.Close(db)

	providerRepo := repository.NewProviderRepository(db)
	modelRepo := repository.NewModelRepository(db)
	statsRepo := repository.NewStatsRepository(db)

	routingCache := service.NewRoutingCache(modelRepo, providerRepo, zapLogger)
	if err := routingCache.Preload(); err != nil {
		zapLogger.Warn("缓存预热失败，将使用懒加载", zap.Error(err))
	}

	statsBuffer := service.NewStatsBuffer(statsRepo, zapLogger,
		service.WithFlushInterval(5*time.Second),
		service.WithFlushThreshold(100))
	statsBuffer.Start()

	providerService := service.NewProviderService(providerRepo, modelRepo, routingCache)
	modelService := service.NewModelService(modelRepo, providerRepo, routingCache)
	routingService := service.NewRoutingService(routingCache)
	statsService := service.NewStatsService(statsRepo, statsBuffer)

	registry := conversion.NewMemoryRegistry()
	if err := registry.Register(openai.NewAdapter()); err != nil {
		zapLogger.Fatal("注册 OpenAI 适配器失败", zap.Error(err))
	}
	if err := registry.Register(anthropic.NewAdapter()); err != nil {
		zapLogger.Fatal("注册 Anthropic 适配器失败", zap.Error(err))
	}
	engine := conversion.NewConversionEngine(registry, zapLogger)

	providerClient := provider.NewClient(zapLogger)

	proxyHandler := handler.NewProxyHandler(engine, providerClient, routingService, providerService, statsService, zapLogger)
	providerHandler := handler.NewProviderHandler(providerService)
	modelHandler := handler.NewModelHandler(modelService)
	statsHandler := handler.NewStatsHandler(statsService)

	gin.SetMode(gin.ReleaseMode)
	r := gin.New()

	r.Use(middleware.RequestID())
	r.Use(middleware.Recovery(zapLogger))
	r.Use(middleware.Logging(zapLogger))
	r.Use(middleware.CORS())

	setupRoutes(r, proxyHandler, providerHandler, modelHandler, statsHandler)

	srv := &http.Server{
		Addr:         fmt.Sprintf(":%d", cfg.Server.Port),
		Handler:      r,
		ReadTimeout:  cfg.Server.ReadTimeout,
		WriteTimeout: cfg.Server.WriteTimeout,
	}

	go func() {
		zapLogger.Info("AI Gateway 启动", zap.String("addr", srv.Addr))
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			zapLogger.Fatal("服务器启动失败", zap.Error(err))
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	zapLogger.Info("正在关闭服务器...")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := srv.Shutdown(ctx); err != nil {
		zapLogger.Fatal("服务器强制关闭", zap.Error(err))
	}

	statsBuffer.Stop()

	zapLogger.Info("服务器已关闭")
}

func setupRoutes(r *gin.Engine, proxyHandler *handler.ProxyHandler, providerHandler *handler.ProviderHandler, modelHandler *handler.ModelHandler, statsHandler *handler.StatsHandler) {
	r.Any("/:protocol/*path", proxyHandler.HandleProxy)

	providers := r.Group("/api/providers")
	{
		providers.GET("", providerHandler.ListProviders)
		providers.POST("", providerHandler.CreateProvider)
		providers.GET("/:id", providerHandler.GetProvider)
		providers.PUT("/:id", providerHandler.UpdateProvider)
		providers.DELETE("/:id", providerHandler.DeleteProvider)
	}

	models := r.Group("/api/models")
	{
		models.GET("", modelHandler.ListModels)
		models.POST("", modelHandler.CreateModel)
		models.GET("/:id", modelHandler.GetModel)
		models.PUT("/:id", modelHandler.UpdateModel)
		models.DELETE("/:id", modelHandler.DeleteModel)
	}

	stats := r.Group("/api/stats")
	{
		stats.GET("", statsHandler.GetStats)
		stats.GET("/aggregate", statsHandler.AggregateStats)
	}

	r.GET("/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})
}