uptodownbot/internal/handler/download.go

package handler

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

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

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

// Step state stored per user during the download flow
type stepState struct {
	URL           string
	MediaInfo     *domain.MediaInfo
	Step          int // 0=type, 1=quality, 2=language, 3=container
	MediaType     domain.MediaType
	Quality       string
	Language      string
	Container     string
	FormatIDs     []string
	PlaylistState *domain.PlaylistState
	URLMessageID  int // original URL message ID for editing
}

// getUserStepState retrieves the current step state for a user.
func (h *Handler) getUserStepState(chatID int64) *stepState {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.stepStates[chatID]
}

func (h *Handler) setUserStepState(chatID int64, ss *stepState) {
	h.mu.Lock()
	defer h.mu.Unlock()
	h.stepStates[chatID] = ss
}

func (h *Handler) clearStepState(chatID int64) {
	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.stepStates, chatID)
}

// handleURLInput is the entry point when a user sends a URL or types in Download mode.
func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string, userID int64) {
	url, ok := domain.ValidateURL(text)
	if !ok {
		h.sendText(ctx, chatID, "Please send a valid URL (http/https).")
		return
	}

	// Check for active job
	if job := h.getActiveJob(userID); job != nil {
		h.sendText(ctx, chatID, "You already have an active download. Cancel it first.")
		return
	}

	// Check quota
	info, err := h.YtDlp.GetMediaInfo(url)
	if err != nil {
		errMsg := err.Error()
		if strings.Contains(errMsg, "Unsupported URL") || strings.Contains(errMsg, "is not a valid URL") {
			h.sendText(ctx, chatID, "This URL is not supported.")
		} else if strings.Contains(errMsg, "HTTP Error 403") || strings.Contains(errMsg, "cookies") {
			h.sendText(ctx, chatID, "This content requires authentication (cookies).")
		} else if strings.Contains(errMsg, "live") {
			h.sendText(ctx, chatID, "Cannot download live streams.")
		} else {
			h.sendText(ctx, chatID, fmt.Sprintf("Error: %s", errMsg))
		}
		return
	}

	if info.IsPlaylist {
		h.handlePlaylist(ctx, chatID, userID, url, info)
		return
	}

	if len(info.Formats) == 0 {
		h.sendText(ctx, chatID, "No formats available for this URL.")
		return
	}

	// Check quota before starting flow
	q, err := h.Repo.GetOrCreateQuotas(userID)
	if err == nil {
		q = h.resetQuotasIfNeeded(q)
		estimatedMB := info.EstimatedSize / (1024 * 1024)
		if estimatedMB > 0 {
			if q.DailyUsedMB+estimatedMB > q.DailyLimitMB ||
				q.WeeklyUsedMB+estimatedMB > q.WeeklyLimitMB ||
				q.MonthlyUsedMB+estimatedMB > q.MonthlyLimitMB {
				h.sendText(ctx, chatID, fmt.Sprintf(
					"Quota exceeded. Daily: %d/%d MB, Weekly: %d/%d MB, Monthly: %d/%d MB.",
					q.DailyUsedMB, q.DailyLimitMB, q.WeeklyUsedMB, q.WeeklyLimitMB, q.MonthlyUsedMB, q.MonthlyLimitMB))
				return
			}
		}
	}

	ss := &stepState{
		URL:       url,
		MediaInfo: info,
		Step:      0,
	}
	h.setUserStepState(chatID, ss)

	// Show thumbnail and metadata preview
	h.sendMediaPreview(ctx, chatID, info)

	if info.IsAudioOnlyContent() {
		// Audio-only sites: skip type selection, go straight to quality
		ss.MediaType = domain.MediaTypeAudio
		ss.Step = 1
		filtered := dl.FilterFormatsByType(info.Formats, domain.MediaTypeAudio)
		if len(filtered) == 0 {
			h.sendText(ctx, chatID, "No audio formats available for this URL.")
			return
		}
		deduped := dl.DeduplicateResolutions(filtered)
		ss.FormatIDs = make([]string, len(deduped))
		labels := make([]string, len(deduped))
		for i, f := range deduped {
			ss.FormatIDs[i] = f.ID
			labels[i] = formatLabelForDisplay(f)
		}
		ss.FormatIDs = append([]string{"best"}, ss.FormatIDs...)
		labels = append([]string{"Best"}, labels...)
		h.setUserStepState(chatID, ss)
		kb := formatKeyboard("quality_", labels)
		h.sendWithKB(ctx, chatID, "Select audio quality:", kb)
		return
	}

	kb := mediaTypeKeyboard()
	h.sendWithKB(ctx, chatID, fmt.Sprintf("Title: %s\n\nSelect media type:", info.Title), kb)
}

