nex/backend/cmd/desktop/main.go

package main

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

	"github.com/gin-gonic/gin"
	"github.com/getlantern/systray"
	"github.com/gofrs/flock"
	"github.com/pressly/goose/v3"
	"go.uber.org/zap"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"nex/backend/internal/config"
	"nex/backend/internal/conversion"
	"nex/backend/internal/conversion/anthropic"
	"nex/backend/internal/conversion/openai"
	"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"

	"nex/embedfs"
)

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

func main() {
	port := 9826

	singleLock := NewSingletonLock(filepath.Join(os.TempDir(), "nex-gateway.lock"))
	if err := singleLock.Lock(); err != nil {
		showError("Nex Gateway", "已有 Nex 实例运行")
		os.Exit(1)
	}
	defer singleLock.Unlock()

	if err := checkPortAvailable(port); err != nil {
		showError("Nex Gateway", err.Error())
		os.Exit(1)
	}

	cfg, err := config.LoadConfig()
	if err != nil {
		showError("Nex Gateway", fmt.Sprintf("加载配置失败: %v", err))
		os.Exit(1)
	}

	zapLogger, err = pkgLogger.New(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 {
		showError("Nex Gateway", fmt.Sprintf("初始化日志失败: %v", err))
		os.Exit(1)
	}
	defer zapLogger.Sync()

	db, err := initDatabase(cfg)
	if err != nil {
		showError("Nex Gateway", fmt.Sprintf("初始化数据库失败: %v", err))
		os.Exit(1)
	}
	defer 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 := registry.Register(openai.NewAdapter()); err != nil {
		zapLogger.Fatal("注册 OpenAI 适配器失败", zap.String("error", err.Error()))
	}
	if err := registry.Register(anthropic.NewAdapter()); err != nil {
		zapLogger.Fatal("注册 Anthropic 适配器失败", zap.String("error", err.Error()))
	}
	engine := conversion.NewConversionEngine(registry, zapLogger)

	providerClient := provider.NewClient()

	proxyHandler := handler.NewProxyHandler(engine, providerClient, routingService, providerService, statsService)
	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)
	setupStaticFiles(r)

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

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

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

	go func() {
		time.Sleep(500 * time.Millisecond)
		if err := openBrowser(fmt.Sprintf("http://localhost:%d", port)); err != nil {
			zapLogger.Warn("无法打开浏览器", zap.String("error", err.Error()))
		}
	}()

	setupSystray(port)
}

func initDatabase(cfg *config.Config) (*gorm.DB, error) {
	dbDir := filepath.Dir(cfg.Database.Path)
	if err := os.MkdirAll(dbDir, 0755); err != nil {
		return nil, fmt.Errorf("创建数据库目录失败: %w", err)
	}

	db, err := gorm.Open(sqlite.Open(cfg.Database.Path), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Info),
	})
	if err != nil {
		return nil, err
	}

	if err := runMigrations(db); err != nil {
		return nil, fmt.Errorf("数据库迁移失败: %w", err)
	}

	if err := db.Exec("PRAGMA journal_mode=WAL").Error; err != nil {
		log.Printf("警告: 启用 WAL 模式失败: %v", err)
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}
	sqlDB.SetMaxIdleConns(cfg.Database.MaxIdleConns)
	sqlDB.SetMaxOpenConns(cfg.Database.MaxOpenConns)
	sqlDB.SetConnMaxLifetime(cfg.Database.ConnMaxLifetime)

	return db, nil
}

func runMigrations(db *gorm.DB) error {
	sqlDB, err := db.DB()
	if err != nil {
		return err
	}

	migrationsDir := getMigrationsDir()
	if _, err := os.Stat(migrationsDir); os.IsNotExist(err) {
		return fmt.Errorf("迁移目录不存在: %s", migrationsDir)
	}

	goose.SetDialect("sqlite3")
	if err := goose.Up(sqlDB, migrationsDir); err != nil {
		return err
	}

	return nil
}

func getMigrationsDir() string {
	_, filename, _, ok := runtime.Caller(0)
	if ok {
		dir := filepath.Join(filepath.Dir(filename), "..", "..", "migrations")
		if abs, err := filepath.Abs(dir); err == nil {
			return abs
		}
	}
	return "./migrations"
}

func closeDB(db *gorm.DB) {
	sqlDB, err := db.DB()
	if err != nil {
		return
	}
	sqlDB.Close()
}

