176 lines
4.1 KiB
Go
176 lines
4.1 KiB
Go
package bot
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"os"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/pyed/transmission"
|
|
tgbotapi "gopkg.in/telegram-bot-api.v4"
|
|
"transmission-telegram/internal/config"
|
|
)
|
|
|
|
// rateLimiter limits API calls to Telegram (30 messages per second)
|
|
// Uses token bucket algorithm for accurate rate limiting
|
|
type rateLimiter struct {
|
|
tokens float64
|
|
maxTokens float64
|
|
refillRate float64 // tokens per second
|
|
lastRefill time.Time
|
|
mu sync.Mutex
|
|
}
|
|
|
|
var (
|
|
// telegramRateLimiter limits API calls to 30 per second
|
|
telegramRateLimiter = func() *rateLimiter {
|
|
rl := &rateLimiter{
|
|
tokens: 30.0, // Start with full bucket
|
|
maxTokens: 30.0, // Maximum 30 tokens
|
|
refillRate: 30.0, // 30 tokens per second
|
|
lastRefill: time.Now(),
|
|
}
|
|
return rl
|
|
}()
|
|
)
|
|
|
|
// wait waits for the next available slot using token bucket algorithm
|
|
func (rl *rateLimiter) wait() {
|
|
rl.mu.Lock()
|
|
|
|
// Refill tokens based on elapsed time
|
|
now := time.Now()
|
|
elapsed := now.Sub(rl.lastRefill).Seconds()
|
|
rl.tokens = min(rl.maxTokens, rl.tokens+elapsed*rl.refillRate)
|
|
rl.lastRefill = now
|
|
|
|
// If we have tokens, consume one and return immediately
|
|
if rl.tokens >= 1.0 {
|
|
rl.tokens -= 1.0
|
|
rl.mu.Unlock()
|
|
return
|
|
}
|
|
|
|
// Calculate how long to wait for the next token
|
|
waitTime := (1.0 - rl.tokens) / rl.refillRate
|
|
rl.mu.Unlock()
|
|
|
|
// Wait for the token to become available
|
|
time.Sleep(time.Duration(waitTime * float64(time.Second)))
|
|
|
|
rl.mu.Lock()
|
|
rl.tokens = 0.0 // Consume the token
|
|
rl.lastRefill = time.Now()
|
|
rl.mu.Unlock()
|
|
}
|
|
|
|
// min returns the minimum of two float64 values
|
|
func min(a, b float64) float64 {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|
|
|
|
// rateLimitWait waits for the next available API call slot
|
|
func rateLimitWait() {
|
|
telegramRateLimiter.wait()
|
|
}
|
|
|
|
// sendMessage sends a message to Telegram, splitting if necessary
|
|
func sendMessage(bot *tgbotapi.BotAPI, chatID int64, text string, markdown bool) int {
|
|
// Rate limit: wait before sending
|
|
telegramRateLimiter.wait()
|
|
|
|
// Set typing action
|
|
action := tgbotapi.NewChatAction(chatID, tgbotapi.ChatTyping)
|
|
if _, err := bot.Send(action); err != nil {
|
|
// Log error but continue - typing action is not critical
|
|
fmt.Fprintf(os.Stderr, "[WARN] Failed to send typing action: %v\n", err)
|
|
}
|
|
|
|
var lastMsgID int
|
|
|
|
// Convert to runes for proper UTF-8 handling
|
|
runes := []rune(text)
|
|
|
|
// Split message if too long
|
|
for len(runes) > config.TelegramMaxMessageLength {
|
|
stop := config.TelegramMaxMessageLength - 1
|
|
|
|
// Find the last newline before the limit
|
|
for stop > 0 && runes[stop] != '\n' {
|
|
stop--
|
|
}
|
|
|
|
// If no newline found, just cut at the limit
|
|
if stop == 0 {
|
|
stop = config.TelegramMaxMessageLength - 1
|
|
}
|
|
|
|
chunk := string(runes[:stop+1])
|
|
runes = runes[stop+1:]
|
|
|
|
telegramRateLimiter.wait()
|
|
msg := tgbotapi.NewMessage(chatID, chunk)
|
|
msg.DisableWebPagePreview = true
|
|
if markdown {
|
|
msg.ParseMode = tgbotapi.ModeMarkdown
|
|
}
|
|
|
|
resp, err := bot.Send(msg)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "[ERROR] Failed to send message chunk: %v\n", err)
|
|
} else {
|
|
lastMsgID = resp.MessageID
|
|
}
|
|
}
|
|
|
|
// Send remaining text
|
|
if len(runes) > 0 {
|
|
telegramRateLimiter.wait()
|
|
msg := tgbotapi.NewMessage(chatID, string(runes))
|
|
msg.DisableWebPagePreview = true
|
|
if markdown {
|
|
msg.ParseMode = tgbotapi.ModeMarkdown
|
|
}
|
|
|
|
resp, err := bot.Send(msg)
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "[ERROR] Failed to send message: %v\n", err)
|
|
} else {
|
|
lastMsgID = resp.MessageID
|
|
}
|
|
}
|
|
|
|
return lastMsgID
|
|
}
|
|
|
|
// parseTorrentIDs parses torrent IDs from tokens
|
|
func parseTorrentIDs(tokens []string) ([]int, []string) {
|
|
var ids []int
|
|
var errors []string
|
|
|
|
for _, token := range tokens {
|
|
if id, err := strconv.Atoi(token); err == nil {
|
|
ids = append(ids, id)
|
|
} else {
|
|
errors = append(errors, fmt.Sprintf("%s is not a number", token))
|
|
}
|
|
}
|
|
|
|
return ids, errors
|
|
}
|
|
|
|
// formatTorrentList formats a list of torrents as a simple list
|
|
func formatTorrentList(torrents []*transmission.Torrent) string {
|
|
var buf bytes.Buffer
|
|
for _, torrent := range torrents {
|
|
buf.WriteString(fmt.Sprintf("<%d> %s\n", torrent.ID, torrent.Name))
|
|
}
|
|
return buf.String()
|
|
}
|
|
|