transmission-telegram/internal/bot/bot.go

package bot

import (
	"context"
	"strings"
	"sync"

	tgbotapi "gopkg.in/telegram-bot-api.v4"
	"transmission-telegram/internal/config"
	"transmission-telegram/internal/logger"
	"transmission-telegram/internal/monitor"
	transmissionClient "transmission-telegram/internal/transmission"
)

// Bot represents the Telegram bot
type Bot struct {
	api     *tgbotapi.BotAPI
	client  transmissionClient.Client
	cfg     *config.Config
	logger  *logger.Logger
	updates <-chan tgbotapi.Update
	monitor *monitor.Monitor
	chatID  int64
	chatMu  sync.RWMutex
}

// NewBot creates a new bot instance
func NewBot(api *tgbotapi.BotAPI, client transmissionClient.Client, cfg *config.Config, log *logger.Logger, updates <-chan tgbotapi.Update, mon *monitor.Monitor) *Bot {
	return &Bot{
		api:     api,
		client:  client,
		cfg:     cfg,
		logger:  log,
		updates: updates,
		monitor: mon,
	}
}

// Run starts the bot's main loop
func (b *Bot) Run(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case update := <-b.updates:
			b.handleUpdate(update)
		}
	}
}

// SendMessage sends a message to a chat
func (b *Bot) SendMessage(chatID int64, text string, markdown bool) int {
	return sendMessage(b.api, chatID, text, markdown)
}

// SendCompletionNotification sends a completion notification to the registered chat
func (b *Bot) SendCompletionNotification(msg string) {
	b.chatMu.RLock()
	chatID := b.chatID
	b.chatMu.RUnlock()
	
	if chatID != 0 {
		b.SendMessage(chatID, msg, false)
	}
}

// safeHandler wraps a handler function with panic recovery
func (b *Bot) safeHandler(handler func()) {
	defer func() {
		if r := recover(); r != nil {
			b.logger.Printf("[ERROR] Panic recovered in handler: %v", r)
		}
	}()
	handler()
}

// handleUpdate processes a Telegram update
func (b *Bot) handleUpdate(update tgbotapi.Update) {
	if update.Message == nil {
		return
	}

	// Check if user is master
	if !b.cfg.IsMaster(update.Message.From.UserName) {
		b.logger.Printf("[INFO] Ignored a message from: %s", update.Message.From.String())
		return
	}

	// Update chatID for completion notifications
	b.chatMu.Lock()
	if b.chatID != update.Message.Chat.ID {
		b.chatID = update.Message.Chat.ID
		if b.monitor != nil {
			b.monitor.SetChatID(update.Message.Chat.ID)
		}
	}
	b.chatMu.Unlock()

	// Validate message text
	if update.Message.Text == "" {
		return
	}

	// Tokenize the update
	tokens := strings.Split(update.Message.Text, " ")

	// Preprocess message based on URL schema
	if len(tokens) > 0 && (strings.HasPrefix(tokens[0], "magnet") || strings.HasPrefix(tokens[0], "http")) {
		tokens = append([]string{"add"}, tokens...)
	}

	command := strings.ToLower(tokens[0])

	// Route to appropriate handler with panic recovery
	switch command {
	case "list", "/list", "li", "/li", "/ls", "ls":
		go b.safeHandler(func() { b.handleList(update, tokens[1:]) })
	case "head", "/head", "he", "/he":
		go b.safeHandler(func() { b.handleHead(update, tokens[1:]) })
	case "tail", "/tail", "ta", "/ta":
		go b.safeHandler(func() { b.handleTail(update, tokens[1:]) })
	case "downs", "/downs", "dg", "/dg":
		go b.safeHandler(func() { b.handleDowns(update) })
	case "seeding", "/seeding", "sd", "/sd":
		go b.safeHandler(func() { b.handleSeeding(update) })
	case "paused", "/paused", "pa", "/pa":
		go b.safeHandler(func() { b.handlePaused(update) })
	case "checking", "/checking", "ch", "/ch":
		go b.safeHandler(func() { b.handleChecking(update) })
	case "active", "/active", "ac", "/ac":
		go b.safeHandler(func() { b.handleActive(update) })
	case "errors", "/errors", "er", "/er":
		go b.safeHandler(func() { b.handleErrors(update) })
	case "sort", "/sort", "so", "/so":
		go b.safeHandler(func() { b.handleSort(update, tokens[1:]) })
	case "trackers", "/trackers", "tr", "/tr":
		go b.safeHandler(func() { b.handleTrackers(update) })
	case "downloaddir", "dd":
		go b.safeHandler(func() { b.handleDownloadDir(update, tokens[1:]) })
	case "add", "/add", "ad", "/ad":
		go b.safeHandler(func() { b.handleAdd(update, tokens[1:]) })
	case "search", "/search", "se", "/se":
		go b.safeHandler(func() { b.handleSearch(update, tokens[1:]) })
	case "latest", "/latest", "la", "/la":
		go b.safeHandler(func() { b.handleLatest(update, tokens[1:]) })
	case "info", "/info", "in", "/in":
		go b.safeHandler(func() { b.handleInfo(update, tokens[1:]) })
	case "stop", "/stop", "sp", "/sp":
		go b.safeHandler(func() { b.handleStop(update, tokens[1:]) })
	case "start", "/start", "st", "/st":
		go b.safeHandler(func() { b.handleStart(update, tokens[1:]) })
	case "check", "/check", "ck", "/ck":
		go b.safeHandler(func() { b.handleCheck(update, tokens[1:]) })
	case "stats", "/stats", "sa", "/sa":
		go b.safeHandler(func() { b.handleStats(update) })
	case "downlimit", "dl":
		go b.safeHandler(func() { b.handleDownLimit(update, tokens[1:]) })
	case "uplimit", "ul":
		go b.safeHandler(func() { b.handleUpLimit(update, tokens[1:]) })
	case "speed", "/speed", "ss", "/ss":
		go b.safeHandler(func() { b.handleSpeed(update) })
	case "count", "/count", "co", "/co":
		go b.safeHandler(func() { b.handleCount(update) })
	case "del", "/del", "rm", "/rm":
		go b.safeHandler(func() { b.handleDel(update, tokens[1:]) })
	case "deldata", "/deldata":
		go b.safeHandler(func() { b.handleDelData(update, tokens[1:]) })
	case "help", "/help":
		go b.safeHandler(func() { b.sendMessage(update.Message.Chat.ID, HelpText, true) })
	case "version", "/version", "ver", "/ver":
		go b.safeHandler(func() { b.handleVersion(update) })
	case "":
		// might be a file received
		go b.safeHandler(func() { b.handleReceiveTorrent(update) })
	default:
		// no such command
		go b.safeHandler(func() { b.sendMessage(update.Message.Chat.ID, "No such command, try /help", false) })
	}
}