nex/backend/cmd/desktop/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"io/fs"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"nex/embedfs"

	"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"
	"nex/backend/pkg/buildinfo"

	"github.com/gin-gonic/gin"
	"github.com/gofrs/flock"
	"go.uber.org/zap"
	"gorm.io/gorm"

	pkgLogger "nex/backend/pkg/logger"
)

var (
	server         *http.Server
	zapLogger      *zap.Logger
	shutdownCtx    context.Context
	shutdownCancel context.CancelFunc
	desktopHooks   = defaultDesktopRuntimeHooks()
)

type singletonLocker interface {
	Lock() error
	Unlock() error
}

type desktopRuntimeHooks struct {
	loadConfig       func() (*config.Config, config.ConfigMetadata, error)
	newLock          func(string) singletonLocker
	listen           func(int) (net.Listener, error)
	upgradeLogger    func(*zap.Logger, pkgLogger.Config) (*zap.Logger, error)
	initDB           func(*config.DatabaseConfig, *zap.Logger) (*gorm.DB, error)
	closeDB          func(*gorm.DB)
	registerAdapters func(conversion.AdapterRegistry) error
	setupStaticFiles func(*gin.Engine) error
	startServer      func(*http.Server, net.Listener, chan<- error, *zap.Logger)
	setupSystray     func(int, <-chan error) error
}

func defaultDesktopRuntimeHooks() desktopRuntimeHooks {
	return desktopRuntimeHooks{
		loadConfig:       config.LoadDesktopConfigWithMetadata,
		newLock:          func(lockPath string) singletonLocker { return NewSingletonLock(lockPath) },
		listen:           listenDesktopPort,
		upgradeLogger:    pkgLogger.Upgrade,
		initDB:           database.Init,
		closeDB:          database.Close,
		registerAdapters: registerDesktopAdapters,
		setupStaticFiles: setupStaticFiles,
		startServer:      startDesktopServer,
		setupSystray:     setupSystray,
	}
}

func main() {
	minimalLogger := pkgLogger.NewMinimal()
	if err := runDesktop(minimalLogger); err != nil {
		reportStartupFailure(err, dialogLogger())
		os.Exit(1)
	}
}

