gemini-banlancer/internal/router/router.go

// Filename: internal/router/router.go
package router

import (
	"gemini-balancer/internal/config"
	"gemini-balancer/internal/domain/proxy"
	"gemini-balancer/internal/domain/upstream"
	"gemini-balancer/internal/handlers"
	"gemini-balancer/internal/middleware"
	"gemini-balancer/internal/pongo"
	"gemini-balancer/internal/service"
	"gemini-balancer/internal/settings"
	"gemini-balancer/internal/webhandlers"
	"net/http"
	"os"
	"time"

	"github.com/gin-contrib/cors"
	"github.com/gin-gonic/gin"
)

func NewRouter(
	// Core Services
	cfg *config.Config,
	securityService *service.SecurityService,
	settingsManager *settings.SettingsManager,
	// Core Handlers
	proxyHandler *handlers.ProxyHandler,
	apiAuthHandler *handlers.APIAuthHandler,
	// Admin API Handlers
	keyGroupHandler *handlers.KeyGroupHandler,
	apiKeyHandler *handlers.APIKeyHandler,
	tokensHandler *handlers.TokensHandler,
	logHandler *handlers.LogHandler,
	settingHandler *handlers.SettingHandler,
	dashboardHandler *handlers.DashboardHandler,
	taskHandler *handlers.TaskHandler,
	// Web Page Handlers
	webAuthHandler *webhandlers.WebAuthHandler,
	pageHandler *webhandlers.PageHandler,
	// === Domain Modules ===
	upstreamModule *upstream.Module,
	proxyModule *proxy.Module,
) *gin.Engine {
	if cfg.Log.Level != "debug" {
		gin.SetMode(gin.ReleaseMode)
	}
	router := gin.Default()

	router.Static("/static", "./web/static")
	// CORS 配置
	config := cors.Config{
		// 允许前端的来源。在生产环境中，需改为实际域名
		AllowOrigins:     []string{"http://localhost:9000"},
		AllowMethods:     []string{"GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"},
		AllowHeaders:     []string{"Origin", "Content-Type", "Accept", "Authorization"},
		ExposeHeaders:    []string{"Content-Length"},
		AllowCredentials: true,
		MaxAge:           12 * time.Hour,
	}
	router.Use(cors.New(config))
	isDebug := gin.Mode() != gin.ReleaseMode
	router.HTMLRender = pongo.New("web/templates", isDebug)

	// --- 基础设施 ---
	router.GET("/", func(c *gin.Context) { c.Redirect(http.StatusMovedPermanently, "/dashboard") })
	router.GET("/health", func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"status": "ok"}) })
	// --- 统一的认证管道 ---
	apiAdminAuth := middleware.APIAdminAuthMiddleware(securityService)
	webAdminAuth := middleware.WebAdminAuthMiddleware(securityService)

	router.Use(gin.RecoveryWithWriter(os.Stdout))
	// --- 将正确的依赖和中间件管道传递下去 ---
	registerProxyRoutes(router, proxyHandler, securityService)
	registerAdminRoutes(router, apiAdminAuth, keyGroupHandler, tokensHandler, apiKeyHandler, logHandler, settingHandler, dashboardHandler, taskHandler, upstreamModule, proxyModule)
	registerPublicAPIRoutes(router, apiAuthHandler, securityService, settingsManager)
	registerWebRoutes(router, webAdminAuth, webAuthHandler, pageHandler)
	return router
}

func registerProxyRoutes(
	router *gin.Engine, proxyHandler *handlers.ProxyHandler, securityService *service.SecurityService,
) {
	// 通用的代理认证中间件
	proxyAuthMiddleware := middleware.ProxyAuthMiddleware(securityService)
	// --- 模式一: 智能聚合模式 (根路径) ---
	// /v1 和 /v1beta 路径作为默认入口，服务于 BasePool 聚合逻辑
	v1 := router.Group("/v1")
	v1.Use(proxyAuthMiddleware)
	{
		v1.Any("/*path", proxyHandler.HandleProxy)
	}
	v1beta := router.Group("/v1beta")
	v1beta.Use(proxyAuthMiddleware)
	{
		v1beta.Any("/*path", proxyHandler.HandleProxy)
	}
	// --- 模式二: 精确路由模式 (/proxy/:group_name) ---
	// 创建一个新的、物理隔离的路由组，用于按组名精确路由
	proxyGroup := router.Group("/proxy/:group_name")
	proxyGroup.Use(proxyAuthMiddleware)
	{
		// 捕获所有子路径 (例如 /v1/chat/completions)，并全部交给同一个 ProxyHandler。
		proxyGroup.Any("/*path", proxyHandler.HandleProxy)
	}
}

