uptodownbot/internal/handler/handler.go

package handler

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	tgbot "github.com/go-telegram/bot"
	"github.com/go-telegram/bot/models"
	"github.com/google/uuid"

	"uptodownBot/internal/dl"
	"uptodownBot/internal/domain"
	"uptodownBot/internal/repo"
)

type PendingKind string

const (
	PendingURL    PendingKind = "url"
	PendingDRM    PendingKind = "drm"
	PendingSearch PendingKind = "search"
)

type PendingInput struct {
	Kind      PendingKind
	UserID    int64
	CreatedAt time.Time
	MsgID     int
	Data      map[string]string
}

const maxConcurrentDownloads = 3

type Handler struct {
	Bot          *tgbot.Bot
	Repo         repo.Repository
	YtDlp        *dl.Client
	AllowedUsers map[int64]bool

	activeJobs   map[int64]*domain.DownloadJob
	stepStates   map[int64]*stepState
	lastMsgTime  map[int64]time.Time
	pendingInput map[int64]*PendingInput
	cancelFuncs  map[int64]context.CancelFunc

	activeJobsMu   sync.Mutex
	stepStatesMu   sync.Mutex
	rateLimiterMu  sync.Mutex
	pendingInputMu sync.Mutex
	cancelFuncsMu  sync.Mutex

	downloadDir string
	maxFileSize int64
	downloadSem chan struct{}
}

func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed map[int64]bool, downloadDir string, maxFileSize int64) *Handler {
	h := &Handler{
		Bot:          bot,
		Repo:         store,
		YtDlp:        ytdlp,
		AllowedUsers: allowed,
		activeJobs:   make(map[int64]*domain.DownloadJob),
		stepStates:   make(map[int64]*stepState),
		lastMsgTime:  make(map[int64]time.Time),
		pendingInput: make(map[int64]*PendingInput),
		cancelFuncs:  make(map[int64]context.CancelFunc),
		downloadDir:  downloadDir,
		maxFileSize:  maxFileSize,
		downloadSem:  make(chan struct{}, maxConcurrentDownloads),
	}
	go h.periodicCleanup()
	return h
}

func (h *Handler) periodicCleanup() {
	for {
		time.Sleep(domain.CleanupInterval)

		func() {
			h.rateLimiterMu.Lock()
			defer h.rateLimiterMu.Unlock()
			cutoff := time.Now().Add(-domain.RateLimitInterval * 2)
			for uid, t := range h.lastMsgTime {
				if t.Before(cutoff) {
					delete(h.lastMsgTime, uid)
				}
			}
		}()

		func() {
			h.pendingInputMu.Lock()
			defer h.pendingInputMu.Unlock()
			pendingCutoff := time.Now().Add(-5 * time.Minute)
			for chatID, pi := range h.pendingInput {
				if pi.CreatedAt.Before(pendingCutoff) {
					delete(h.pendingInput, chatID)
				}
			}
		}()

		h.cleanOldTempFiles()
	}
}

func (h *Handler) cleanOldTempFiles() {
	entries, err := os.ReadDir(h.downloadDir)
	if err != nil {
		return
	}
	cutoff := time.Now().Add(-domain.FileRetention)
	for _, entry := range entries {
		info, err := entry.Info()
		if err != nil {
			continue
		}
		if info.ModTime().Before(cutoff) {
			path := filepath.Join(h.downloadDir, entry.Name())
			if err := os.Remove(path); err != nil {
				slog.Warn("failed to remove old temp file", "path", path, "error", err)
			}
		}
	}
}

func (h *Handler) isRateLimited(userID int64) bool {
	h.rateLimiterMu.Lock()
	defer h.rateLimiterMu.Unlock()
	now := time.Now()
	if last, ok := h.lastMsgTime[userID]; ok && now.Sub(last) < domain.RateLimitInterval {
		return true
	}
	h.lastMsgTime[userID] = now
	return false
}

func (h *Handler) getOrCreateUser(chatID int64) (int64, error) {
	return h.Repo.GetOrCreateUser(chatID)
}

