streamrip-go/internal/download/downloader.go

package download

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync/atomic"

	"github.com/vbauerster/mpb/v8"
	"github.com/vbauerster/mpb/v8/decor"
	"golang.org/x/term"

	"streamrip-go/internal/netutil"
)

type Downloader struct {
	http         *http.Client
	showProgress bool
	progress     *mpb.Progress
	barStarted   atomic.Int32
}

func New() *Downloader {
	return NewWithOptions(true, true)
}

func NewWithVerifySSL(verifySSL bool) *Downloader {
	return NewWithOptions(verifySSL, true)
}

func NewWithOptions(verifySSL bool, showProgress bool) *Downloader {
	forceProgress := strings.EqualFold(os.Getenv("STREAMRIP_GO_FORCE_PROGRESS"), "1") || strings.EqualFold(os.Getenv("STREAMRIP_GO_FORCE_PROGRESS"), "true")
	interactive := showProgress && (forceProgress || (term.IsTerminal(int(os.Stderr.Fd())) && strings.ToLower(os.Getenv("TERM")) != "dumb"))
	d := &Downloader{http: netutil.NewHTTPClient(0, verifySSL), showProgress: interactive}
	if interactive {
		d.progress = mpb.New(mpb.WithWidth(40), mpb.WithOutput(os.Stderr))
	}
	return d
}

func (d *Downloader) File(ctx context.Context, sourceURL, outputPath string) error {
	return d.file(ctx, sourceURL, outputPath, true, false)
}

func (d *Downloader) FileNoProgress(ctx context.Context, sourceURL, outputPath string) error {
	return d.file(ctx, sourceURL, outputPath, false, false)
}

func (d *Downloader) FileVideo(ctx context.Context, sourceURL, outputPath string) error {
	return d.file(ctx, sourceURL, outputPath, true, true)
}

func (d *Downloader) file(ctx context.Context, sourceURL, outputPath string, allowProgress bool, includeVideo bool) error {
	if err := os.MkdirAll(filepath.Dir(outputPath), 0o755); err != nil {
		return err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL, nil)
	if err != nil {
		return err
	}

	resp, err := d.http.Do(req)
	if err != nil {
		return err
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("download failed: status=%d", resp.StatusCode)
	}

	reader := bufio.NewReader(resp.Body)
	peek, _ := reader.Peek(1024)
	if isManifestResponse(resp.Header.Get("Content-Type"), peek) {
		_ = resp.Body.Close()
		return d.streamManifestWithFFmpeg(ctx, sourceURL, outputPath, includeVideo)
	}

	out, err := os.Create(outputPath)
	if err != nil {
		return err
	}
	defer func() { _ = out.Close() }()

	if d.ProgressEnabled() && allowProgress && resp.ContentLength > 0 {
		d.barStarted.Store(1)
		desc := shortenName(filepath.Base(outputPath), 54)
		bar := d.progress.AddBar(
			resp.ContentLength,
			mpb.PrependDecorators(
				decor.Name(desc+" ", decor.WC{W: 56, C: decor.DSyncWidth | decor.DindentRight}),
				decor.Percentage(decor.WCSyncWidthR),
			),
			mpb.AppendDecorators(
				decor.CountersKibiByte("% .1f / % .1f", decor.WCSyncWidthR),
				decor.Name(" | ", decor.WCSyncWidth),
				decor.AverageSpeed(decor.SizeB1024(0), "% .1f", decor.WCSyncWidthR),
				decor.Name(" | ETA ", decor.WCSyncWidth),
				decor.AverageETA(decor.ET_STYLE_GO, decor.WCSyncWidthR),
			),
			mpb.BarRemoveOnComplete(),
		)
		buf := make([]byte, 256*1024)
		for {
			n, readErr := reader.Read(buf)
			if n > 0 {
				if _, writeErr := out.Write(buf[:n]); writeErr != nil {
					return writeErr
				}
				bar.IncrBy(n)
			}
			if readErr != nil {
				if readErr == io.EOF {
					break
				}
				return readErr
			}
		}
	} else {
		if _, err = io.Copy(out, reader); err != nil {
			return err
		}
	}

	return nil
}

func (d *Downloader) Close() {
	if d.progress != nil {
		d.progress.Wait()
	}
}

func (d *Downloader) ProgressEnabled() bool {
	return d.showProgress && d.progress != nil
}

func (d *Downloader) Logf(format string, args ...any) {
	msg := fmt.Sprintf(format, args...)
	if d.ProgressEnabled() && d.barStarted.Load() == 1 {
		_, _ = d.progress.Write([]byte(msg))
		return
	}
	fmt.Print(msg)
}

func shortenName(name string, max int) string {
	if max <= 0 {
		return name
	}
	r := []rune(name)
	if len(r) <= max {
		return name
	}
	if max <= 3 {
		return string(r[:max])
	}
	return string(r[:max-3]) + "..."
}

func (d *Downloader) streamManifestWithFFmpeg(ctx context.Context, sourceURL, outputPath string, includeVideo bool) error {
	if _, err := exec.LookPath("ffmpeg"); err != nil {
		return fmt.Errorf("ffmpeg not found for manifest stream: %w", err)
	}

	args := []string{
		"-y",
		"-protocol_whitelist", "file,http,https,tcp,tls,crypto,data",
		"-i", sourceURL,
	}
	if includeVideo {
		args = append(args, "-map", "0")
	} else {
		args = append(args, "-map", "0:a:0")
	}
	args = append(args, "-c", "copy", outputPath)

	cmd := exec.CommandContext(ctx, "ffmpeg", args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("ffmpeg stream copy failed: %w: %s", err, string(output))
	}
	return nil
}

func isManifestResponse(contentType string, peek []byte) bool {
	ct := strings.ToLower(contentType)
	if strings.Contains(ct, "dash+xml") || strings.Contains(ct, "mpegurl") || strings.Contains(ct, "vnd.apple.mpegurl") {
		return true
	}
	s := strings.TrimSpace(strings.ToLower(string(peek)))
	if strings.HasPrefix(s, "<?xml") && strings.Contains(s, "<mpd") {
		return true
	}
	if strings.HasPrefix(s, "#extm3u") {
		return true
	}
	return false
}