Ztorrent/internal/tgbot/bot.go

package tgbot

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"

	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
	"github.com/veggiedefender/torrent-client/internal/config"
	"github.com/veggiedefender/torrent-client/internal/logger"
)

type Controller interface {
	DownloadMagnet(magnetURI string) error
	// You can add more methods if you want to support downloading .torrent files or getting status.
}

type Bot struct {
	api        *tgbotapi.BotAPI
	controller Controller
}

func NewBot(ctrl Controller, token string) (*Bot, error) {
	if token == "" {
		// Bot is disabled if no token is provided.
		logger.Info("TGBOT", "TG_BOT_TOKEN not set, bot integration is disabled")
		return nil, nil
	}

	api, err := tgbotapi.NewBotAPI(token)
	if err != nil {
		return nil, fmt.Errorf("failed to create bot: %w", err)
	}

	logger.Info("TGBOT", "Authorized on account %s", api.Self.UserName)

	return &Bot{
		api:        api,
		controller: ctrl,
	}, nil
}

func (b *Bot) Start() {
	if b == nil || b.api == nil {
		return
	}

	u := tgbotapi.NewUpdate(0)
	u.Timeout = 60

	updates := b.api.GetUpdatesChan(u)

	go func() {
		cfg, _ := config.Load()
		var ownerID int64
		if cfg != nil {
			ownerID = cfg.TelegramOwnerID
		}

		for update := range updates {
			if update.Message == nil {
				continue
			}

			// Authorization Check
			if ownerID == 0 {
				if update.Message.Text == "/start" {
					ownerID = update.Message.From.ID
					if cfg == nil {
						cfg = &config.Config{}
					}
					cfg.TelegramOwnerID = ownerID
					_ = config.Save(cfg)
					b.api.Send(tgbotapi.NewMessage(update.Message.Chat.ID, "You are now registered as the owner of this Ztorrent instance."))
				} else {
					b.api.Send(tgbotapi.NewMessage(update.Message.Chat.ID, "Ztorrent is waiting for the owner to send /start."))
				}
				continue
			} else if update.Message.From.ID != ownerID {
				logger.Warn("TGBOT", "Unauthorized access attempt from %s", update.Message.From.UserName)
				continue
			}

			// Check for documents (.torrent files)
			if update.Message.Document != nil {
				doc := update.Message.Document
				if strings.HasSuffix(strings.ToLower(doc.FileName), ".torrent") {
					logger.Info("TGBOT", "Received .torrent file from %s: %s", update.Message.From.UserName, doc.FileName)
					
					// Get file URL from Telegram
					fileURL, err := b.api.GetFileDirectURL(doc.FileID)
					if err != nil {
						b.api.Send(tgbotapi.NewMessage(update.Message.Chat.ID, fmt.Sprintf("Error getting file: %v", err)))
						continue
					}

					b.api.Send(tgbotapi.NewMessage(update.Message.Chat.ID, "Torrent file received, starting download..."))
					
					// Download file contents async
					go func(chatId int64, url, name string) {
						err := b.downloadAndStartTorrent(chatId, url, name)
						if err != nil {
							b.api.Send(tgbotapi.NewMessage(chatId, fmt.Sprintf("Error starting torrent: %v", err)))
						} else {
							b.api.Send(tgbotapi.NewMessage(chatId, "Download started successfully!"))
						}
					}(update.Message.Chat.ID, fileURL, doc.FileName)
					
					continue
				}
			}

			text := strings.TrimSpace(update.Message.Text)
			
			if strings.HasPrefix(text, "magnet:?") {
				logger.Info("TGBOT", "Received magnet link from %s", update.Message.From.UserName)
				msg := tgbotapi.NewMessage(update.Message.Chat.ID, "Magnet link received, starting download...")
				b.api.Send(msg)
				
				err := b.controller.DownloadMagnet(text)
				if err != nil {
					errMsg := tgbotapi.NewMessage(update.Message.Chat.ID, fmt.Sprintf("Error starting download: %v", err))
					b.api.Send(errMsg)
				} else {
					succMsg := tgbotapi.NewMessage(update.Message.Chat.ID, "Download started successfully!")
					b.api.Send(succMsg)
				}
				continue
			}
			
			if text == "/start" {
				b.api.Send(tgbotapi.NewMessage(update.Message.Chat.ID, "Welcome back, Owner! Send me a magnet link or a .torrent file."))
			}
		}
	}()
}

func (b *Bot) downloadAndStartTorrent(chatID int64, fileURL, fileName string) error {
	resp, err := http.Get(fileURL)
	if err != nil {
		return fmt.Errorf("failed to download file from telegram: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("bad status code: %d", resp.StatusCode)
	}

	out, err := os.CreateTemp("", "*_"+fileName)
	if err != nil {
		return fmt.Errorf("failed to create temp file: %w", err)
	}
	tmpPath := out.Name()

	_, err = io.Copy(out, resp.Body)
	out.Close()
	if err != nil {
		return fmt.Errorf("failed to write temp file: %w", err)
	}

	// We pass the local path to controller
	return b.controller.DownloadMagnet(tmpPath)
}