func (h *Handler) HandleMessage(ctx context.Context, msg *models.Message) {
	if len(h.AllowedUsers) > 0 && !h.AllowedUsers[msg.Chat.ID] {
		return
	}
	if msg.Text == "" {
		return
	}
	userID, err := h.getOrCreateUser(msg.Chat.ID)
	if err != nil {
		return
	}
	if h.isRateLimited(userID) {
		return
	}
	slog.Info("message", "chat_id", msg.Chat.ID, "text", msg.Text)

	if msg.Text[0] == '/' {
		h.pendingInputMu.Lock()
		delete(h.pendingInput, msg.Chat.ID)
		h.pendingInputMu.Unlock()
		h.routeCommand(ctx, msg, userID)
		return
	}

	h.pendingInputMu.Lock()
	pi, hasPI := h.pendingInput[msg.Chat.ID]
	if hasPI {
		delete(h.pendingInput, msg.Chat.ID)
	}
	h.pendingInputMu.Unlock()

	if hasPI {
		h.processPendingInput(ctx, msg, userID, pi)
		return
	}

	// Treat plain text as a potential URL
	h.handleURLInput(ctx, msg.Chat.ID, msg.Text, userID)
}

func (h *Handler) routeCommand(ctx context.Context, msg *models.Message, userID int64) {
	cmd := msg.Text
	if i := strings.IndexByte(msg.Text, ' '); i >= 0 {
		cmd = msg.Text[:i]
	}
	switch cmd {
	case "/start":
		h.handleStart(ctx, msg, userID)
	case "/settings":
		h.handleSettings(ctx, msg, userID)
	case "/help":
		h.handleHelp(ctx, msg)
	case "/history":
		h.handleHistory(ctx, msg, userID)
	case "/download":
		h.handleDownloadCommand(ctx, msg, userID)
	case "/search":
		h.handleSearchCommand(ctx, msg, userID)
	case "/music":
		h.handleMusicSearchCommand(ctx, msg, userID)
	default:
		h.sendText(ctx, msg.Chat.ID, "Unknown command")
	}
}

func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery) {
	chatID := cb.Message.Message.Chat.ID
	if len(h.AllowedUsers) > 0 && !h.AllowedUsers[chatID] {
		h.answerCbWithText(ctx, cb.ID, "Unauthorized")
		return
	}
	userID, err := h.getOrCreateUser(chatID)
	if err != nil {
		h.answerCb(ctx, cb.ID)
		return
	}
	if h.isRateLimited(userID) {
		h.answerCb(ctx, cb.ID)
		return
	}
	slog.Info("callback", "chat_id", chatID, "data", cb.Data)

	h.pendingInputMu.Lock()
	delete(h.pendingInput, chatID)
	h.pendingInputMu.Unlock()

	msgID := cb.Message.Message.ID
	data := cb.Data

	switch {
	case data == "noop":
		h.answerCb(ctx, cb.ID)
	case data == "cancel_dl":
		h.handleCancelDownload(ctx, chatID, msgID, cb.ID, userID)
	case data == "back_step":
		h.handleBackStep(ctx, chatID, msgID, cb.ID, userID)
	case data == "back_menu":
		h.backToMenu(ctx, chatID, msgID, cb.ID, userID)
	case data == "download":
		h.handleDownloadPrompt(ctx, chatID, msgID, cb.ID, userID)
		h.answerCb(ctx, cb.ID)
	case data == "search":
		h.promptForInput(ctx, chatID, msgID, userID, PendingSearch, "Enter a YouTube search query:", map[string]string{"source": "youtube"})
		h.answerCb(ctx, cb.ID)
	case data == "settings":
		h.settingsCallback(ctx, chatID, msgID, cb.ID, userID)
	case data == "help":
		h.handleHelp(ctx, cb.Message.Message)
		h.answerCb(ctx, cb.ID)
	case data == "history":
		h.historyCallback(ctx, chatID, msgID, cb.ID, userID)
	case data == "deleteaccount":
		h.deleteAccountPrompt(ctx, chatID, msgID, userID)
		h.answerCb(ctx, cb.ID)
	case data == "drm_search":
		h.promptForInput(ctx, chatID, msgID, userID, PendingDRM, "Enter a search query to find this content on YouTube:", nil)
		h.answerCb(ctx, cb.ID)
	case data == "pending_cancel":
		h.backToMenu(ctx, chatID, msgID, cb.ID, userID)
	default:
		h.routePrefixedCallback(ctx, chatID, msgID, cb.ID, userID, data)
	}
}