// handleTypeSelection processes the media type choice.
func (h *Handler) handleTypeSelection(ctx context.Context, chatID int64, msgID int, userID int64, choice string) {
	ss := h.getUserStepState(chatID)
	if ss == nil {
		return
	}
	switch choice {
	case "audio":
		ss.MediaType = domain.MediaTypeAudio
	case "video":
		ss.MediaType = domain.MediaTypeVideo
	case "both":
		ss.MediaType = domain.MediaTypeVideoAudio
	default:
		return
	}
	ss.Step = 1

	// Filter and deduplicate formats by type
	filtered := dl.FilterFormatsByType(ss.MediaInfo.Formats, ss.MediaType)
	if len(filtered) == 0 {
		h.editText(ctx, chatID, msgID, fmt.Sprintf("No %s formats available for this URL.", choice), &models.InlineKeyboardMarkup{
			InlineKeyboard: [][]models.InlineKeyboardButton{
				{{Text: "Back", CallbackData: "back_step"}, {Text: "Cancel", CallbackData: "cancel_dl"}},
			},
		})
		return
	}
	deduped := dl.DeduplicateResolutions(filtered)
	ss.FormatIDs = make([]string, len(deduped))
	labels := make([]string, len(deduped))
	for i, f := range deduped {
		ss.FormatIDs[i] = f.ID
		labels[i] = formatLabelForDisplay(f)
	}
	// Add "Best" option first
	ss.FormatIDs = append([]string{"best"}, ss.FormatIDs...)
	labels = append([]string{"Best"}, labels...)

	h.setUserStepState(chatID, ss)
	kb := formatKeyboard("quality_", labels)
	h.editText(ctx, chatID, msgID, fmt.Sprintf("Title: %s\n\nSelect quality:", ss.MediaInfo.Title), &kb)
}

func formatLabelForDisplay(f domain.Format) string {
	if f.Resolution != "" {
		label := f.Resolution
		if f.Filesize > 0 {
			label += fmt.Sprintf(" (%s)", formatBytes(f.Filesize))
		}
		return label
	}
	if f.Bitrate != "" {
		label := f.Bitrate
		if f.Filesize > 0 {
			label += fmt.Sprintf(" (%s)", formatBytes(f.Filesize))
		}
		return label
	}
	return f.ID
}

// handleQualitySelection processes the quality choice.
func (h *Handler) handleQualitySelection(ctx context.Context, chatID int64, msgID int, userID int64, quality string) {
	ss := h.getUserStepState(chatID)
	if ss == nil {
		return
	}
	ss.Quality = quality
	ss.Step = 2

	// Check if we need language selection
	if len(ss.MediaInfo.Languages) > 1 {
		kb := languageKeyboard(ss.MediaInfo.Languages)
		h.editText(ctx, chatID, msgID, "Select language:", &kb)
		return
	}

	// Skip to container selection
	ss.Step = 3
	containers := dl.ContainerOptions(ss.MediaType)
	kb := containerKeyboard(containers)
	h.editText(ctx, chatID, msgID, "Select container format:", &kb)
}

// handleLanguageSelection processes the language choice.
func (h *Handler) handleLanguageSelection(ctx context.Context, chatID int64, msgID int, userID int64, language string) {
	ss := h.getUserStepState(chatID)
	if ss == nil {
		return
	}
	ss.Language = language
	ss.Step = 3
	h.setUserStepState(chatID, ss)

	containers := dl.ContainerOptions(ss.MediaType)
	kb := containerKeyboard(containers)
	h.editText(ctx, chatID, msgID, "Select container format:", &kb)
}

// handleContainerSelection processes the container format choice and starts the download.
func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, msgID int, userID int64, container string) {
	ss := h.getUserStepState(chatID)
	if ss == nil {
		return
	}
	ss.Container = container

	// Build the format object
	format := &domain.Format{ID: ss.Quality}
	for _, f := range ss.MediaInfo.Formats {
		if f.ID == ss.Quality {
			format = &f
			break
		}
	}

	// Create the download job
	job := &domain.DownloadJob{
		ID:             h.generateJobID(),
		UserID:         userID,
		ChatID:         chatID,
		MessageID:      msgID,
		URL:            ss.URL,
		MediaType:      ss.MediaType,
		SelectedFormat: format,
		Container:      container,
		Language:       ss.Language,
		Status:         domain.StatusDownloading,
		Title:          ss.MediaInfo.Title,
		FileSize:       format.Filesize,
	}

	h.setActiveJob(userID, job)
	h.clearStepState(chatID)

	// Start download in background
	go h.runDownload(context.Background(), job)
}

