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"
)

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

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
	mu           sync.Mutex
	downloadDir  string
	maxFileSize  int64
}

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),
		downloadDir:  downloadDir,
		maxFileSize:  maxFileSize,
	}
	go h.periodicCleanup()
	return h
}

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

		// Clean rate limit entries
		cutoff := time.Now().Add(-domain.RateLimitInterval * 2)
		for uid, t := range h.lastMsgTime {
			if t.Before(cutoff) {
				delete(h.lastMsgTime, uid)
			}
		}

		// Clean stale pending inputs
		pendingCutoff := time.Now().Add(-5 * time.Minute)
		for chatID, pi := range h.pendingInput {
			if pi.CreatedAt.Before(pendingCutoff) {
				delete(h.pendingInput, chatID)
			}
		}

		h.mu.Unlock()

		// Clean old temp files (>1h)
		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.mu.Lock()
	defer h.mu.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.mu.Lock()
		delete(h.pendingInput, msg.Chat.ID)
		h.mu.Unlock()
		h.routeCommand(ctx, msg, userID)
		return
	}

	h.mu.Lock()
	pi, hasPI := h.pendingInput[msg.Chat.ID]
	if hasPI {
		delete(h.pendingInput, msg.Chat.ID)
	}
	h.mu.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)
	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.mu.Lock()
	delete(h.pendingInput, chatID)
	h.mu.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 == "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.handleHistory(ctx, cb.Message.Message, userID)
		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, "type_"):
		h.handleTypeSelection(ctx, chatID, msgID, userID, data[5:])
	case strings.HasPrefix(data, "quality_"):
		h.handleQualitySelection(ctx, chatID, msgID, userID, data[8:])
	case strings.HasPrefix(data, "lang_"):
		h.handleLanguageSelection(ctx, chatID, msgID, userID, data[5:])
	case strings.HasPrefix(data, "container_"):
		h.handleContainerSelection(ctx, chatID, msgID, userID, data[10:])
	case strings.HasPrefix(data, "pl_page_"):
		h.handlePlaylistPage(ctx, chatID, msgID, userID, data[8:])
	case strings.HasPrefix(data, "pl_entry_"):
		h.handlePlaylistEntry(ctx, chatID, msgID, userID, data[9:])
	case data == "pl_select_all":
		h.handlePlaylistSelectAll(ctx, chatID, msgID, userID)
	case data == "pl_confirm":
		h.handlePlaylistConfirm(ctx, chatID, msgID, userID)
	case data == "back_settings":
		h.backToSettings(ctx, chatID, msgID, userID)
	case data == "delaccount_prompt":
		h.deleteAccountPrompt(ctx, chatID, msgID, userID)
	case strings.HasPrefix(data, "delaccount_"):
		h.deleteAccountConfirm(ctx, chatID, msgID, userID)
	case strings.HasPrefix(data, "settings_"):
		h.handleSettingsSelection(ctx, chatID, msgID, userID, data[9:])
	case strings.HasPrefix(data, "settings_set_"):
		h.handleSettingsSet(ctx, chatID, msgID, userID, data[13:])
	}
}

func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, data map[string]string) {
	h.mu.Lock()
	h.pendingInput[chatID] = &PendingInput{
		Kind:      kind,
		UserID:    userID,
		CreatedAt: time.Now(),
		MsgID:     msgID,
		Data:      data,
	}
	h.mu.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)
	default:
		h.sendText(ctx, msg.Chat.ID, "Unknown input type")
	}
}

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:", nil)
	h.sendWithKB(ctx, chatID, "Main Menu:", mainKeyboard())
}

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

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

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

func formatBytes(bytes int64) string {
	const unit = 1024
	if bytes < unit {
		return strconv.FormatInt(bytes, 10) + "B"
	}
	div, exp := int64(unit), 0
	for n := bytes / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	switch exp {
	case 0:
		return fmt.Sprintf("%.1fKB", float64(bytes)/float64(div))
	case 1:
		return fmt.Sprintf("%.1fMB", float64(bytes)/float64(div))
	case 2:
		return fmt.Sprintf("%.1fGB", float64(bytes)/float64(div))
	default:
		return fmt.Sprintf("%.1fTB", float64(bytes)/float64(div))
	}
}

func (h *Handler) generateJobID() string {
	return uuid.NewString()[:8]
}

// handleStart shows the main menu.
func (h *Handler) handleStart(ctx context.Context, msg *models.Message, userID int64) {
	slog.Info("start", "user_id", userID, "chat_id", 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

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)
}

// 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 {
		return
	}
	q = h.resetQuotasIfNeeded(q)
	sizeMB := fileSize / (1024 * 1024)
	if sizeMB < 1 {
		sizeMB = 1
	}
	q.DailyUsedMB += sizeMB
	q.WeeklyUsedMB += sizeMB
	q.MonthlyUsedMB += sizeMB
	_ = h.Repo.UpdateQuotas(userID, q)
}

// 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
}