func (h *Handler) routePrefixedCallback(ctx context.Context, chatID int64, msgID int, cbID string, userID int64, data string) {
	defer h.answerCb(ctx, cbID)

	switch {
	case strings.HasPrefix(data, "format_"):
		h.handleFormatSelection(ctx, chatID, msgID, userID, data[7:])
	case strings.HasPrefix(data, "secondary_"):
		h.handleSecondarySelection(ctx, chatID, msgID, userID, data[10:])
	case strings.HasPrefix(data, "lang_"):
		h.handleLanguageSelection(ctx, chatID, msgID, userID, data[5:])
	case strings.HasPrefix(data, "container_"):
		// container selection was removed — yt-dlp handles it
	case strings.HasPrefix(data, "delaccount_"):
		if uid, err := strconv.ParseInt(data[11:], 10, 64); err == nil && uid > 0 {
			h.deleteAccountConfirm(ctx, chatID, msgID, uid)
		}
	case data == "back_quality":
		h.handleSettingsSelection(ctx, chatID, msgID, userID, "quality")
	case data == "back_settings":
		h.backToSettings(ctx, chatID, msgID, userID)
	case strings.HasPrefix(data, "settings_set_"):
		h.handleSettingsSet(ctx, chatID, msgID, userID, data[13:])
	case strings.HasPrefix(data, "settings_"):
		h.handleSettingsSelection(ctx, chatID, msgID, userID, data[9:])
	case strings.HasPrefix(data, "search_"):
		h.handleSearchCallback(ctx, chatID, msgID, userID, data[7:])
	}
}

func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, data map[string]string) {
	h.pendingInputMu.Lock()
	h.pendingInput[chatID] = &PendingInput{
		Kind:      kind,
		UserID:    userID,
		CreatedAt: time.Now(),
		MsgID:     msgID,
		Data:      data,
	}
	h.pendingInputMu.Unlock()
}

func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message, userID int64, pi *PendingInput) {
	text := strings.TrimSpace(msg.Text)
	switch pi.Kind {
	case PendingURL:
		h.handleURLInput(ctx, msg.Chat.ID, text, userID)
	case PendingDRM:
		h.handleDRMSearch(ctx, msg.Chat.ID, text, userID)
	case PendingSearch:
		source := "youtube"
		if pi.Data != nil && pi.Data["source"] != "" {
			source = pi.Data["source"]
		}
		h.handleSearchQuery(ctx, msg.Chat.ID, text, userID, source)
	default:
		h.sendText(ctx, msg.Chat.ID, "Unknown input type")
	}
}

func (h *Handler) handleDRMSearch(ctx context.Context, chatID int64, text string, userID int64) {
	youtubeURL, ytTitle, err := h.YtDlp.SearchYoutube(text)
	if err != nil {
		h.sendText(ctx, chatID, fmt.Sprintf("No results found for '%s'. Try a different search query or send a YouTube URL directly.", text))
		return
	}
	h.sendText(ctx, chatID, fmt.Sprintf("Found on YouTube: %s", ytTitle))
	h.handleURLInput(ctx, chatID, youtubeURL, userID)
}

func (h *Handler) promptForInput(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, prompt string, data map[string]string) {
	kb := models.InlineKeyboardMarkup{
		InlineKeyboard: [][]models.InlineKeyboardButton{
			{{Text: "Cancel", CallbackData: "pending_cancel"}},
		},
	}
	if msgID == 0 {
		msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{
			ChatID: chatID, Text: prompt, ReplyMarkup: kb,
		})
		if err != nil {
			slog.Warn("send prompt failed", "chat_id", chatID, "error", err)
			return
		}
		h.setPending(ctx, chatID, msg.ID, userID, kind, data)
	} else {
		h.editText(ctx, chatID, msgID, prompt, &kb)
		h.setPending(ctx, chatID, msgID, userID, kind, data)
	}
}

func (h *Handler) sendOrEdit(ctx context.Context, chatID int64, msgID int, text string, kb *models.InlineKeyboardMarkup) {
	if msgID == 0 {
		if kb != nil {
			h.sendWithKB(ctx, chatID, text, *kb)
		} else {
			h.sendText(ctx, chatID, text)
		}
	} else {
		h.editText(ctx, chatID, msgID, text, kb)
	}
}

// -- helpers --

