uptodownbot/internal/dl/ytdlp.go

package dl

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"uptodownBot/internal/domain"
)

// Client wraps calls to the yt-dlp executable.
type Client struct {
	binaryPath  string
	cookiesFile string
	proxyURL    string
	cache       map[string]*cachedMediaInfo
	cacheMu     sync.Mutex
}

type cachedMediaInfo struct {
	info *domain.MediaInfo
	ts   time.Time
}

const cacheTTL = 30 * time.Minute
const cacheMaxEntries = 200

func (c *Client) getCached(url string) *domain.MediaInfo {
	c.cacheMu.Lock()
	defer c.cacheMu.Unlock()
	entry, ok := c.cache[url]
	if !ok || time.Since(entry.ts) > cacheTTL {
		delete(c.cache, url)
		return nil
	}
	return entry.info
}

func (c *Client) setCache(url string, info *domain.MediaInfo) {
	c.cacheMu.Lock()
	defer c.cacheMu.Unlock()
	if len(c.cache) >= cacheMaxEntries {
		c.cache = make(map[string]*cachedMediaInfo)
	}
	c.cache[url] = &cachedMediaInfo{info: info, ts: time.Now()}
}

// NewClient creates a new yt-dlp client. It checks for the binary.
func NewClient(cookiesFile string) (*Client, error) {
	path, err := exec.LookPath("yt-dlp")
	if err != nil {
		return nil, fmt.Errorf("yt-dlp not found in PATH: %w", err)
	}
	slog.Info("yt-dlp found", "path", path)

	proxyURL := os.Getenv("HTTP_PROXY")
	if proxyURL == "" {
		proxyURL = os.Getenv("HTTPS_PROXY")
	}
	if proxyURL != "" {
		slog.Info("using proxy for yt-dlp", "proxy", proxyURL)
	}

	return &Client{
		binaryPath:  path,
		cookiesFile: cookiesFile,
		proxyURL:    proxyURL,
		cache:       make(map[string]*cachedMediaInfo),
	}, nil
}

type ytdlpFormat struct {
	FormatID       string  `json:"format_id"`
	Ext            string  `json:"ext"`
	Width          int     `json:"width"`
	Height         int     `json:"height"`
	Filesize       int64   `json:"filesize"`
	FilesizeApprox int64   `json:"filesize_approx"`
	Codec          string  `json:"vcodec"`
	ACodec         string  `json:"acodec"`
	TBR            float64 `json:"tbr"`
	ABR            float64 `json:"abr"`
	Language       string  `json:"language"`
	Resolution     string  `json:"resolution"`
}

type ytdlpInfo struct {
	Title            string            `json:"title"`
	Uploader         string            `json:"uploader"`
	Thumbnail        string            `json:"thumbnail"`
	Duration         float64           `json:"duration"`
	IsLive           bool              `json:"is_live"`
	Formats          []ytdlpFormat     `json:"formats"`
	Playlist         string            `json:"playlist"`
	PlaylistID       string            `json:"playlist_id"`
	PlaylistCount    int               `json:"playlist_count"`
	Entries          []json.RawMessage `json:"entries"`
	RequestedFormats []ytdlpFormat     `json:"requested_formats"`
}

type ytdlpPlaylistEntry struct {
	ID       string  `json:"id"`
	Title    string  `json:"title"`
	URL      string  `json:"url"`
	Duration float64 `json:"duration"`
}

// GetMediaInfo fetches format information for a given URL.
func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) {
	if cached := c.getCached(url); cached != nil {
		slog.Debug("cache hit", "url", url)
		return cached, nil
	}

	args := []string{
		"-J", "--no-download",
		"--no-warnings",
		url,
	}
	if c.cookiesFile != "" {
		args = append(args, "--cookies", c.cookiesFile)
	}
	if c.proxyURL != "" {
		args = append(args, "--proxy", c.proxyURL)
	}

	slog.Info("fetching media info", "url", url)
	out, err := c.run(args)
	if err != nil {
		return nil, fmt.Errorf("yt-dlp info failed: %w", err)
	}

	var info ytdlpInfo
	if err := json.Unmarshal(out, &info); err != nil {
		return nil, fmt.Errorf("parse yt-dlp output: %w", err)
	}

	if info.IsLive {
		return nil, fmt.Errorf("cannot download live streams")
	}

	mi := &domain.MediaInfo{
		URL:       url,
		Title:     info.Title,
		Uploader:  info.Uploader,
		Thumbnail: info.Thumbnail,
		Duration:  info.Duration,
	}

	langSet := make(map[string]bool)
	var maxSize int64
	for _, f := range info.Formats {
		ff := c.convertFormat(f)
		mi.Formats = append(mi.Formats, ff)
		if (ff.HasAudio || ff.HasVideo) && ff.Language != "" && !langSet[ff.Language] {
			langSet[ff.Language] = true
			mi.Languages = append(mi.Languages, ff.Language)
		}
		if ff.Filesize > maxSize {
			maxSize = ff.Filesize
		}
	}
	mi.EstimatedSize = maxSize

	c.setCache(url, mi)
	return mi, nil
}

