yt-dlp-telegram-bot/dl.go

package main

import (
	"context"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync/atomic"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/wader/goutubedl"
)

const downloadAndConvertTimeout = 30 * time.Minute
const telegramUploadThreshold = 512 * 1024 * 1024 // 512MB
const downloadProgressUpdateInterval = time.Second

type ConvertStartCallbackFunc func(ctx context.Context, videoCodecs, audioCodecs, convertActionsNeeded string)
type UpdateProgressPercentCallbackFunc func(progressStr string, progressPercent int)

type Downloader struct {
	ConvertStartFunc          ConvertStartCallbackFunc
	UpdateProgressPercentFunc UpdateProgressPercentCallbackFunc
}

type goYouTubeDLLogger struct{}

func (l goYouTubeDLLogger) Print(v ...interface{}) {
	fmt.Print("  yt-dlp dbg:")
	fmt.Println(v...)
}

type DownloadResult struct {
	Title    string
	FilePath string
	FileSize int64
}

// progressWriter wraps an io.Writer and tracks bytes written
type progressWriter struct {
	writer       io.Writer
	written      int64
	total        int64 // estimated total size, 0 if unknown
	callback     UpdateProgressPercentCallbackFunc
	lastUpdate   time.Time
	updateInterval time.Duration
}

func newProgressWriter(w io.Writer, total int64, callback UpdateProgressPercentCallbackFunc) *progressWriter {
	return &progressWriter{
		writer:         w,
		total:          total,
		callback:       callback,
		lastUpdate:     time.Now(),
		updateInterval: downloadProgressUpdateInterval,
	}
}

func (pw *progressWriter) Write(p []byte) (n int, err error) {
	n, err = pw.writer.Write(p)
	if n > 0 {
		atomic.AddInt64(&pw.written, int64(n))

		now := time.Now()
		if now.Sub(pw.lastUpdate) >= pw.updateInterval {
			pw.lastUpdate = now
			written := atomic.LoadInt64(&pw.written)

			if pw.total > 0 && pw.callback != nil {
				percent := int(float64(written) * 100 / float64(pw.total))
				if percent > 100 {
					percent = 100
				}
				pw.callback("⬇️ Downloading", percent)
			} else if pw.callback != nil {
				// Unknown total size, just show bytes downloaded
				pw.callback(fmt.Sprintf("⬇️ Downloaded %s", humanize.Bytes(uint64(written))), -1)
			}
		}
	}
	return n, err
}

func (pw *progressWriter) Written() int64 {
	return atomic.LoadInt64(&pw.written)
}

func (d *Downloader) downloadURL(dlCtx context.Context, url string) (*DownloadResult, error) {
	// Use 4K quality for saving, but fall back to best available
	result, err := goutubedl.New(dlCtx, url, goutubedl.Options{
		Type:              goutubedl.TypeSingle,
		DebugLog:          goYouTubeDLLogger{},
		MergeOutputFormat: "mkv",      // This handles VP9 properly. yt-dlp uses mp4 by default, which doesn't.
		SortingFormat:     "res:2160", // Prefer videos up to 4K (2160p)
	})
	if err != nil {
		return nil, fmt.Errorf("preparing download %q: %w", url, err)
	}

	// Create filename with date and timestamp format: 2025-11-22-{timestamp}.mkv
	now := time.Now()
	dateStr := now.Format("2006-01-02")
	timestamp := now.Unix()
	fileName := fmt.Sprintf("%s-%d.mkv", dateStr, timestamp)
	filePath := filepath.Join(params.SaveDir, fileName)

	// Check if file already exists, if so add a suffix
	if _, err := os.Stat(filePath); err == nil {
		for i := 1; i < 1000; i++ {
			fileName = fmt.Sprintf("%s-%d-%d.mkv", dateStr, timestamp, i)
			filePath = filepath.Join(params.SaveDir, fileName)
			if _, err := os.Stat(filePath); os.IsNotExist(err) {
				break
			}
		}
	}

	dlResult, err := result.Download(dlCtx, "")
	if err != nil {
		return nil, fmt.Errorf("downloading %q: %w", url, err)
	}
	defer dlResult.Close()

	// Create file
	file, err := os.Create(filePath)
	if err != nil {
		return nil, fmt.Errorf("creating file %q: %w", filePath, err)
	}
	defer file.Close()

	// Get estimated file size from format info
	var estimatedSize int64
	if len(result.Info.Formats) > 0 {
		// Try to get filesize from the selected format
		for _, f := range result.Info.Formats {
			if f.Filesize > 0 {
				fs := int64(f.Filesize)
				if fs > estimatedSize {
					estimatedSize = fs
				}
			} else if f.FilesizeApprox > 0 && estimatedSize == 0 {
				estimatedSize = int64(f.FilesizeApprox)
			}
		}
	}
	if estimatedSize == 0 && result.Info.Filesize > 0 {
		estimatedSize = int64(result.Info.Filesize)
	}
	if estimatedSize == 0 && result.Info.FilesizeApprox > 0 {
		estimatedSize = int64(result.Info.FilesizeApprox)
	}

	// Create progress writer
	pw := newProgressWriter(file, estimatedSize, d.UpdateProgressPercentFunc)

	// Copy data to file with progress tracking
	_, err = io.Copy(pw, dlResult)
	if err != nil {
		os.Remove(filePath)
		return nil, fmt.Errorf("writing to file %q: %w", filePath, err)
	}

	written := pw.Written()
	fmt.Printf("  saved to %s (%s)\n", filePath, humanize.Bytes(uint64(written)))

	return &DownloadResult{
		Title:    result.Info.Title,
		FilePath: filePath,
		FileSize: written,
	}, nil
}

func (d *Downloader) DownloadAndConvertURL(ctx context.Context, url, format string) (*DownloadResult, error) {
	return d.downloadURL(ctx, url)
}