func (h *Handler) sendText(ctx context.Context, chatID int64, text string) {
	if _, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ChatID: chatID, Text: text}); err != nil {
		slog.Warn("send text failed", "chat_id", chatID, "error", err)
	}
}

func (h *Handler) editText(ctx context.Context, chatID int64, msgID int, text string, kb *models.InlineKeyboardMarkup) {
	p := &tgbot.EditMessageTextParams{ChatID: chatID, MessageID: msgID, Text: text}
	if kb != nil {
		p.ReplyMarkup = kb
	}
	if _, err := h.Bot.EditMessageText(ctx, p); err != nil {
		slog.Warn("edit text failed", "chat_id", chatID, "msg_id", msgID, "error", err)
	}
}

func (h *Handler) sendWithKB(ctx context.Context, chatID int64, text string, kb models.InlineKeyboardMarkup) {
	if _, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ChatID: chatID, Text: text, ReplyMarkup: kb}); err != nil {
		slog.Warn("send with keyboard failed", "chat_id", chatID, "error", err)
	}
}

func (h *Handler) answerCb(ctx context.Context, callbackID string) {
	if _, err := h.Bot.AnswerCallbackQuery(ctx, &tgbot.AnswerCallbackQueryParams{CallbackQueryID: callbackID}); err != nil {
		slog.Debug("answer callback failed", "error", err)
	}
}

func (h *Handler) answerCbWithText(ctx context.Context, callbackID, text string) {
	if _, err := h.Bot.AnswerCallbackQuery(ctx, &tgbot.AnswerCallbackQueryParams{CallbackQueryID: callbackID, Text: text}); err != nil {
		slog.Debug("answer callback with text failed", "error", err)
	}
}

func (h *Handler) backToMenu(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
	defer h.answerCb(ctx, cbID)
	h.clearActiveJob(userID)
	h.editText(ctx, chatID, msgID, "Main Menu:", &models.InlineKeyboardMarkup{
		InlineKeyboard: mainKeyboard().InlineKeyboard,
	})
	h.clearStepState(chatID)
}

func (h *Handler) clearActiveJob(userID int64) {
	h.activeJobsMu.Lock()
	job, ok := h.activeJobs[userID]
	if ok {
		dl.CancelDownload(job)
		delete(h.activeJobs, userID)
	}
	h.activeJobsMu.Unlock()

	h.cancelFuncsMu.Lock()
	if cancel, ok := h.cancelFuncs[userID]; ok {
		cancel()
		delete(h.cancelFuncs, userID)
	}
	h.cancelFuncsMu.Unlock()
}

func (h *Handler) getActiveJob(userID int64) *domain.DownloadJob {
	h.activeJobsMu.Lock()
	defer h.activeJobsMu.Unlock()
	return h.activeJobs[userID]
}

func (h *Handler) setActiveJob(userID int64, job *domain.DownloadJob) {
	h.activeJobsMu.Lock()
	defer h.activeJobsMu.Unlock()
	h.activeJobs[userID] = job
}

func (h *Handler) setCancelFunc(userID int64, cancel context.CancelFunc) {
	h.cancelFuncsMu.Lock()
	defer h.cancelFuncsMu.Unlock()
	h.cancelFuncs[userID] = cancel
}

func formatBytes(bytes int64) string {
	const unit = 1024
	if bytes < unit {
		return strconv.FormatInt(bytes, 10) + "B"
	}
	units := []string{"KB", "MB", "GB", "TB"}
	size := float64(bytes)
	for _, u := range units {
		size /= unit
		if size < unit {
			return fmt.Sprintf("%.1f%s", size, u)
		}
	}
	return fmt.Sprintf("%.1fTB", size)
}

func (h *Handler) generateJobID() string {
	return uuid.NewString()
}

// handleStart shows the main menu and cleans up any stale state.
func (h *Handler) handleStart(ctx context.Context, msg *models.Message, userID int64) {
	slog.Info("start", "user_id", userID, "chat_id", msg.Chat.ID)
	h.clearActiveJob(userID)
	h.clearStepState(msg.Chat.ID)
	h.sendWithKB(ctx, msg.Chat.ID, "Main Menu:", mainKeyboard())
}

