cellbot/internal/di/providers.go

package di

import (
	"cellbot/internal/adapter/milky"
	"cellbot/internal/adapter/onebot11"
	"cellbot/internal/config"
	"cellbot/internal/engine"
	"cellbot/internal/plugins/echo"
	"cellbot/internal/protocol"
	"cellbot/pkg/net"
	"context"

	"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) *engine.Dispatcher {
	return engine.NewDispatcher(eventBus, 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
}

func ProvideEchoPlugin(logger *zap.Logger, botManager *protocol.BotManager, dispatcher *engine.Dispatcher) {
	echoPlugin := echo.NewEchoPlugin(logger, botManager)
	dispatcher.RegisterHandler(echoPlugin)
	logger.Info("Echo plugin registered")
}

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