hako/internal/archival/archiver/yt_dlp.go

package archiver

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"git.nakama.town/fmartingr/hako/internal/model"
	"git.nakama.town/fmartingr/hako/internal/storage"
)

// YtDlpConfig defines the configuration for the yt-dlp archiver
type YtDlpConfig struct {
	Timeout        string `json:"timeout"`
	WriteThumbnail bool   `json:"write_thumbnail"`
	WriteSubs      bool   `json:"write_subs"` // Download subtitles as separate files
	EmbedSubs      bool   `json:"embed_subs"` // Embed subtitles into the video file
}

// ToMap converts the YtDlpConfig to a map[string]any for use with ApplyConfig
func (c YtDlpConfig) ToMap() map[string]any {
	data, err := json.Marshal(c)
	if err != nil {
		return nil
	}
	var out map[string]any
	if err := json.Unmarshal(data, &out); err != nil {
		return nil
	}
	return out
}

// YtDlpArchiver implements video download archival using the yt-dlp binary
type YtDlpArchiver struct {
	binaryPath string
	config     YtDlpConfig
}

// NewYtDlpArchiver creates a new YtDlpArchiver
func NewYtDlpArchiver() *YtDlpArchiver {
	return &YtDlpArchiver{
		config: YtDlpConfig{
			Timeout:        "10m",
			WriteThumbnail: true,
			WriteSubs:      true,
			EmbedSubs:      true,
		},
	}
}

// Key returns the unique identifier for this archiver
func (a *YtDlpArchiver) Key() string {
	return "yt_dlp"
}

// Name returns the human-readable name for this archiver
func (a *YtDlpArchiver) Name() string {
	return "yt-dlp"
}

// IsEnabled checks if the yt-dlp binary is available
func (a *YtDlpArchiver) IsEnabled(ctx context.Context) bool {
	path, err := exec.LookPath("yt-dlp")
	if err != nil {
		return false
	}
	a.binaryPath = path
	return true
}

// Init initializes the archiver (resolves binary path)
func (a *YtDlpArchiver) Init() error {
	path, err := exec.LookPath("yt-dlp")
	if err != nil {
		return fmt.Errorf("yt-dlp not found in PATH: %w", err)
	}
	a.binaryPath = path
	return nil
}

// GetDefaultConfig returns the default configuration
func (a *YtDlpArchiver) GetDefaultConfig() any {
	return YtDlpConfig{
		Timeout:        "10m",
		WriteThumbnail: true,
		WriteSubs:      true,
		EmbedSubs:      true,
	}
}

// ApplyConfig applies the provided configuration
func (a *YtDlpArchiver) ApplyConfig(config map[string]any) error {
	if timeoutStr, ok := config["timeout"].(string); ok && timeoutStr != "" {
		if _, err := time.ParseDuration(timeoutStr); err != nil {
			return fmt.Errorf("invalid timeout format: %w", err)
		}
		a.config.Timeout = timeoutStr
	}
	if v, ok := config["write_thumbnail"].(bool); ok {
		a.config.WriteThumbnail = v
	}
	if v, ok := config["write_subs"].(bool); ok {
		a.config.WriteSubs = v
	}
	if v, ok := config["embed_subs"].(bool); ok {
		a.config.EmbedSubs = v
	}
	return nil
}

// ytDlpMetadata holds the title field from yt-dlp --dump-json output
type ytDlpMetadata struct {
	Title string `json:"title"`
}

func (a *YtDlpArchiver) fetchTitle(ctx context.Context, url string) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, 2*time.Minute)
	defer cancel()
	cmd := exec.CommandContext(ctx, a.binaryPath, "--dump-json", "--no-download", "--no-warnings", url)
	out, err := cmd.Output()
	if err != nil {
		return "", err
	}
	var meta ytDlpMetadata
	if err := json.Unmarshal(out, &meta); err != nil {
		return "", err
	}
	return strings.TrimSpace(meta.Title), nil
}