// runDownload manages the download lifecycle, progress updates, and completion.
func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
	// Show initial progress message
	kb := progressKeyboard(0, "", "")
	msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{
		ChatID:      job.ChatID,
		Text:        fmt.Sprintf("Downloading: %s", job.Title),
		ReplyMarkup: &kb,
	})
	if err != nil {
		slog.Error("send progress msg failed", "error", err)
		return
	}
	job.MessageID = msg.ID

	updates, err := h.YtDlp.StartDownload(job, h.downloadDir)
	if err != nil {
		h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", err.Error()), nil)
		job.Status = domain.StatusFailed
		h.mu.Lock()
		delete(h.activeJobs, job.UserID)
		h.mu.Unlock()
		return
	}

	lastProgress := 0.0
	for upd := range updates {
		h.mu.Lock()
		currentJob := h.activeJobs[job.UserID]
		h.mu.Unlock()
		if currentJob == nil {
			// Job was cancelled
			return
		}

		if upd.Status == domain.StatusCancelled {
			h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil)
			h.mu.Lock()
			delete(h.activeJobs, job.UserID)
			h.mu.Unlock()
			return
		}

		if upd.Status == domain.StatusFailed {
			h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil)
			h.mu.Lock()
			delete(h.activeJobs, job.UserID)
			h.mu.Unlock()
			return
		}

		// Update progress every ~5%
		if upd.Progress-lastProgress >= domain.ProgressThreshold || upd.Status == domain.StatusCompleted {
			lastProgress = upd.Progress

			// Apply quota at >50%
			if upd.Progress > domain.QuotaProgressThreshold && !job.QuotaApplied && job.FileSize > 0 {
				job.QuotaApplied = true
				h.applyQuota(job.UserID, job.FileSize)
			}

			kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA)
			text := fmt.Sprintf("Downloading: %s", job.Title)
			h.editText(ctx, job.ChatID, job.MessageID, text, &kb)
		}
	}

	// Download completed
	h.mu.Lock()
	delete(h.activeJobs, job.UserID)
	h.mu.Unlock()

	if job.Status == domain.StatusCompleted {
		// Find the downloaded file using job ID prefix
		pattern := filepath.Join(h.downloadDir, job.ID+"_*")
		files, err := filepath.Glob(pattern)
		if err == nil && len(files) > 0 {
			f := files[0] // take the first match
			fileInfo, err := os.Stat(f)
			if err == nil {
				if fileInfo.Size() > h.maxFileSize {
					h.editText(ctx, job.ChatID, job.MessageID,
						fmt.Sprintf("File too large (%s). Max: %s", formatBytes(fileInfo.Size()), formatBytes(h.maxFileSize)), nil)
					os.Remove(f)
					return
				}

				if !job.QuotaApplied && job.FileSize > 0 {
					job.QuotaApplied = true
					h.applyQuota(job.UserID, job.FileSize)
				}

				formatID := job.SelectedFormat.ID
				h.sendDocument(ctx, job.ChatID, f, job.Title)
				os.Remove(f)

				_ = h.Repo.AddHistory(job.UserID, &domain.DownloadRecord{
					URL:       job.URL,
					Title:     job.Title,
					MediaType: job.MediaType.String(),
					FormatID:  formatID,
					Container: job.Container,
					FileSize:  job.FileSize,
					Status:    "completed",
				})

				h.editText(ctx, job.ChatID, job.MessageID, "Download completed!", nil)
				return
			}
		}
		h.editText(ctx, job.ChatID, job.MessageID, "Download completed but file not found.", nil)
	}
}

func (h *Handler) sendDocument(ctx context.Context, chatID int64, filePath, title string) {
	f, err := os.Open(filePath)
	if err != nil {
		slog.Error("open file for send", "path", filePath, "error", err)
		return
	}
	defer f.Close()

	doc := &models.InputFileUpload{
		Filename: title + filepath.Ext(filePath),
		Data:     f,
	}

	if _, err := h.Bot.SendDocument(ctx, &tgbot.SendDocumentParams{
		ChatID:   chatID,
		Document: doc,
	}); err != nil {
		slog.Error("send document failed", "error", err)
	}
}

