uptodownbot/internal/dl/ytdlp.go

package dl

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

	"uptodownBot/internal/domain"
)

// Client wraps calls to the yt-dlp executable.
type Client struct {
	binaryPath  string
	cookiesFile string
}

// 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)
	return &Client{binaryPath: path, cookiesFile: cookiesFile}, 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"`
	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) {
	args := []string{
		"-J", "--no-download",
		"--no-warnings",
		url,
	}
	if c.cookiesFile != "" {
		args = append(args, "--cookies", c.cookiesFile)
	}

	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,
		Duration:   info.Duration,
		IsPlaylist: info.Playlist != "",
	}

	if mi.IsPlaylist {
		mi.PlaylistCount = info.PlaylistCount
		mi.PlaylistEntries = c.parsePlaylistEntries(info.Entries)
		return mi, nil
	}

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

	return mi, nil
}

// GetPlaylistInfo fetches a flat playlist listing (fast, no format detail).
func (c *Client) GetPlaylistInfo(url string) (*domain.MediaInfo, error) {
	args := []string{
		"-J", "--flat-playlist", "--no-download",
		"--no-warnings",
		url,
	}
	if c.cookiesFile != "" {
		args = append(args, "--cookies", c.cookiesFile)
	}

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

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

	mi := &domain.MediaInfo{
		URL:             url,
		Title:           info.Title,
		IsPlaylist:      true,
		PlaylistCount:   info.PlaylistCount,
		PlaylistEntries: c.parsePlaylistEntries(info.Entries),
	}

	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

	// Build format string for merge
	switch job.MediaType {
	case domain.MediaTypeAudio:
		args = append(args, "-f", formatID, "-x")
		if job.Container != "" {
			args = append(args, "--audio-format", job.Container)
		}
	case domain.MediaTypeVideo:
		args = append(args, "-f", formatID)
	case domain.MediaTypeVideoAudio:
		// Download best video + best audio that match
		args = append(args, "-f", formatID+"+bestaudio/best")
		if job.Container != "" {
			args = append(args, "--merge-output-format", job.Container)
		}
	}

	// Output template
	outputPath := filepath.Join(downloadDir, "%(id)s.%(ext)s")
	args = append(args, "--output", outputPath)

	// Progress reporting
	args = append(args, "--newline", "--no-progress")

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

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

	args = append(args, job.URL)

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

	cmd := exec.Command(c.binaryPath, args...)

	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 {
				job.Status = domain.StatusCompleted
				updates <- job
			}
			close(updates)
		}()

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

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

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
		}
		return "audio"
	}
	if f.Resolution != "" {
		return f.Resolution
	}
	if f.Height > 0 {
		return fmt.Sprintf("%dp", f.Height)
	}
	return f.ID
}

// ContainerOptions returns the container format options for a given media type.
func ContainerOptions(mt domain.MediaType) []string {
	switch mt {
	case domain.MediaTypeAudio:
		return []string{"mp3", "m4a", "opus"}
	case domain.MediaTypeVideo:
		return []string{"mp4", "mkv", "webm"}
	case domain.MediaTypeVideoAudio:
		return []string{"mp4", "mkv"}
	default:
		return []string{"mp4"}
	}
}