// registerAdminRoutes
func registerAdminRoutes(
	router *gin.Engine,
	authMiddleware gin.HandlerFunc,
	keyGroupHandler *handlers.KeyGroupHandler,
	tokensHandler *handlers.TokensHandler,
	apiKeyHandler *handlers.APIKeyHandler,
	logHandler *handlers.LogHandler,
	settingHandler *handlers.SettingHandler,
	dashboardHandler *handlers.DashboardHandler,
	taskHandler *handlers.TaskHandler,
	upstreamModule *upstream.Module,
	proxyModule *proxy.Module,
) {
	admin := router.Group("/admin", authMiddleware)
	{
		// --- KeyGroup Base Routes ---
		admin.POST("/keygroups", keyGroupHandler.CreateKeyGroup)
		admin.GET("/keygroups", keyGroupHandler.GetKeyGroups)
		admin.PUT("/keygroups/order", keyGroupHandler.UpdateKeyGroupOrder)
		// --- KeyGroup Specific Routes (by :id) ---
		admin.GET("/keygroups/:id", keyGroupHandler.GetKeyGroups)
		admin.PUT("/keygroups/:id", keyGroupHandler.UpdateKeyGroup)
		admin.DELETE("/keygroups/:id", keyGroupHandler.DeleteKeyGroup)
		admin.POST("/keygroups/:id/clone", keyGroupHandler.CloneKeyGroup)
		admin.GET("/keygroups/:id/stats", keyGroupHandler.GetKeyGroupStats)
		admin.POST("/keygroups/:id/bulk-actions", apiKeyHandler.HandleBulkAction)
		// --- APIKey Sub-resource Routes under a KeyGroup ---
		keyGroupAPIKeys := admin.Group("/keygroups/:id/apikeys")
		{
			keyGroupAPIKeys.GET("", apiKeyHandler.ListKeysForGroup)
			keyGroupAPIKeys.GET("/export", apiKeyHandler.ExportKeysForGroup)
			keyGroupAPIKeys.POST("/bulk", apiKeyHandler.AddMultipleKeysToGroup)
			keyGroupAPIKeys.DELETE("/bulk", apiKeyHandler.UnlinkMultipleKeysFromGroup)
			keyGroupAPIKeys.POST("/test", apiKeyHandler.TestKeysForGroup)
			keyGroupAPIKeys.PUT("/:keyId", apiKeyHandler.UpdateGroupAPIKeyMapping)
		}

		// Global key operations
		admin.GET("/apikeys", apiKeyHandler.ListAPIKeys)
		// admin.PUT("/apikeys/:id", apiKeyHandler.UpdateAPIKey) // DEPRECATED: Status is now contextual
		admin.POST("/apikeys/test", apiKeyHandler.TestMultipleKeys)           // Test keys globally
		admin.DELETE("/apikeys/:id", apiKeyHandler.HardDeleteAPIKey)          // Hard delete a single key
		admin.DELETE("/apikeys/bulk", apiKeyHandler.HardDeleteMultipleKeys)   // Hard delete multiple keys
		admin.PUT("/apikeys/bulk/restore", apiKeyHandler.RestoreMultipleKeys) // Restore multiple keys globally

		// --- Global Routes ---
		admin.GET("/tokens", tokensHandler.GetAllTokens)
		admin.PUT("/tokens", tokensHandler.UpdateTokens)
		admin.GET("/logs", logHandler.GetLogs)
		admin.GET("/settings", settingHandler.GetSettings)
		admin.PUT("/settings", settingHandler.UpdateSettings)
		admin.PUT("/settings/reset", settingHandler.ResetSettingsToDefaults)

		// 用于查询异步任务的状态
		admin.GET("/tasks/:id", taskHandler.GetTaskStatus)

		// 领域模块
		upstreamModule.RegisterRoutes(admin)
		proxyModule.RegisterRoutes(admin)
		// --- 全局仪表盘路由 ---
		dashboard := admin.Group("/dashboard")
		{
			dashboard.GET("/overview", dashboardHandler.GetOverview)
			dashboard.GET("/chart", dashboardHandler.GetChart)
			dashboard.GET("/stats/:period", dashboardHandler.GetRequestStats) // 点击详情
		}
	}
}

// registerWebRoutes
func registerWebRoutes(
	router *gin.Engine,
	authMiddleware gin.HandlerFunc,
	webAuthHandler *webhandlers.WebAuthHandler,
	pageHandler *webhandlers.PageHandler,
) {
	router.GET("/login", webAuthHandler.ShowLoginPage)
	router.POST("/login", webAuthHandler.HandleLogin)
	router.GET("/logout", webAuthHandler.HandleLogout)
	// For Test only router.Run("127.0.0.1:9000")
	// 受保护的Admin Web界面
	webGroup := router.Group("/", authMiddleware)
	webGroup.Use(authMiddleware)
	{
		webGroup.GET("/keys", pageHandler.ShowKeysPage)
		webGroup.GET("/settings", pageHandler.ShowConfigEditorPage)
		webGroup.GET("/logs", pageHandler.ShowErrorLogsPage)
		webGroup.GET("/dashboard", pageHandler.ShowDashboardPage)
		webGroup.GET("/tasks", pageHandler.ShowTasksPage)
		webGroup.GET("/chat", pageHandler.ShowChatPage)
	}

}

// registerPublicAPIRoutes 无需后台登录的公共API路由
func registerPublicAPIRoutes(router *gin.Engine, apiAuthHandler *handlers.APIAuthHandler, securityService *service.SecurityService, settingsManager *settings.SettingsManager) {
	ipBanMiddleware := middleware.IPBanMiddleware(securityService, settingsManager)
	publicAPIGroup := router.Group("/api")
	{
		publicAPIGroup.POST("/login", ipBanMiddleware, apiAuthHandler.HandleLogin)
	}
}