opencatd-open/router/setRouter.go

package router

import (
	"context"
	"embed"
	"fmt"
	"io/fs"
	"log"
	"net/http"
	"opencatd-open/internal/api"
	"opencatd-open/internal/channel"
	"opencatd-open/internal/dao"
	"opencatd-open/internal/proxy"
	"opencatd-open/internal/usage"
	"opencatd-open/middleware"
	"opencatd-open/pkg/config"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

func SetRouter(cfg *config.Config, db *gorm.DB, web *embed.FS) {
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup

	if cfg == nil || db == nil {
		panic("cfg or db is nil")
	}

	sqlDB, err := db.DB()
	if err != nil {
		log.Fatalf("Failed to get underlying *sql.DB: %v", err)
	}

	// Initialize DAOs
	userDAO := dao.NewUserDAO(db)
	apiKeyDAO := dao.NewApiKeyDAO(db)
	usageDAO := dao.NewUsageDAO(db)
	dailyDAO := dao.NewDailyUsageDAO(db)
	channelDAO := dao.NewChannelDAO(db)
	modelDAO := dao.NewModelDAO(db)

	// Initialize channel service
	channelSvc := channel.NewService(channelDAO, modelDAO)

	// Initialize health checker and start periodic checks
	healthChecker := channel.NewHealthChecker(channelDAO, channelSvc)
	go healthChecker.StartPeriodicCheck(ctx, 5*time.Minute)

	// Initialize usage recorder and start background worker
	usageRecorder := usage.NewRecorder(usageDAO, dailyDAO)
	usageRecorder.Start(ctx)
	defer usageRecorder.Stop()

	// Initialize gateway
	gateway := proxy.NewGateway(ctx, cfg, db, &wg, userDAO, apiKeyDAO, usageDAO, dailyDAO)
	gateway.SetChannelService(channelSvc)

	// Initialize API handler
	apiHandler := api.NewHandler(db)

	r := gin.Default()
	r.Use(middleware.CORS())

	// Public auth routes
	public := r.Group("/api/auth")
	{
		public.POST("/register", apiHandler.Register)
		public.POST("/login", apiHandler.Login)
	}

	// API routes (authenticated)
	apiGroup := r.Group("/api", middleware.Auth(db))
	{
		// User profile
		apiGroup.GET("/me", apiHandler.Me)
		apiGroup.GET("/profile", apiHandler.Me)

		// User management (admin)
		apiGroup.GET("/users", apiHandler.ListUsers)
		apiGroup.POST("/users", apiHandler.CreateUser)
		apiGroup.DELETE("/users/:id", apiHandler.DeleteUser)

		// API Key management
		apiGroup.GET("/keys", apiHandler.ListApiKeys)
		apiGroup.POST("/keys", apiHandler.CreateApiKey)
		apiGroup.DELETE("/keys/:id", apiHandler.DeleteApiKey)

		// Channel management
		apiGroup.GET("/channels", apiHandler.ListChannels)
		apiGroup.POST("/channels", apiHandler.CreateChannel)
		apiGroup.PUT("/channels/:id", apiHandler.UpdateChannel)
		apiGroup.DELETE("/channels/:id", apiHandler.DeleteChannel)
		apiGroup.GET("/channels/:id/models", apiHandler.GetChannelModels)
		apiGroup.POST("/channels/:id/models", apiHandler.BindChannelModels)

		// Model management
		apiGroup.GET("/models", apiHandler.ListModels)
		apiGroup.POST("/models", apiHandler.CreateModel)
		apiGroup.PUT("/models/:id", apiHandler.UpdateModel)
		apiGroup.DELETE("/models/:id", apiHandler.DeleteModel)
	}

	// LLM proxy routes
	v1 := r.Group("/v1")
	v1.Use(middleware.AuthLLM(db))
	{
		v1.POST("/chat/completions", gateway.HandleChat)
		v1.POST("/messages", gateway.HandleMessages)
		v1.POST("/responses", gateway.HandleResponses)
		v1.GET("/models", gateway.HandleModels)
	}

	// SPA fallback
	idxFS, err := fs.Sub(web, "dist")
	if err != nil {
		panic(err)
	}

	assetsFS, err := fs.Sub(web, "dist/assets")
	if err != nil {
		panic(err)
	}
	r.StaticFS("/assets", http.FS(assetsFS))

	r.NoRoute(func(c *gin.Context) {
		if c.Writer.Status() == http.StatusNotFound {
			c.FileFromFS("/", http.FS(idxFS))
		}
	})

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

	go func() {
		fmt.Println("Starting server at port:", cfg.Port)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("listen: %s\n", err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	fmt.Println("\nShutdown Server ...")
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer shutdownCancel()

	if err := srv.Shutdown(shutdownCtx); err != nil {
		log.Fatalln("Server Shutdown:", err)
	}

	cancel()
	sqlDB.Close()

	waitChan := make(chan struct{})
	go func() {
		wg.Wait()
		close(waitChan)
	}()

	select {
	case <-waitChan:
		fmt.Println("All goroutines have finished")
	case <-shutdownCtx.Done():
		fmt.Println("⚠️ Shutdown timeout")
	}

	fmt.Println("Server exited")
}