func setupRoutes(r *gin.Engine, proxyHandler *handler.ProxyHandler, providerHandler *handler.ProviderHandler, modelHandler *handler.ModelHandler, statsHandler *handler.StatsHandler) {
	r.Any("/v1/*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"})
	})
}

func setupStaticFiles(r *gin.Engine) {
	distFS, err := fs.Sub(embedfs.FrontendDist, "frontend-dist")
	if err != nil {
		zapLogger.Fatal("无法加载前端资源", zap.String("error", err.Error()))
	}

	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("/favicon.svg", func(c *gin.Context) {
		data, err := fs.ReadFile(distFS, "favicon.svg")
		if err != nil {
			c.Status(404)
			return
		}
		c.Data(200, "image/svg+xml", data)
	})

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

		if strings.HasPrefix(path, "/api/") ||
			strings.HasPrefix(path, "/v1/") ||
			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 setupSystray(port int) {
	systray.Run(func() {
		var icon []byte
		var err error
		if runtime.GOOS == "windows" {
			icon, err = embedfs.Assets.ReadFile("assets/icon.ico")
		} else {
			icon, err = embedfs.Assets.ReadFile("assets/icon.png")
		}
		if err != nil {
			zapLogger.Error("无法加载托盘图标", zap.String("error", err.Error()))
		}
		systray.SetIcon(icon)
		systray.SetTitle("Nex Gateway")
		systray.SetTooltip("AI Gateway")

		mOpen := systray.AddMenuItem("打开管理界面", "在浏览器中打开")
		systray.AddSeparator()
		mStatus := systray.AddMenuItem("状态: 运行中", "")
		mStatus.Disable()
		mPort := systray.AddMenuItem(fmt.Sprintf("端口: %d", port), "")
		mPort.Disable()
		systray.AddSeparator()
		mAbout := systray.AddMenuItem("关于", "")
		systray.AddSeparator()
		mQuit := systray.AddMenuItem("退出", "停止服务并退出")

		go func() {
			for {
				select {
				case <-mOpen.ClickedCh:
					openBrowser(fmt.Sprintf("http://localhost:%d", port))
				case <-mAbout.ClickedCh:
					showAbout()
				case <-mQuit.ClickedCh:
					doShutdown()
					systray.Quit()
					return
				}
			}
		}()
	}, nil)
}

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

	if server != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		server.Shutdown(ctx)
	}

	if shutdownCancel != nil {
		shutdownCancel()
	}
}

func checkPortAvailable(port int) error {
	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return fmt.Errorf("端口 %d 已被占用\n\n可能原因:\n- 已有 Nex 实例运行\n- 其他程序占用了该端口\n\n请检查并关闭占用端口的程序", port)
	}
	ln.Close()
	return nil
}

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

func showError(title, message string) {
	switch runtime.GOOS {
	case "darwin":
		script := fmt.Sprintf(`display dialog "%s" buttons {"OK"} default button "OK" with title "%s"`, message, title)
		exec.Command("osascript", "-e", script).Run()
	case "windows":
		messageBox(title, message, MB_ICONERROR)
	case "linux":
		exec.Command("zenity", "--error", fmt.Sprintf("--title=%s", title), fmt.Sprintf("--text=%s", message)).Run()
	}
}

func showAbout() {
	message := "Nex Gateway\n\nAI Gateway - 统一的大模型 API 网关\n\nhttps://github.com/nex/gateway"
	switch runtime.GOOS {
	case "darwin":
		script := fmt.Sprintf(`display dialog "%s" buttons {"OK"} default button "OK" with title "关于 Nex Gateway"`, message)
		exec.Command("osascript", "-e", script).Run()
	case "windows":
		messageBox("关于 Nex Gateway", message, MB_ICONINFORMATION)
	case "linux":
		exec.Command("zenity", "--info", "--title=关于 Nex Gateway", fmt.Sprintf("--text=%s", message)).Run()
	}
}

const (
	MB_ICONERROR      = 0x10
	MB_ICONINFORMATION = 0x40
)

var (
	user32          = syscall.NewLazyDLL("user32.dll")
	procMessageBoxW = user32.NewProc("MessageBoxW")
)

func messageBox(title, message string, flags uint) {
	titlePtr, _ := syscall.UTF16PtrFromString(title)
	messagePtr, _ := syscall.UTF16PtrFromString(message)
	procMessageBoxW.Call(
		0,
		uintptr(unsafe.Pointer(messagePtr)),
		uintptr(unsafe.Pointer(titlePtr)),
		uintptr(flags),
	)
}