func (c *Client) parsePlaylistEntries(entries []json.RawMessage) []domain.PlaylistEntry {
	var result []domain.PlaylistEntry
	for _, raw := range entries {
		var entry ytdlpPlaylistEntry
		if err := json.Unmarshal(raw, &entry); err != nil {
			slog.Warn("failed to parse playlist entry", "error", err)
			continue
		}
		result = append(result, domain.PlaylistEntry{
			ID:       entry.ID,
			Title:    entry.Title,
			URL:      entry.URL,
			Duration: entry.Duration,
		})
	}
	return result
}

func (c *Client) convertFormat(f ytdlpFormat) domain.Format {
	ff := domain.Format{
		ID:         f.FormatID,
		Extension:  f.Ext,
		Width:      f.Width,
		Height:     f.Height,
		Resolution: f.Resolution,
		Codec:      f.Codec,
		Language:   f.Language,
	}
	if f.Filesize > 0 {
		ff.Filesize = f.Filesize
	} else {
		ff.Filesize = f.FilesizeApprox
	}
	if f.TBR > 0 {
		ff.Bitrate = fmt.Sprintf("%.0fk", f.TBR)
	} else if f.ABR > 0 {
		ff.Bitrate = fmt.Sprintf("%.0fk", f.ABR)
	}
	ff.HasVideo = f.Codec != "none" && f.Codec != ""
	ff.HasAudio = f.ACodec != "none" && f.ACodec != ""
	ff.IsAudioOnly = !ff.HasVideo && ff.HasAudio
	ff.IsVideoOnly = ff.HasVideo && !ff.HasAudio
	return ff
}

// StartDownload begins a download and sends progress updates to the channel.
// Returns the job and an updates channel. Caller must consume from updates until closed.
func (c *Client) StartDownload(job *domain.DownloadJob, downloadDir string) (<-chan *domain.DownloadJob, error) {
	if err := os.MkdirAll(downloadDir, 0755); err != nil {
		return nil, fmt.Errorf("create download dir: %w", err)
	}

	formatID := job.SelectedFormat.ID
	var args []string

	// Use format string for the download — yt-dlp handles container choice
	if job.FormatString != "" {
		args = append(args, "-f", job.FormatString)
	} else {
		switch job.MediaType {
		case domain.MediaTypeAudio:
			args = append(args, "-f", formatID, "-x")
		case domain.MediaTypeVideo:
			args = append(args, "-f", formatID)
		case domain.MediaTypeVideoAudio:
			args = append(args, "-f", formatID+"+bestaudio/best")
		}
	}

	// Output template with job ID prefix for reliable file finding
	outputPath := filepath.Join(downloadDir, job.ID+"_%(id)s.%(ext)s")
	args = append(args, "--output", outputPath)

	// Remux to Telegram-friendly containers: mp4 for video, m4a for audio-only.
	if job.MediaType == domain.MediaTypeAudio {
		args = append(args, "--merge-output-format", "m4a")
	} else {
		args = append(args, "--merge-output-format", "mp4")
	}

	// --progress forces progress output in non-TTY (pipe) environments
	// --newline uses \n instead of \r so bufio.Scanner can read line-by-line
	args = append(args, "--progress", "--newline")

	// Allow resuming partial downloads
	args = append(args, "--continue")

	// Enforce max file size limit
	if job.MaxFileSize > 0 {
		args = append(args, "--max-filesize", strconv.FormatInt(job.MaxFileSize, 10))
	}

	// Language selection if specified
	if job.Language != "" {
		args = append(args, "--sub-langs", job.Language)
	}

	// Cookies
	if c.cookiesFile != "" {
		args = append(args, "--cookies", c.cookiesFile)
	}
	if c.proxyURL != "" {
		args = append(args, "--proxy", c.proxyURL)
	}

	args = append(args, job.URL)

	slog.Info("starting download", "url", job.URL, "format", formatID)

	cmd := exec.Command(c.binaryPath, args...)
	cmd.Env = append(os.Environ(), "PYTHONUNBUFFERED=1")

	// yt-dlp outputs [download] progress lines to stdout when saving to a file
	// and to stderr when using -o -. Since we save to a file, read stdout for progress.
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("create stdout pipe: %w", err)
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("create stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start yt-dlp: %w", err)
	}

	job.Cmd = cmd
	job.CancelCh = make(chan struct{})

	updates := make(chan *domain.DownloadJob, 100)

	go func() {
		defer func() {
			if job.Status != domain.StatusCancelled {
				updates <- job
			}
			close(updates)
		}()

		progressDone := make(chan struct{})
		go func() {
			readProgress(stdout, updates, job)
			close(progressDone)
		}()

		// Capture stderr (warnings, errors) for diagnostics
		go func() {
			slurp, _ := io.ReadAll(stderr)
			job.StderrLog = string(slurp)
		}()

		// Wait for either completion or cancellation
		done := make(chan error, 1)
		go func() {
			done <- cmd.Wait()
		}()

		select {
		case <-job.CancelCh:
			slog.Info("cancelling download", "url", job.URL)
			cmd.Process.Signal(syscall.SIGTERM)
			select {
			case <-done:
			case <-time.After(3 * time.Second):
				cmd.Process.Kill()
			}
			<-progressDone
			job.Status = domain.StatusCancelled
		case err := <-done:
			<-progressDone
			if err != nil {
				job.Status = domain.StatusFailed
				slog.Error("download failed", "url", job.URL, "error", err, "stderr", job.StderrLog)
			} else {
				job.Status = domain.StatusCompleted
			}
		}
	}()

	return updates, nil
}

