ProxyPool/internal/telegram/bot.go

package telegram

import (
	"context"
	"log/slog"
	"sync"
	"time"

	"proxyrotator/internal/config"
	"proxyrotator/internal/store"

	tele "gopkg.in/telebot.v3"
)

// Bot Telegram Bot 管理器
type Bot struct {
	mu    sync.RWMutex
	bot   *tele.Bot
	cfg   *config.Config
	store store.ProxyStore

	scheduler *Scheduler
	notifier  *Notifier

	running  bool
	stopChan chan struct{}
}

// Status Bot 状态
type Status struct {
	Running   bool   `json:"running"`
	Connected bool   `json:"connected"`
	Username  string `json:"username,omitempty"`
}

// NewBot 创建 Bot 实例
func NewBot(cfg *config.Config, proxyStore store.ProxyStore) *Bot {
	return &Bot{
		cfg:      cfg,
		store:    proxyStore,
		stopChan: make(chan struct{}),
	}
}

// Start 启动 Bot
func (b *Bot) Start(ctx context.Context) error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.running {
		return nil
	}

	if b.cfg.TelegramBotToken == "" {
		slog.Info("telegram bot token not configured, skipping")
		return nil
	}

	return b.startInternal()
}

// startInternal 内部启动（需要持有锁）
func (b *Bot) startInternal() error {
	pref := tele.Settings{
		Token:  b.cfg.TelegramBotToken,
		Poller: &tele.LongPoller{Timeout: 10 * time.Second},
	}

	bot, err := tele.NewBot(pref)
	if err != nil {
		slog.Error("failed to create telegram bot", "error", err)
		return err
	}

	b.bot = bot
	b.notifier = NewNotifier(bot, b.cfg.TelegramNotifyChatID)
	b.scheduler = NewScheduler(b.store, b.notifier, b.cfg)

	// 注册命令处理器
	b.registerCommands(b.cfg.TelegramAdminIDs)

	// 启动调度器
	b.scheduler.Start()

	// 注册命令菜单
	commands := []tele.Command{
		{Text: "stats", Description: "查看代理池统计"},
		{Text: "groups", Description: "查看分组统计"},
		{Text: "get", Description: "获取可用代理 (默认1个，如 /get 5)"},
		{Text: "import", Description: "导入代理 (如 /import groupname)"},
		{Text: "test", Description: "触发测活 (如 /test groupname)"},
		{Text: "purge", Description: "清理死代理"},
		{Text: "help", Description: "显示帮助信息"},
	}
	if err := bot.SetCommands(commands); err != nil {
		slog.Warn("failed to set bot commands", "error", err)
	}

	// 启动 Bot
	b.stopChan = make(chan struct{})
	go func() {
		slog.Info("telegram bot started", "username", bot.Me.Username)
		bot.Start()
	}()

	b.running = true
	return nil
}

// Stop 停止 Bot
func (b *Bot) Stop() {
	b.mu.Lock()
	defer b.mu.Unlock()

	b.stopInternal()
}

// stopInternal 内部停止（需要持有锁）
func (b *Bot) stopInternal() {
	if !b.running {
		return
	}

	if b.scheduler != nil {
		b.scheduler.Stop()
	}

	if b.bot != nil {
		b.bot.Stop()
		slog.Info("telegram bot stopped")
	}

	close(b.stopChan)
	b.running = false
}

// Status 获取 Bot 状态
func (b *Bot) Status() Status {
	b.mu.RLock()
	defer b.mu.RUnlock()

	status := Status{
		Running: b.running,
	}

	if b.bot != nil && b.running {
		status.Connected = true
		status.Username = b.bot.Me.Username
	}

	return status
}

// TriggerTest 手动触发测活
func (b *Bot) TriggerTest(ctx context.Context) error {
	b.mu.RLock()
	defer b.mu.RUnlock()

	if b.scheduler == nil {
		return nil
	}

	return b.scheduler.RunTest(ctx)
}

// registerCommands 注册命令
func (b *Bot) registerCommands(adminIDs []int64) {
	// 管理员权限中间件
	adminOnly := func(next tele.HandlerFunc) tele.HandlerFunc {
		return func(c tele.Context) error {
			if len(adminIDs) == 0 {
				return next(c)
			}
			userID := c.Sender().ID
			for _, id := range adminIDs {
				if id == userID {
					return next(c)
				}
			}
			return c.Send("⛔ 无权限访问")
		}
	}

	// 创建命令处理器
	cmds := NewCommands(b.store, b.scheduler)

	b.bot.Handle("/start", adminOnly(cmds.HandleStart))
	b.bot.Handle("/help", adminOnly(cmds.HandleHelp))
	b.bot.Handle("/stats", adminOnly(cmds.HandleStats))
	b.bot.Handle("/groups", adminOnly(cmds.HandleGroups))
	b.bot.Handle("/get", adminOnly(cmds.HandleGet))
	b.bot.Handle("/test", adminOnly(cmds.HandleTest))
	b.bot.Handle("/purge", adminOnly(cmds.HandlePurge))
	b.bot.Handle("/import", adminOnly(cmds.HandleImport))
	b.bot.Handle(tele.OnDocument, adminOnly(cmds.HandleDocument))
	b.bot.Handle(tele.OnText, adminOnly(cmds.HandleText))
}