// Archive downloads the video using yt-dlp and saves files to storage
func (a *YtDlpArchiver) Archive(ctx context.Context, link *model.Link, stor storage.Storage) (*ArchiveResult, error) {
	if a.binaryPath == "" {
		if path, err := exec.LookPath("yt-dlp"); err != nil {
			return nil, fmt.Errorf("yt-dlp not found in PATH: %w", err)
		} else {
			a.binaryPath = path
		}
	}

	timeout := 10 * time.Minute
	if a.config.Timeout != "" {
		if d, err := time.ParseDuration(a.config.Timeout); err == nil {
			timeout = d
		}
	}

	title, _ := a.fetchTitle(ctx, link.URL)

	tempDir, err := os.MkdirTemp("", "hako-yt-dlp-*")
	if err != nil {
		return nil, fmt.Errorf("failed to create temp dir: %w", err)
	}
	defer func() { _ = os.RemoveAll(tempDir) }()

	outputTemplate := filepath.Join(tempDir, "%(title)s.%(ext)s")
	args := []string{
		"-o", outputTemplate,
		"--no-warnings",
		"--no-playlist",
	}
	if a.config.WriteThumbnail {
		args = append(args, "--write-thumbnail")
	}
	if a.config.WriteSubs {
		args = append(args, "--write-subs", "--write-auto-subs", "--sub-langs", "en.*")
	}
	if a.config.EmbedSubs {
		args = append(args, "--embed-subs", "--sub-langs", "en.*")
	}
	args = append(args, link.URL)

	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	cmd := exec.CommandContext(ctx, a.binaryPath, args...)
	var combinedOutput strings.Builder
	cmd.Stdout = &combinedOutput
	cmd.Stderr = &combinedOutput

	if err := cmd.Run(); err != nil {
		errMsg := err.Error()
		if exitErr, ok := err.(*exec.ExitError); ok {
			errMsg = fmt.Sprintf("exit code %d", exitErr.ExitCode())
		}
		return nil, fmt.Errorf("yt-dlp failed (%s): %s", errMsg, strings.TrimSpace(combinedOutput.String()))
	}

	var files []FileInfo
	err = filepath.Walk(tempDir, func(path string, info os.FileInfo, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if info.IsDir() || path == tempDir {
			return nil
		}
		rel, _ := filepath.Rel(tempDir, path)
		if rel == "" || rel == "." {
			return nil
		}
		filename := filepath.Base(path)

		f, err := os.Open(path)
		if err != nil {
			return fmt.Errorf("open %s: %w", filename, err)
		}
		hash := sha256.New()
		tee := io.TeeReader(f, hash)
		storagePath, err := stor.Save(ctx, link.ID, filename, tee)
		_ = f.Close()
		if err != nil {
			return fmt.Errorf("save %s: %w", filename, err)
		}

		mimeType := mimeTypeFromExt(filepath.Ext(filename))
		// Add +thumbnail suffix to image files when WriteThumbnail is enabled
		if a.config.WriteThumbnail && isImageMimeType(mimeType) {
			mimeType = mimeType + "+thumbnail"
		}
		files = append(files, FileInfo{
			Filename:   filename,
			MimeType:   mimeType,
			FileSize:   info.Size(),
			HashSha256: hex.EncodeToString(hash.Sum(nil)),
			Path:       storagePath,
		})
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("processing downloaded files: %w", err)
	}

	if len(files) == 0 {
		return nil, fmt.Errorf("yt-dlp produced no files. output: %s", strings.TrimSpace(combinedOutput.String()))
	}

	if title == "" {
		title = link.URL
	}

	return &ArchiveResult{
		Title: title,
		Files: files,
	}, nil
}

func mimeTypeFromExt(ext string) string {
	ext = strings.ToLower(ext)
	m := map[string]string{
		".mp4": "video/mp4", ".webm": "video/webm", ".mkv": "video/x-matroska",
		".m4a": "audio/mp4", ".opus": "audio/opus",
		".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".webp": "image/webp",
		".vtt": "text/vtt", ".srt": "text/plain", ".ass": "text/x-ssa",
	}
	if mt, ok := m[ext]; ok {
		return mt
	}
	return "application/octet-stream"
}

func isImageMimeType(mimeType string) bool {
	return strings.HasPrefix(mimeType, "image/")
}