// CancelDownload sends a cancellation signal to a running download.
func CancelDownload(job *domain.DownloadJob) {
	if job == nil {
		return
	}
	select {
	case job.CancelCh <- struct{}{}:
	default:
		// Already cancelled or not started
	}
	if job.Cmd != nil && job.Cmd.Process != nil {
		job.Cmd.Process.Signal(syscall.SIGTERM)
	}
}

// ytdlpSearchResult holds a single flat search result from yt-dlp.
type ytdlpSearchResult struct {
	Title string `json:"title"`
	URL   string `json:"url"`
}

// SearchYoutube searches YouTube via yt-dlp ytsearch and returns the first result URL and title.
func (c *Client) SearchYoutube(query string) (url, title string, err error) {
	args := []string{
		"--flat-playlist", "-J", "--no-download",
		"--no-warnings",
		"ytsearch1:" + query,
	}
	if c.cookiesFile != "" {
		args = append(args, "--cookies", c.cookiesFile)
	}
	if c.proxyURL != "" {
		args = append(args, "--proxy", c.proxyURL)
	}

	out, err := c.run(args)
	if err != nil {
		return "", "", fmt.Errorf("youtube search failed: %w", err)
	}

	var info struct {
		Entries []ytdlpSearchResult `json:"entries"`
	}
	if err := json.Unmarshal(out, &info); err != nil {
		return "", "", fmt.Errorf("parse search result: %w", err)
	}
	if len(info.Entries) == 0 {
		return "", "", fmt.Errorf("no results found")
	}
	return info.Entries[0].URL, info.Entries[0].Title, nil
}

// Search does a yt-dlp search and returns entries.
func (c *Client) Search(query, source string, limit int) ([]domain.PlaylistEntry, error) {
	searchPrefix := "ytsearch"
	if source == "ytmusic" {
		searchPrefix = "ytmusicsearch"
	}

	args := []string{
		"--flat-playlist", "-J", "--no-download",
		"--no-warnings",
		fmt.Sprintf("%s%d:%s", searchPrefix, limit, query),
	}
	if c.cookiesFile != "" {
		args = append(args, "--cookies", c.cookiesFile)
	}
	if c.proxyURL != "" {
		args = append(args, "--proxy", c.proxyURL)
	}

	slog.Info("searching", "query", query, "source", source)
	out, err := c.run(args)
	if err != nil {
		return nil, fmt.Errorf("search failed: %w", err)
	}

	var info ytdlpInfo
	if err := json.Unmarshal(out, &info); err != nil {
		return nil, fmt.Errorf("parse search result: %w", err)
	}

	return c.parsePlaylistEntries(info.Entries), nil
}

func (c *Client) run(args []string) ([]byte, error) {
	cmd := exec.Command(c.binaryPath, args...)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	out, err := cmd.Output()
	if err != nil {
		stderrStr := strings.TrimSpace(stderr.String())
		slog.Warn("yt-dlp run failed", "stderr", stderrStr, "error", err)
		// Extract user-friendly error message
		if stderrStr != "" {
			return nil, errors.New(stderrStr)
		}
		return nil, err
	}
	return out, nil
}

// FilterFormatsByType filters formats by the requested media type.
func FilterFormatsByType(formats []domain.Format, mt domain.MediaType) []domain.Format {
	var filtered []domain.Format
	for _, f := range formats {
		switch mt {
		case domain.MediaTypeAudio:
			if f.IsAudioOnly {
				filtered = append(filtered, f)
			}
		case domain.MediaTypeVideo:
			if f.IsVideoOnly {
				filtered = append(filtered, f)
			}
		case domain.MediaTypeVideoAudio:
			if f.HasVideo && f.HasAudio {
				filtered = append(filtered, f)
			}
		}
	}
	return filtered
}

// DeduplicateResolutions reduces formats to unique resolution labels.
func DeduplicateResolutions(formats []domain.Format) []domain.Format {
	seen := make(map[string]bool)
	var result []domain.Format
	for _, f := range formats {
		label := formatLabel(f)
		if !seen[label] {
			seen[label] = true
			result = append(result, f)
		}
	}
	return result
}

// formatLabel returns a human-readable label for a format.
func formatLabel(f domain.Format) string {
	if f.IsAudioOnly {
		if f.Bitrate != "" {
			return f.Bitrate
		}
		if f.Extension != "" {
			return f.ID + " (" + f.Extension + ")"
		}
		return f.ID
	}
	if f.Resolution != "" {
		return f.Resolution
	}
	if f.Height > 0 {
		return fmt.Sprintf("%dp", f.Height)
	}
	return f.ID
}