spider/internal/handler/router.go

package handler

import (
	"context"
	"fmt"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"

	"spider/internal/config"
	"spider/internal/notification"
	"spider/internal/store"
	"spider/internal/task"
	"spider/internal/telegram"
)

// rateLimitMiddleware limits requests per IP using Redis.
// maxRequests per window duration.
func rateLimitMiddleware(rdb *redis.Client, prefix string, maxRequests int64, window time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := fmt.Sprintf("spider:ratelimit:%s:%s", prefix, c.ClientIP())
		ctx := context.Background()

		count, _ := rdb.Incr(ctx, key).Result()
		if count == 1 {
			rdb.Expire(ctx, key, window)
		}

		if count > maxRequests {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, Response{
				Code:    429,
				Message: "请求过于频繁，请稍后再试",
			})
			return
		}
		c.Next()
	}
}

// scheduleHandler is a package-level reference so main.go can set the scheduler after creation.
var scheduleHandler *ScheduleHandler

// GetScheduleHandler returns the package-level ScheduleHandler so main.go
// can call SetScheduler after the scheduler is created.
func GetScheduleHandler() *ScheduleHandler {
	return scheduleHandler
}

// SetupRouter builds and returns the Gin engine with all routes registered.
func SetupRouter(s *store.Store, taskMgr *task.Manager, rdb *redis.Client, tgMgr *telegram.AccountManager, tgAccountHandler *TgAccountHandler) *gin.Engine {
	r := gin.Default()

	// CORS middleware
	cfg := config.Get()
	var corsOrigins []string
	if cfg != nil {
		corsOrigins = cfg.Server.CORSOrigins
	}
	r.Use(corsMiddleware(corsOrigins))

	// Upload size limit
	maxUpload := 10
	if cfg != nil && cfg.Server.MaxUploadMB > 0 {
		maxUpload = cfg.Server.MaxUploadMB
	}
	r.MaxMultipartMemory = int64(maxUpload) << 20

	// Health check (public, no auth)
	r.GET("/ping", func(c *gin.Context) {
		health := gin.H{"status": "ok"}
		httpStatus := http.StatusOK

		sqlDB, err := s.DB.DB()
		if err != nil {
			health["mysql"] = "error: " + err.Error()
			health["status"] = "degraded"
			httpStatus = http.StatusServiceUnavailable
		} else {
			ctx, cancel := context.WithTimeout(c.Request.Context(), 3*time.Second)
			defer cancel()
			if err := sqlDB.PingContext(ctx); err != nil {
				health["mysql"] = "error: " + err.Error()
				health["status"] = "degraded"
				httpStatus = http.StatusServiceUnavailable
			} else {
				health["mysql"] = "ok"
			}
		}

		ctx, cancel := context.WithTimeout(c.Request.Context(), 3*time.Second)
		defer cancel()
		if err := rdb.Ping(ctx).Err(); err != nil {
			health["redis"] = "error: " + err.Error()
			health["status"] = "degraded"
			httpStatus = http.StatusServiceUnavailable
		} else {
			health["redis"] = "ok"
		}

		c.JSON(httpStatus, health)
	})

	api := r.Group("/api/v1")

	// Set Redis for token blacklist and DB for permissions
	SetAuthRedis(rdb)
	setPermissionDB(s.DB)

	// ── Public routes (no auth) ──
	auth := &AuthHandler{store: s}
	api.POST("/auth/login", rateLimitMiddleware(rdb, "login", 10, time.Minute), auth.Login)

	// App info (public, needed for branding)
	api.GET("/app/info", func(c *gin.Context) {
		appCfg := config.Get()
		name := "商户采集系统"
		if appCfg != nil && appCfg.App.Name != "" {
			name = appCfg.App.Name
		}
		OK(c, gin.H{"name": name})
	})

	// Level map (needed by login page for display)
	sh := &SettingHandler{store: s}
	api.GET("/settings/level-map", sh.GetLevelMap)

	// ── Protected routes (require login) ──
	protected := api.Group("")
	protected.Use(JWTAuth())

	// Dashboard & System Health
	dh := &DashboardHandler{store: s, rdb: rdb}
	protected.GET("/dashboard", dh.Get)
	protected.GET("/system/health", dh.Health)

	// Auth - self
	protected.GET("/auth/profile", auth.GetProfile)
	protected.PUT("/auth/profile", auth.UpdateProfile)
	protected.PUT("/auth/password", auth.ChangePassword)
	protected.POST("/auth/logout", auth.Logout)

	// User's own permissions
	ph := &PermissionHandler{store: s}
	protected.GET("/auth/permissions", ph.GetMyPermissions)

	// Keywords
	kw := &KeywordHandler{store: s}
	protected.GET("/keywords", kw.List)
	protected.POST("/keywords", RequireRole("admin", "operator"), RequireAction("keyword_manage"), kw.Create)
	protected.POST("/keywords/import", RequireRole("admin", "operator"), RequireAction("keyword_manage"), kw.ImportCSV)
	protected.PUT("/keywords/:id", RequireRole("admin", "operator"), RequireAction("keyword_manage"), kw.Update)
	protected.DELETE("/keywords/:id", RequireRole("admin", "operator"), RequireAction("keyword_manage"), kw.Delete)

	// Merchants
	mc := &MerchantHandler{store: s}
	protected.GET("/merchants/stats", mc.Stats)
	protected.GET("/merchants/raw", mc.ListRaw)
	protected.GET("/merchants/clean", mc.ListClean)
	protected.GET("/merchants/clean/export", mc.ExportCSV)
	protected.GET("/merchants/clean/users", mc.ListUsers)
	protected.GET("/merchants/clean/industry-tags", mc.ListIndustryTags)
	protected.POST("/merchants/clean/import", RequireRole("admin", "operator"), RequireAction("merchant_import"), mc.ImportCSV)
	protected.GET("/merchants/:id", mc.GetByID)
	protected.PUT("/merchants/clean/:id", RequireRole("admin", "operator"), RequireAction("merchant_edit"), mc.UpdateClean)
	protected.PUT("/merchants/clean/:id/follow-status", RequireRole("admin", "operator"), RequireAction("merchant_edit"), mc.UpdateFollowStatus)
	protected.PUT("/merchants/clean/:id/assign", RequireRole("admin", "operator"), RequireAction("merchant_assign"), mc.AssignMerchant)
	protected.POST("/merchants/clean/:id/recheck", RequireRole("admin", "operator"), RequireAction("merchant_edit"), mc.RecheckMerchant)
	protected.GET("/merchants/clean/:id/notes", mc.ListNotes)
	protected.POST("/merchants/clean/:id/notes", RequireRole("admin", "operator"), mc.AddNote)
	protected.DELETE("/merchants/raw/batch", RequireRole("admin", "operator"), RequireAction("merchant_delete"), mc.BatchDeleteRaw)
	protected.DELETE("/merchants/clean/batch", RequireRole("admin", "operator"), RequireAction("merchant_delete"), mc.BatchDeleteClean)
	protected.PUT("/merchants/clean/batch-assign", RequireRole("admin", "operator"), mc.BatchAssign)
	protected.PUT("/merchants/clean/batch-follow-status", RequireRole("admin", "operator"), mc.BatchFollowStatus)
	protected.PUT("/merchants/clean/batch-level", RequireRole("admin", "operator"), mc.BatchLevel)
	protected.POST("/merchants/clean/merge", RequireRole("admin", "operator"), mc.MergeMerchants)
	protected.POST("/merchants/clean/batch-recheck", RequireRole("admin", "operator"), mc.BatchRecheck)
	protected.POST("/merchants/archive", RequireRole("admin"), mc.ArchiveMerchants)
	protected.GET("/merchants/archived", mc.ListArchived)
	protected.POST("/merchants/archived/:id/restore", RequireRole("admin"), mc.RestoreArchived)

	// Tasks
	th := &TaskHandler{store: s, taskMgr: taskMgr}
	protected.GET("/tasks", th.List)
	protected.POST("/tasks/start", RequireRole("admin", "operator"), RequireAction("task_start"), th.Start)
	protected.GET("/tasks/:id", th.Get)
	protected.POST("/tasks/:id/stop", RequireRole("admin", "operator"), RequireAction("task_stop"), th.Stop)
	protected.POST("/tasks/:id/retry", RequireRole("admin", "operator"), RequireAction("task_start"), th.Retry)
	protected.GET("/tasks/:id/logs", th.Logs)
	protected.GET("/tasks/:id/details", th.Details)
	protected.GET("/plugins", th.Plugins)

	// Analytics
	an := &AnalyticsHandler{store: s}
	protected.GET("/analytics/funnel", an.Funnel)
	protected.GET("/analytics/source-efficiency", an.SourceEfficiency)
	protected.GET("/analytics/trends", an.Trends)
	protected.GET("/analytics/trends/export", an.ExportTrends)

	// Proxies
	px := &ProxyHandler{store: s, taskMgr: taskMgr}
	protected.GET("/proxies", px.List)
	protected.GET("/proxies/enabled", px.ListEnabled)
	protected.GET("/proxies/pool-status", px.PoolStatus)
	protected.POST("/proxies", RequireRole("admin"), px.Create)
	protected.POST("/proxies/test-all", RequireRole("admin"), px.TestAll)
	protected.PUT("/proxies/:id", RequireRole("admin"), px.Update)
	protected.DELETE("/proxies/:id", RequireRole("admin"), px.Delete)
	protected.POST("/proxies/:id/test", RequireRole("admin"), px.Test)

	// Channels
	ch := &ChannelHandler{store: s}
	protected.GET("/channels", ch.List)
	protected.GET("/channels/stats", ch.Stats)
	protected.PUT("/channels/:id/status", RequireRole("admin", "operator"), ch.UpdateStatus)
	protected.DELETE("/channels/:id", RequireRole("admin"), ch.Delete)

	// Groups
	gh := &GroupHandler{store: s, tgManager: tgMgr}
	protected.GET("/groups", gh.ListGroups)
	protected.GET("/groups/:username/members", gh.ListMembers)
	protected.GET("/groups/:username/members/export", gh.ExportMembers)
	protected.GET("/member-groups/:username", gh.ListMemberGroups)
	protected.GET("/members/search", gh.SearchMembers)
	protected.POST("/groups/:username/clone-members", RequireRole("admin", "operator"), gh.CloneMembers)

	// Settings (admin only)
	protected.GET("/settings/grading", sh.GetGrading)
	protected.PUT("/settings/grading", RequireRole("admin"), sh.UpdateGrading)
	protected.POST("/settings/grading/reset", RequireRole("admin"), sh.ResetGrading)

	// Audit logs (admin only)
	ah := &AuditHandler{store: s}

	// User management (admin only)
	uh := &UserHandler{store: s}
	adminOnly := protected.Group("", RequireRole("admin"))
	adminOnly.GET("/audit-logs", ah.List)

	// Permissions management (admin only)
	adminOnly.GET("/permissions", ph.ListAll)
	adminOnly.PUT("/permissions/:role", ph.Update)
	adminOnly.POST("/permissions/reset", ph.Reset)

	// Notification configs (admin only)
	notifMgr := notification.NewManager(s.DB)
	nh := &NotificationHandler{store: s, manager: notifMgr}
	adminOnly.GET("/notification-configs", nh.List)
	adminOnly.POST("/notification-configs", nh.Create)
	adminOnly.PUT("/notification-configs/:id", nh.Update)
	adminOnly.DELETE("/notification-configs/:id", nh.Delete)
	adminOnly.POST("/notification-configs/:id/test", nh.Test)
	adminOnly.GET("/users", uh.List)
	adminOnly.POST("/users", uh.Create)
	adminOnly.PUT("/users/:id", uh.Update)
	adminOnly.POST("/users/:id/reset-password", uh.ResetPassword)
	adminOnly.POST("/users/:id/force-logout", uh.ForceLogout)

	// Backup (admin only)
	bh := &BackupHandler{store: s}
	adminOnly.GET("/backup/export", bh.ExportJSON)
	adminOnly.GET("/backup/stats", bh.Stats)
	adminOnly.DELETE("/users/:id", uh.Delete)

	// Schedules (admin only)
	scheduleHandler = &ScheduleHandler{store: s, taskMgr: taskMgr}
	adminOnly.GET("/schedules", scheduleHandler.List)
	adminOnly.POST("/schedules", scheduleHandler.Create)
	adminOnly.PUT("/schedules/:id", scheduleHandler.Update)
	adminOnly.DELETE("/schedules/:id", scheduleHandler.Delete)
	adminOnly.POST("/schedules/:id/run", scheduleHandler.RunNow)

	// TG account management
	ta := tgAccountHandler
	protected.GET("/tg-accounts", ta.List) // all logged-in users can view
	adminOnly.POST("/tg-accounts", ta.Create)
	adminOnly.POST("/tg-accounts/import", ta.Import)
	adminOnly.PUT("/tg-accounts/:id", ta.Update)
	adminOnly.POST("/tg-accounts/:id/test", ta.Test)
	adminOnly.POST("/tg-accounts/:id/reveal-2fa", ta.RevealTwoFA)
	adminOnly.DELETE("/tg-accounts/:id", ta.Delete)

	return r
}

// corsMiddleware handles Cross-Origin Resource Sharing.
func corsMiddleware(allowedOrigins []string) gin.HandlerFunc {
	return func(c *gin.Context) {
		origin := c.GetHeader("Origin")
		if len(allowedOrigins) == 0 {
			c.Header("Access-Control-Allow-Origin", "*")
		} else {
			for _, o := range allowedOrigins {
				if o == origin || o == "*" {
					c.Header("Access-Control-Allow-Origin", origin)
					break
				}
			}
		}
		c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		c.Header("Access-Control-Allow-Headers", "Origin, Content-Type, Authorization")
		c.Header("Access-Control-Max-Age", "86400")

		if c.Request.Method == "OPTIONS" {
			c.AbortWithStatus(http.StatusNoContent)
			return
		}

		c.Next()
	}
}

// ServerAddr returns the listen address string for the given port.
func ServerAddr(port int) string {
	return fmt.Sprintf(":%d", port)
}