// handleCancelDownload cancels the current download for a user.
func (h *Handler) handleCancelDownload(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
	defer h.answerCb(ctx, cbID)

	job := h.getActiveJob(userID)
	if job != nil {
		dl.CancelDownload(job)
		h.mu.Lock()
		delete(h.activeJobs, userID)
		h.mu.Unlock()
		h.editText(ctx, chatID, msgID, "Download cancelled.", nil)
		return
	}
	h.clearStepState(chatID)
	h.editText(ctx, chatID, msgID, "Cancelled.", nil)
}

// handleBackStep goes back one step in the download flow.
func (h *Handler) handleBackStep(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
	defer h.answerCb(ctx, cbID)
	ss := h.getUserStepState(chatID)
	if ss == nil {
		return
	}

	ss.Step--
	switch ss.Step {
	case -1:
		// Go back to URL input
		h.clearStepState(chatID)
		h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID)
	case 0:
		if ss.MediaInfo != nil && ss.MediaInfo.IsAudioOnlyContent() {
			// Audio-only: go back to URL input directly
			h.clearStepState(chatID)
			h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID)
			return
		}
		h.editText(ctx, chatID, msgID, "Select media type:", &models.InlineKeyboardMarkup{
			InlineKeyboard: mediaTypeKeyboard().InlineKeyboard,
		})
	case 1:
		// Re-show quality selector
		filtered := dl.FilterFormatsByType(ss.MediaInfo.Formats, ss.MediaType)
		deduped := dl.DeduplicateResolutions(filtered)
		labels := make([]string, len(deduped))
		for i, f := range deduped {
			labels[i] = formatLabelForDisplay(f)
		}
		labels = append([]string{"Best"}, labels...)
		kb := formatKeyboard("quality_", labels)
		h.editText(ctx, chatID, msgID, "Select quality:", &kb)
	case 2:
		if len(ss.MediaInfo.Languages) > 1 {
			kb := languageKeyboard(ss.MediaInfo.Languages)
			h.editText(ctx, chatID, msgID, "Select language:", &kb)
		} else {
			ss.Step = 1
			h.handleBackStep(ctx, chatID, msgID, cbID, userID)
		}
	}
}

// handleDownloadCommand handles the /download command, prompting for a URL.
func (h *Handler) handleDownloadCommand(ctx context.Context, msg *models.Message, userID int64) {
	h.promptForInput(ctx, msg.Chat.ID, 0, userID, PendingURL,
		"Send me the URL to download:", nil)
}

// handleDownloadPrompt shows the download prompt (asks for URL).
func (h *Handler) handleDownloadPrompt(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) {
	h.promptForInput(ctx, chatID, msgID, userID, PendingURL,
		"Send me the URL to download:", nil)
}

// sendMediaPreview sends a thumbnail and metadata info before the selection flow.
func (h *Handler) sendMediaPreview(ctx context.Context, chatID int64, info *domain.MediaInfo) {
	caption := info.Title
	if info.Uploader != "" || info.Duration > 0 {
		caption += "\n"
		if info.Uploader != "" {
			caption += fmt.Sprintf("Channel: %s", info.Uploader)
		}
		if info.Duration > 0 {
			if info.Uploader != "" {
				caption += " | "
			}
			caption += fmt.Sprintf("Duration: %s", formatDuration(info.Duration))
		}
	}

	if info.Thumbnail != "" {
		photo := &models.InputFileString{Data: info.Thumbnail}
		_, err := h.Bot.SendPhoto(ctx, &tgbot.SendPhotoParams{
			ChatID:  chatID,
			Photo:   photo,
			Caption: caption,
		})
		if err == nil {
			return
		}
		slog.Warn("send thumbnail failed, falling back to text", "error", err)
	}

	h.sendText(ctx, chatID, caption)
}

func formatDuration(seconds float64) string {
	if seconds <= 0 {
		return "0:00"
	}
	d := time.Duration(seconds) * time.Second
	h := int(d.Hours())
	m := int(math.Mod(d.Minutes(), 60))
	s := int(math.Mod(d.Seconds(), 60))
	if h > 0 {
		return fmt.Sprintf("%d:%02d:%02d", h, m, s)
	}
	return fmt.Sprintf("%d:%02d", m, s)
}