func runDesktop(minimalLogger *zap.Logger) error {
	if minimalLogger == nil {
		minimalLogger = pkgLogger.NewMinimal()
	}

	cfg, cfgMeta, err := desktopHooks.loadConfig()
	if err != nil {
		return newStartupError(phaseConfig, desktopConfigErrorMessage(getDesktopConfigPath(), err), err)
	}

	port := cfg.Server.Port

	singleLock := desktopHooks.newLock(filepath.Join(os.TempDir(), "nex-gateway.lock"))
	if err := singleLock.Lock(); err != nil {
		return newStartupError(phaseSingleton, "已有 Nex 实例运行", err)
	}
	defer func() {
		if err := singleLock.Unlock(); err != nil {
			minimalLogger.Warn("释放实例锁失败", zap.Error(err))
		}
	}()

	listener, err := desktopHooks.listen(port)
	if err != nil {
		return newStartupError(phasePort, desktopPortUnavailableMessage(port), err)
	}
	defer listener.Close()

	zapLogger, err = desktopHooks.upgradeLogger(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 {
		return newStartupError(phaseLogger, fmt.Sprintf("初始化日志失败\n\n日志目录: %s\n\n请检查目录权限或磁盘空间", cfg.Log.Path), err)
	}
	defer func() {
		if err := zapLogger.Sync(); err != nil {
			minimalLogger.Warn("同步日志失败", zap.Error(err))
		}
	}()

	cfg.PrintSummary(zapLogger)

	db, err := desktopHooks.initDB(&cfg.Database, zapLogger)
	if err != nil {
		phase := phaseDatabase
		message := fmt.Sprintf("数据库初始化失败\n\n请检查数据库配置、文件权限或连接状态\n\n%v", err)
		if errors.Is(err, database.ErrMigration) {
			phase = phaseMigration
			message = fmt.Sprintf("数据库迁移失败\n\n请查看日志或检查数据库迁移权限\n\n%v", err)
		}
		return newStartupError(phase, message, err)
	}
	defer desktopHooks.closeDB(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()
	defer statsBuffer.Stop()

	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 := desktopHooks.registerAdapters(registry); err != nil {
		return newStartupError(phaseAdapter, startupInternalErrorMessage(), 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)
	versionHandler := handler.NewVersionHandler()
	settingsHandler := handler.NewSettingsHandler(cfg, "desktop", true, cfgMeta.ConfigPath)

	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, versionHandler, settingsHandler)
	if err := desktopHooks.setupStaticFiles(r); err != nil {
		return newStartupError(phaseStaticResource, startupInternalErrorMessage(), err)
	}

	server = &http.Server{
		Addr:         desktopListenAddr(port),
		Handler:      r,
		ReadTimeout:  cfg.Server.ReadTimeout,
		WriteTimeout: cfg.Server.WriteTimeout,
	}

	shutdownCtx, shutdownCancel = context.WithCancel(context.Background())
	defer doShutdown()

	serverErrCh := make(chan error, 1)
	desktopHooks.startServer(server, listener, serverErrCh, zapLogger)
	select {
	case err := <-serverErrCh:
		return newStartupError(phaseServer, startupServerErrorMessage(), err)
	case <-time.After(50 * time.Millisecond):
	}

	if err := desktopHooks.setupSystray(port, serverErrCh); err != nil {
		return err
	}

	select {
	case err := <-serverErrCh:
		return newStartupError(phaseServer, startupServerErrorMessage(), err)
	default:
		return nil
	}
}

func registerDesktopAdapters(registry conversion.AdapterRegistry) error {
	if err := registry.Register(openai.NewAdapter()); err != nil {
		return err
	}
	return registry.Register(anthropic.NewAdapter())
}

func startDesktopServer(server *http.Server, listener net.Listener, serverErrCh chan<- error, logger *zap.Logger) {
	go func() {
		logger.Info("AI Gateway 启动",
			zap.String("addr", server.Addr),
			zap.String("version", buildinfo.Version()),
			zap.String("commit", buildinfo.Commit()),
			zap.String("build_time", buildinfo.BuildTime()))
		if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
			serverErrCh <- err
		}
	}()
}

func setupRoutes(r *gin.Engine, proxyHandler *handler.ProxyHandler, providerHandler *handler.ProviderHandler, modelHandler *handler.ModelHandler, statsHandler *handler.StatsHandler, versionHandler *handler.VersionHandler, settingsHandler *handler.SettingsHandler) {
	r.Any("/openai/*path", withProtocol("openai", proxyHandler.HandleProxy))
	r.Any("/anthropic/*path", withProtocol("anthropic", proxyHandler.HandleProxy))
	r.GET("/api/version", versionHandler.GetVersion)

	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)
	}

	settings := r.Group("/api/settings")
	{
		settings.GET("/startup", settingsHandler.GetStartupSettings)
		settings.PUT("/startup", settingsHandler.SaveStartupSettings)
	}

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

func withProtocol(protocol string, next gin.HandlerFunc) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Params = append(c.Params, gin.Param{Key: "protocol", Value: protocol})
		next(c)
	}
}

func setupStaticFiles(r *gin.Engine) error {
	distFS, err := frontendDistFS()
	if err != nil {
		return err
	}
	setupStaticFilesWithFS(r, distFS)
	return nil
}

func frontendDistFS() (fs.FS, error) {
	return fs.Sub(embedfs.FrontendDist, "frontend-dist")
}

