aiaa-notification-server/cmd/server/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"aiaa-notification-service/internal/adapter"
	"aiaa-notification-service/internal/cache"
	"aiaa-notification-service/internal/config"
	"aiaa-notification-service/internal/engine"
	"aiaa-notification-service/internal/handler"
	"aiaa-notification-service/internal/model"
	"aiaa-notification-service/internal/notify"
	"aiaa-notification-service/internal/safew"
	"aiaa-notification-service/internal/store"
	"aiaa-notification-service/internal/subscriber"

	"github.com/gin-gonic/gin"
	"github.com/logbull/logbull-go/logbull"
)

func main() {
	// Load config first (need logbull settings for logger setup)
	cfg, err := config.Load("config/config.yaml")
	if err != nil {
		slog.Error("failed to load config", "error", err)
		os.Exit(1)
	}

	// Setup dual logger: stdout JSON + Logbull remote
	stdoutHandler := slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo})

	var logbullHandler slog.Handler
	if cfg.Logbull.Host != "" {
		lbCfg := logbull.Config{
			Host:      cfg.Logbull.Host,
			ProjectID: cfg.Logbull.ProjectID,
			APIKey:    cfg.Logbull.APIKey,
		}
		switch cfg.Logbull.LogLevel {
		case "DEBUG":
			lbCfg.LogLevel = logbull.DEBUG
		case "WARNING":
			lbCfg.LogLevel = logbull.WARNING
		case "ERROR":
			lbCfg.LogLevel = logbull.ERROR
		default:
			lbCfg.LogLevel = logbull.INFO
		}
		h, err := logbull.NewSlogHandler(lbCfg)
		if err != nil {
			slog.Warn("logbull init failed, using stdout only", "error", err)
		} else {
			logbullHandler = h
		}
	}

	var baseHandler slog.Handler
	if logbullHandler != nil {
		baseHandler = &teeHandler{handlers: []slog.Handler{stdoutHandler, logbullHandler}}
	} else {
		baseHandler = stdoutHandler
	}

	logger := slog.New(baseHandler)
	slog.SetDefault(logger)

	// Flush logbull on exit
	defer func() {
		if lbh, ok := logbullHandler.(interface{ Flush() }); ok {
			lbh.Flush()
		}
		if lbh, ok := logbullHandler.(interface{ Shutdown() }); ok {
			lbh.Shutdown()
		}
		time.Sleep(2 * time.Second) // allow async flush
	}()

	// Connect MySQL
	st, err := store.NewStore(cfg.Database)
	if err != nil {
		slog.Error("failed to connect to MySQL", "error", err)
		os.Exit(1)
	}
	defer st.Close()

	// Connect Redis
	var redisCache *cache.Cache
	if cfg.Redis.Host != "" {
		redisCache, err = cache.NewCache(cfg.Redis)
		if err != nil {
			slog.Warn("redis connection failed, running without cache", "error", err)
			redisCache = nil
		}
	} else {
		slog.Warn("redis not configured, running without cache")
	}
	if redisCache != nil {
		defer redisCache.Close()
	}

	// Build engine
	matcher := engine.NewMatcher(st, redisCache)
	renderer := engine.NewRenderer()

	// DingTalk per-robot rate limit (queue/wait when over 18/min by default)
	var dingtalkLimiter *adapter.DingTalkLimiter
	limitPerMin := cfg.RateLimit.DingTalkPerMin
	if redisCache != nil {
		dingtalkLimiter = adapter.NewDingTalkLimiter(redisCache, limitPerMin)
	} else {
		dingtalkLimiter = adapter.NewMemoryDingTalkLimiter(limitPerMin)
		slog.Warn("dingtalk rate limit using in-memory store (single instance only)")
	}

	// Build sender factory
	senderFactory := func(channelType string) (adapter.ChannelSender, error) {
		return adapter.NewSender(channelType, &cfg.SMTP, dingtalkLimiter)
	}
	router := engine.NewRouter(st, redisCache, senderFactory)

	// Build handlers
	notifySvc := notify.NewService(matcher, st, renderer, router, st)
	notifyH := handler.NewNotifyHandler(notifySvc)
	sourceH := handler.NewSourceHandler(st, redisCache)
	templateH := handler.NewTemplateHandler(st, redisCache)

	poller := func(token string, offset int64, timeout int) ([]adapter.SafewChat, int64, error) {
		return (&adapter.SafeWSender{}).PollGroupChats(token, offset, timeout)
	}
	var watcherStore safew.ChatStore
	if redisCache != nil {
		watcherStore = redisCache
	} else {
		watcherStore = safew.NewMemStore()
		slog.Warn("safew chats using in-memory store")
	}
	safewWatcher := safew.NewWatcher(watcherStore, poller)
	defer safewWatcher.Stop()
	if tokens, err := st.ListSafewTokens(context.Background()); err != nil {
		slog.Warn("list safew tokens", "error", err)
	} else {
		safewWatcher.StartTokens(tokens)
		slog.Info("safew watchers ensured", "count", len(tokens))
	}

	channelH := handler.NewChannelHandler(st, redisCache, safewWatcher)
	ruleH := handler.NewRuleHandler(st, redisCache)
	msgLogH := handler.NewMessageLogHandler(st)

	// Setup Gin
	gin.SetMode(gin.ReleaseMode)
	r := gin.New()
	r.Use(gin.Recovery())

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

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

	// Notify endpoint (source auth + rate limit)
	notifyGroup := api.Group("/notify")
	notifyGroup.Use(handler.SourceAuth(st))
	if redisCache != nil {
		notifyGroup.Use(handler.RateLimit(redisCache, cfg.RateLimit.Default))
	}
	notifyGroup.POST("", notifyH.Handle)

	// Management endpoints (admin auth)
	admin := api.Group("")
	admin.Use(handler.AdminAuth(cfg.Server.AdminKey))

	// Sources
	admin.POST("/sources", sourceH.Create)
	admin.GET("/sources", sourceH.List)
	admin.GET("/sources/:id", sourceH.Get)
	admin.PUT("/sources/:id", sourceH.Update)
	admin.DELETE("/sources/:id", sourceH.Delete)

	// Templates
	admin.POST("/templates", templateH.Create)
	admin.GET("/templates", templateH.List)
	admin.GET("/templates/:id", templateH.Get)
	admin.PUT("/templates/:id", templateH.Update)
	admin.DELETE("/templates/:id", templateH.Delete)

	// Channels
	admin.POST("/channels/safew/chats", channelH.ListSafewChats)
	admin.GET("/channels/:id/chats", channelH.ListChannelSafewChats)
	admin.POST("/channels", channelH.Create)
	admin.GET("/channels", channelH.List)
	admin.GET("/channels/:id", channelH.Get)
	admin.PUT("/channels/:id", channelH.Update)
	admin.DELETE("/channels/:id", channelH.Delete)

	// Rules
	admin.POST("/rules", ruleH.Create)
	admin.GET("/rules", ruleH.List)
	admin.GET("/rules/:id", ruleH.Get)
	admin.PUT("/rules/:id", ruleH.Update)
	admin.DELETE("/rules/:id", ruleH.Delete)
	admin.PATCH("/rules/:id/enable", ruleH.Enable)
	admin.PATCH("/rules/:id/disable", ruleH.Disable)
	admin.PATCH("/rules/:id/channels/:channel_id/enable", ruleH.EnableChannel)
	admin.PATCH("/rules/:id/channels/:channel_id/disable", ruleH.DisableChannel)

	// Message Logs
	admin.GET("/message-logs", msgLogH.List)

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

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	lookup := func(ctx context.Context, name string) (*model.Source, error) {
		return st.GetSourceByName(ctx, name)
	}
	deduper := subscriber.NewCacheDeduper(redisCache, cfg.SubscriptionDedupTTL)
	for _, sub := range cfg.ActiveSubscriptions() {
		sub := sub
		cons, err := subscriber.New(sub, lookup, notifySvc.Process, deduper, redisCache)
		if err != nil {
			slog.Error("subscriber init", "name", sub.Name, "error", err)
			os.Exit(1)
		}
		go func() {
			if err := cons.Run(ctx); err != nil && !errors.Is(err, context.Canceled) {
				slog.Error("subscriber stopped", "name", sub.Name, "error", err)
			}
		}()
		slog.Info("subscriber started", "name", sub.Name, "queue", sub.Queue, "source", sub.Source)
	}

	go func() {
		slog.Info("server starting", "port", cfg.Server.Port)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			slog.Error("server error", "error", err)
			os.Exit(1)
		}
	}()

	<-ctx.Done()
	slog.Info("shutting down...")

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	if err := srv.Shutdown(shutdownCtx); err != nil {
		slog.Error("forced shutdown", "error", err)
	}
	slog.Info("server stopped")
}

// teeHandler fans out log records to multiple handlers (stdout + remote).
type teeHandler struct {
	handlers []slog.Handler
}

func (t *teeHandler) Enabled(ctx context.Context, level slog.Level) bool {
	for _, h := range t.handlers {
		if h.Enabled(ctx, level) {
			return true
		}
	}
	return false
}

func (t *teeHandler) Handle(ctx context.Context, r slog.Record) error {
	for _, h := range t.handlers {
		if h.Enabled(ctx, r.Level) {
			if err := h.Handle(ctx, r.Clone()); err != nil {
				return err
			}
		}
	}
	return nil
}

func (t *teeHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	handlers := make([]slog.Handler, len(t.handlers))
	for i, h := range t.handlers {
		handlers[i] = h.WithAttrs(attrs)
	}
	return &teeHandler{handlers: handlers}
}

func (t *teeHandler) WithGroup(name string) slog.Handler {
	handlers := make([]slog.Handler, len(t.handlers))
	for i, h := range t.handlers {
		handlers[i] = h.WithGroup(name)
	}
	return &teeHandler{handlers: handlers}
}