cellbot/internal/di/providers.go

package di

import (
	"context"

	"cellbot/internal/adapter/milky"
	"cellbot/internal/adapter/onebot11"
	"cellbot/internal/config"
	"cellbot/internal/engine"
	_ "cellbot/internal/plugins/echo" // 导入插件以触发 init 函数
	"cellbot/internal/protocol"
	"cellbot/pkg/net"

	"go.uber.org/fx"
	"go.uber.org/zap"
)

func ProvideLogger(cfg *config.Config) (*zap.Logger, error) {
	return config.InitLogger(&cfg.Log)
}

func ProvideConfig() (*config.Config, error) {
	configManager := config.NewConfigManager("configs/config.toml", zap.NewNop())
	if err := configManager.Load(); err != nil {
		return nil, err
	}
	return configManager.Get(), nil
}

func ProvideConfigManager(logger *zap.Logger) (*config.ConfigManager, error) {
	configManager := config.NewConfigManager("configs/config.toml", logger)
	if err := configManager.Load(); err != nil {
		return nil, err
	}
	if err := configManager.Watch(); err != nil {
		logger.Warn("Failed to watch config file", zap.Error(err))
	}
	return configManager, nil
}

func ProvideEventBus(logger *zap.Logger) *engine.EventBus {
	return engine.NewEventBus(logger, 10000)
}

func ProvideDispatcher(eventBus *engine.EventBus, logger *zap.Logger, cfg *config.Config) *engine.Dispatcher {
	dispatcher := engine.NewDispatcher(eventBus, logger)

	// 注册限流中间件
	if cfg.Engine.RateLimit.Enabled {
		rateLimitMiddleware := engine.NewRateLimitMiddleware(
			logger,
			cfg.Engine.RateLimit.RPS,
			cfg.Engine.RateLimit.Burst,
		)
		dispatcher.RegisterMiddleware(rateLimitMiddleware)
		logger.Info("Rate limit middleware registered",
			zap.Int("rps", cfg.Engine.RateLimit.RPS),
			zap.Int("burst", cfg.Engine.RateLimit.Burst))
	}

	return dispatcher
}

func ProvidePluginRegistry(dispatcher *engine.Dispatcher, logger *zap.Logger) *engine.PluginRegistry {
	return engine.NewPluginRegistry(dispatcher, logger)
}

func ProvideBotManager(logger *zap.Logger) *protocol.BotManager {
	return protocol.NewBotManager(logger)
}

func ProvideWebSocketManager(logger *zap.Logger, eventBus *engine.EventBus) *net.WebSocketManager {
	return net.NewWebSocketManager(logger, eventBus)
}

func ProvideServer(cfg *config.Config, logger *zap.Logger, botManager *protocol.BotManager, eventBus *engine.EventBus) *net.Server {
	return net.NewServer(cfg.Server.Host, cfg.Server.Port, logger, botManager, eventBus)
}

func ProvideMilkyBots(cfg *config.Config, logger *zap.Logger, eventBus *engine.EventBus, wsManager *net.WebSocketManager, botManager *protocol.BotManager, lc fx.Lifecycle) error {
	for _, botCfg := range cfg.Bots {
		if botCfg.Protocol == "milky" && botCfg.Enabled {
			logger.Info("Creating Milky bot", zap.String("bot_id", botCfg.ID))

			milkyCfg := &milky.Config{
				ProtocolURL:       botCfg.Milky.ProtocolURL,
				AccessToken:       botCfg.Milky.AccessToken,
				EventMode:         botCfg.Milky.EventMode,
				WebhookListenAddr: botCfg.Milky.WebhookListenAddr,
				Timeout:           botCfg.Milky.Timeout,
				RetryCount:        botCfg.Milky.RetryCount,
			}

			bot := milky.NewBot(botCfg.ID, milkyCfg, eventBus, wsManager, logger)
			botManager.Add(bot)

			lc.Append(fx.Hook{
				OnStart: func(ctx context.Context) error {
					logger.Info("Starting Milky bot", zap.String("bot_id", botCfg.ID))
					return bot.Connect(ctx)
				},
				OnStop: func(ctx context.Context) error {
					logger.Info("Stopping Milky bot", zap.String("bot_id", botCfg.ID))
					return bot.Disconnect(ctx)
				},
			})
		}
	}
	return nil
}