// handleHelp shows the help text.
func (h *Handler) handleHelp(ctx context.Context, msg *models.Message) {
	text := `Commands:
/start - Show main menu
/settings - Open settings
/help - Show this help
/history - View download history
/search - Search YouTube
/music - Search YouTube Music

Just send me a URL to start downloading.`
	h.sendText(ctx, msg.Chat.ID, text)
}

// handleSettings opens the settings menu.
func (h *Handler) handleSettings(ctx context.Context, msg *models.Message, userID int64) {
	prefs, err := h.Repo.GetOrCreatePreferences(userID)
	if err != nil {
		h.sendText(ctx, msg.Chat.ID, "Error loading settings")
		return
	}
	text, kb := h.buildSettingsKeyboard(userID, prefs)
	h.sendWithKB(ctx, msg.Chat.ID, text, kb)
}

// handleHistory opens the download history.
func (h *Handler) handleHistory(ctx context.Context, msg *models.Message, userID int64) {
	records, err := h.Repo.GetHistory(userID, 10, 0)
	if err != nil {
		h.sendText(ctx, msg.Chat.ID, "Error loading history.")
		return
	}
	if len(records) == 0 {
		h.sendText(ctx, msg.Chat.ID, "No download history yet.")
		return
	}
	text := "Recent downloads:\n"
	for i, r := range records {
		text += fmt.Sprintf("%d. %s (%s)\n", i+1, r.Title, formatBytes(r.FileSize))
	}
	h.sendText(ctx, msg.Chat.ID, text)
}

// historyCallback handles the history button from the main menu.
func (h *Handler) historyCallback(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
	defer h.answerCb(ctx, cbID)
	records, err := h.Repo.GetHistory(userID, 10, 0)
	if err != nil {
		h.editText(ctx, chatID, msgID, "Error loading history.", &models.InlineKeyboardMarkup{
			InlineKeyboard: [][]models.InlineKeyboardButton{
				{{Text: "Back to Menu", CallbackData: "back_menu"}},
			},
		})
		return
	}
	if len(records) == 0 {
		h.editText(ctx, chatID, msgID, "No download history yet.", &models.InlineKeyboardMarkup{
			InlineKeyboard: [][]models.InlineKeyboardButton{
				{{Text: "Back to Menu", CallbackData: "back_menu"}},
			},
		})
		return
	}
	text := "Recent downloads:\n"
	for i, r := range records {
		text += fmt.Sprintf("%d. %s (%s)\n", i+1, r.Title, formatBytes(r.FileSize))
	}
	h.editText(ctx, chatID, msgID, text, &models.InlineKeyboardMarkup{
		InlineKeyboard: [][]models.InlineKeyboardButton{
			{{Text: "Back to Menu", CallbackData: "back_menu"}},
		},
	})
}

// applyQuota adds file size to the user's quota counters.
func (h *Handler) applyQuota(userID int64, fileSize int64) {
	q, err := h.Repo.GetOrCreateQuotas(userID)
	if err != nil {
		slog.Error("get quotas for apply", "user_id", userID, "error", err)
		return
	}
	q = h.resetQuotasIfNeeded(q)
	sizeMB := fileSize / (1024 * 1024)
	if sizeMB < 1 {
		sizeMB = 1
	}
	q.DailyUsedMB += sizeMB
	q.WeeklyUsedMB += sizeMB
	q.MonthlyUsedMB += sizeMB
	if err := h.Repo.UpdateQuotas(userID, q); err != nil {
		slog.Error("update quotas", "user_id", userID, "error", err)
	}
}

// resetQuotasIfNeeded resets quota counters if the period has changed.
func (h *Handler) resetQuotasIfNeeded(q *domain.UserQuotas) *domain.UserQuotas {
	today := time.Now().UTC().Format(domain.DateLayout)
	_, weekNum := time.Now().UTC().ISOWeek()
	weekStart := fmt.Sprintf("%d-W%02d", time.Now().UTC().Year(), weekNum)
	monthStart := time.Now().UTC().Format("2006-01")

	if q.DailyDate != today {
		q.DailyUsedMB = 0
		q.DailyDate = today
	}
	if q.WeekStart != weekStart {
		q.WeeklyUsedMB = 0
		q.WeekStart = weekStart
	}
	if q.MonthStart != monthStart {
		q.MonthlyUsedMB = 0
		q.MonthStart = monthStart
	}
	return q
}