func setupStaticFilesWithFS(r *gin.Engine, distFS fs.FS) {
	getContentType := func(path string) string {
		if strings.HasSuffix(path, ".js") {
			return "application/javascript"
		}
		if strings.HasSuffix(path, ".css") {
			return "text/css"
		}
		if strings.HasSuffix(path, ".svg") {
			return "image/svg+xml"
		}
		if strings.HasSuffix(path, ".png") {
			return "image/png"
		}
		if strings.HasSuffix(path, ".ico") {
			return "image/x-icon"
		}
		if strings.HasSuffix(path, ".woff") || strings.HasSuffix(path, ".woff2") {
			return "font/woff2"
		}
		return "application/octet-stream"
	}

	r.GET("/assets/*filepath", func(c *gin.Context) {
		filepath := c.Param("filepath")
		data, err := fs.ReadFile(distFS, "assets"+filepath)
		if err != nil {
			c.Status(404)
			return
		}
		c.Data(200, getContentType(filepath), data)
	})

	r.GET("/icon.png", func(c *gin.Context) {
		data, err := fs.ReadFile(distFS, "icon.png")
		if err != nil {
			c.Status(404)
			return
		}
		c.Data(200, "image/png", data)
	})

	r.NoRoute(func(c *gin.Context) {
		path := c.Request.URL.Path

		if strings.HasPrefix(path, "/api/") ||
			strings.HasPrefix(path, "/openai/") ||
			strings.HasPrefix(path, "/anthropic/") ||
			path == "/openai" ||
			path == "/anthropic" ||
			strings.HasPrefix(path, "/health") {
			c.JSON(404, gin.H{"error": "not found"})
			return
		}

		data, err := fs.ReadFile(distFS, "index.html")
		if err != nil {
			c.Status(500)
			return
		}
		c.Data(200, "text/html; charset=utf-8", data)
	})
}

func doShutdown() {
	if zapLogger != nil {
		zapLogger.Info("正在关闭服务器...")
	}

	if server != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := server.Shutdown(ctx); err != nil && zapLogger != nil {
			zapLogger.Warn("关闭服务器失败", zap.Error(err))
		}
	}

	if shutdownCancel != nil {
		shutdownCancel()
	}
}

func getDesktopConfigPath() string {
	configDir, err := config.GetConfigDir()
	if err != nil {
		return "~/.nex/config.yaml"
	}
	return filepath.Join(configDir, "config.yaml")
}

func desktopConfigErrorMessage(configPath string, err error) string {
	return fmt.Sprintf("加载配置失败\n\n配置文件: %s\n\n%v", configPath, err)
}

func desktopListenAddr(port int) string {
	return fmt.Sprintf(":%d", port)
}

func desktopURL(port int) string {
	return fmt.Sprintf("http://localhost:%d", port)
}

func desktopPortMenuTitle(port int) string {
	return fmt.Sprintf("端口: %d", port)
}

func listenDesktopPort(port int) (net.Listener, error) {
	return net.Listen("tcp", desktopListenAddr(port))
}

func desktopPortUnavailableMessage(port int) string {
	return fmt.Sprintf("端口 %d 已被占用\n\n可能原因:\n- 已有 Nex 实例运行\n- 其他程序占用了该端口\n\n请检查并关闭占用端口的程序", port)
}

type SingletonLock struct {
	flock *flock.Flock
}

func NewSingletonLock(lockPath string) *SingletonLock {
	return &SingletonLock{
		flock: flock.New(lockPath),
	}
}

func (s *SingletonLock) Lock() error {
	locked, err := s.flock.TryLock()
	if err != nil {
		return err
	}
	if !locked {
		return fmt.Errorf("已有实例运行")
	}
	return nil
}

func (s *SingletonLock) Unlock() error {
	return s.flock.Unlock()
}

func openBrowser(url string) error {
	var cmd *exec.Cmd

	switch runtime.GOOS {
	case "darwin":
		cmd = exec.Command("open", url)
	case "windows":
		cmd = exec.Command("rundll32", "url.dll,FileProtocolHandler", url)
	case "linux":
		browsers := []string{"xdg-open", "google-chrome", "firefox"}
		for _, browser := range browsers {
			if _, err := exec.LookPath(browser); err == nil {
				cmd = exec.Command(browser, url)
				break
			}
		}
	}

	if cmd == nil {
		return fmt.Errorf("无法打开浏览器")
	}

	return cmd.Start()
}