func ProvideOneBot11Bots(cfg *config.Config, logger *zap.Logger, wsManager *net.WebSocketManager, eventBus *engine.EventBus, botManager *protocol.BotManager, lc fx.Lifecycle) error {
	for _, botCfg := range cfg.Bots {
		if botCfg.Protocol == "onebot11" && botCfg.Enabled {
			logger.Info("Creating OneBot11 bot", zap.String("bot_id", botCfg.ID))

			ob11Cfg := &onebot11.Config{
				ConnectionType:    botCfg.OneBot11.ConnectionType,
				Host:              botCfg.OneBot11.Host,
				Port:              botCfg.OneBot11.Port,
				AccessToken:       botCfg.OneBot11.AccessToken,
				WSUrl:             botCfg.OneBot11.WSUrl,
				WSReverseUrl:      botCfg.OneBot11.WSReverseUrl,
				Heartbeat:         botCfg.OneBot11.Heartbeat,
				ReconnectInterval: botCfg.OneBot11.ReconnectInterval,
				HTTPUrl:           botCfg.OneBot11.HTTPUrl,
				HTTPPostUrl:       botCfg.OneBot11.HTTPPostUrl,
				Secret:            botCfg.OneBot11.Secret,
				Timeout:           botCfg.OneBot11.Timeout,
				SelfID:            botCfg.OneBot11.SelfID,
				Nickname:          botCfg.OneBot11.Nickname,
			}

			bot := onebot11.NewBot(botCfg.ID, ob11Cfg, logger, wsManager, eventBus)
			botManager.Add(bot)

			lc.Append(fx.Hook{
				OnStart: func(ctx context.Context) error {
					logger.Info("Starting OneBot11 bot", zap.String("bot_id", botCfg.ID))
					return bot.Connect(ctx)
				},
				OnStop: func(ctx context.Context) error {
					logger.Info("Stopping OneBot11 bot", zap.String("bot_id", botCfg.ID))
					return bot.Disconnect(ctx)
				},
			})
		}
	}
	return nil
}

// LoadPlugins 加载所有插件
func LoadPlugins(logger *zap.Logger, botManager *protocol.BotManager, registry *engine.PluginRegistry) {
	// 从全局注册表加载所有插件（通过 RegisterPlugin 注册的）
	plugins := engine.LoadAllPlugins(botManager, logger)
	for _, plugin := range plugins {
		registry.Register(plugin)
	}

	// 加载所有通过 OnXXX().Handle() 注册的处理器（注入依赖）
	handlers := engine.LoadAllHandlers(botManager, logger)
	for _, handler := range handlers {
		registry.Register(handler)
	}

	totalCount := len(plugins) + len(handlers)
	logger.Info("All plugins loaded",
		zap.Int("plugin_count", len(plugins)),
		zap.Int("handler_count", len(handlers)),
		zap.Int("total_count", totalCount),
		zap.Strings("plugins", engine.GetRegisteredPlugins()))
}

var Providers = fx.Options(
	fx.Provide(
		ProvideConfig,
		ProvideConfigManager,
		ProvideLogger,
		ProvideEventBus,
		ProvideDispatcher,
		ProvidePluginRegistry,
		ProvideBotManager,
		ProvideWebSocketManager,
		ProvideServer,
	),
	fx.Invoke(ProvideMilkyBots),
	fx.Invoke(ProvideOneBot11Bots),
	fx.Invoke(LoadPlugins),
)