2 files changed, 43 insertions, 32 deletions

package thumbnailgen

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"os/exec"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"time"
)

type TimeFilter struct {
	Start float64
	End   float64
}

type Options struct {
	Interval      int
	Segments      int
	UseSegments   bool
	Format        string
	Scale         string
	EnableFilter  bool
	MaxThumbnails int
}

func NewDefaultOptions() Options {
	return Options{
		Interval:      5,
		Segments:      0,
		UseSegments:   false,
		Format:        "webp",
		Scale:         "-1:-1",
		EnableFilter:  true,
		MaxThumbnails: 0,
	}
}

func (options *Options) Apply(opts ...func(*Options)) *Options {
	for _, opt := range opts {
		opt(options)
	}

	return options
}

type videoStreamProbe struct {
	Streams []struct {
		CodecType string `json:"codec_type"`
		Duration  string `json:"duration"`
	} `json:"streams"`
}

func getThumbnailDuration(path string) (float64, error) {
	stderr := bytes.NewBuffer(nil)
	stdout := bytes.NewBuffer(nil)

	cmd := exec.Command(
		"ffprobe",
		"-v", "error",
		"-select_streams", "v:0",
		"-show_entries", "stream=codec_type,duration",
		"-of", "json",
		path,
	)
	cmd.Stdout = stdout
	cmd.Stderr = stderr

	if err := cmd.Run(); err != nil {
		return 0, fmt.Errorf(
			"probe video stream: %w: %s",
			err,
			strings.TrimSpace(stderr.String()),
		)
	}

	var probe videoStreamProbe
	if err := json.Unmarshal(stdout.Bytes(), &probe); err != nil {
		return 0, fmt.Errorf("decode video-stream probe: %w", err)
	}
	if len(probe.Streams) == 0 || probe.Streams[0].CodecType != "video" {
		return 0, errors.New("input has no video stream")
	}

	durationText := strings.TrimSpace(probe.Streams[0].Duration)
	if durationText == "" || durationText == "N/A" {
		return GetVideoLength(path) // container-duration fallback for animated WebP
	}

	duration, err := strconv.ParseFloat(durationText, 64)
	if err != nil {
		return 0, fmt.Errorf("parse video-stream duration %q: %w", durationText, err)
	}
	if duration <= 0 {
		return 0, fmt.Errorf("invalid video-stream duration %v", duration)
	}
	return duration, nil
}

func (opts *Options) GetThumbnail(path string) ([][]byte, int, error) {
	var filters []TimeFilter
	if opts.EnableFilter {
		f, err := GetFilter(path)
		if err != nil {
			return nil, 0, err
		}
		filters = f
	}

	length, err := getThumbnailDuration(path)
	if err != nil {
		return nil, 0, err
	}

	var interval int
	if opts.UseSegments {
		interval = int(length / float64(opts.Segments))
	} else {
		interval = opts.Interval
	}

	return GetThumbnailUnderlying(path, opts.MaxThumbnails, filters, length, interval, opts.EnableFilter, opts.Format, opts.Scale)
}

func GetThumbnailUnderlying(path string, maxThumbnails int, filters []TimeFilter, length float64, intervalSeconds int, enableFilter bool, format string, scale string) ([][]byte, int, error) {
	buf := bytes.NewBuffer(nil)
	framesExtracted := 0

	var out [][]byte
	if maxThumbnails > 0 {
		out = make([][]byte, maxThumbnails)
	} else {
		out = make([][]byte, 0)
	}

	var time float64 = 1
	for {
		if (maxThumbnails > 0 && framesExtracted >= maxThumbnails) || time >= length {
			break
		}

		if enableFilter {
			if ok, next := FrameLiesWithinFilter(time, filters); ok {
				time = next + 1
				continue
			}
		}

		err := GetImage(buf, path, int(time), format, scale)
		if err != nil {
			return nil, 0, err
		}

		b := bytes.Clone(buf.Bytes())[0:buf.Len()]

		if maxThumbnails > 0 {
			out[framesExtracted] = b
		} else {
			out = append(out, b)
		}

		time += float64(intervalSeconds)
		framesExtracted += 1
		buf.Reset()
	}
	return out, framesExtracted, nil
}

// Return true, if time is inside filter, also returns next time, when not in filter anymore.
// If False it returns input time
func FrameLiesWithinFilter(time float64, filters []TimeFilter) (bool, float64) {
	for _, filter := range filters {
		if time >= filter.Start && time <= filter.End {
			return true, filter.End
		}

	}
	return false, time
}

func GetImage(buf *bytes.Buffer, path string, timestamp int, format string, scale string) error {
	var t time.Time
	t = t.Add(time.Duration(timestamp) * time.Second)
	stderr := bytes.NewBuffer(nil)
	cmd := exec.Command(
		"ffmpeg",
		"-ss",
		t.Format("15:04:05"),
		"-i",
		path,
		"-vframes",
		"1",
		"-c:v",
		format,
		"-filter:v",
		fmt.Sprintf("scale=%s", scale),
		"-f",
		"image2pipe",
		"-",
	)
	cmd.Stdout = buf
	cmd.Stderr = stderr
	if err := cmd.Run(); err != nil {
		return fmt.Errorf("generate thumbnail at %s: %w: %s", t.Format("15:04:05"), err, strings.TrimSpace(stderr.String()))
	}
	return nil
}

func GetVideoLength(path string) (float64, error) {
	cmd := exec.Command("ffprobe", "-v", "quiet", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", path)
	stderr := bytes.NewBuffer(nil)
	buf := bytes.NewBuffer(nil)
	cmd.Stdout = buf
	cmd.Stderr = stderr
	if err := cmd.Run(); err != nil {
		return 0, fmt.Errorf(
			"probe video duration: %w: %s",
			err,
			strings.TrimSpace(stderr.String()),
		)
	}
	return strconv.ParseFloat(strings.ReplaceAll(buf.String(), "\n", ""), 64)
}

func escapeFilterValue(value string) string {
	return strings.NewReplacer(
		"\\", "\\\\",
		"'", "\\'",
		":", "\\:",
		",", "\\,",
		";", "\\;",
		"[", "\\[",
		"]", "\\]",
	).Replace(value)
}

func GetFilter(path string) ([]TimeFilter, error) {
	buf := bytes.NewBuffer(nil)
	stderr := bytes.NewBuffer(nil)
	filter := fmt.Sprintf("movie=filename='%s',blackdetect[out0]", escapeFilterValue(path))

	cmd := exec.Command("ffprobe", "-f", "lavfi", "-i", filter, "-show_entries", "tags=lavfi.black_start,lavfi.black_end", "-of", "default=nw=1", "-v", "quiet")
	cmd.Stderr = stderr
	cmd.Stdout = buf
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("detect black frames: %w: %s", err, strings.TrimSpace(stderr.String()))
	}
	filterStr := buf.String()
	filters := strings.Split(filterStr, "\n")
	filters = slices.Compact(filters)
	i := 0
	reg := regexp.MustCompile(`([0-9]*\.?[0-9]+)`)
	blackFilters := make([]TimeFilter, len(filters)/2)
	for {
		if i >= len(filters) {
			break
		}
		start := reg.FindString(filters[i])
		var end string
		if i+1 >= len(filters) || filters[i+1] == "" {
			end = start
		} else {
			end = reg.FindString(filters[i+1])
		}

		if start == "" || end == "" {
			break
		}

		s, err := strconv.ParseFloat(start, 64)
		if err != nil {
			return nil, err
		}

		e, err := strconv.ParseFloat(end, 64)
		if err != nil {
			return nil, err
		}

		blackFilters[i/2] = TimeFilter{
			Start: s,
			End:   e,
		}
		i += 2
	}

	return blackFilters, nil
}

func GetFramerate(path string) (float64, error) {
	cmd := exec.Command("ffprobe", "-v", "0", "-of", "csv=p=0", "-select_streams", "v:0", "-show_entries", "stream=r_frame_rate", path)
	buf := bytes.NewBuffer(nil)
	cmd.Stdout = buf
	err := cmd.Run()
	if err != nil {
		return 0, err
	}
	s := buf.String()
	s = strings.ReplaceAll(s, "\n", "")
	s = strings.TrimSpace(s)
	ss := strings.Split(s, "/")
	f1, err := strconv.Atoi(ss[0])
	if err != nil {
		return 0, err
	}

	f2, err := strconv.Atoi(ss[1])
	if err != nil {
		return 0, err
	}
	return float64(f1) / float64(f2), nil
}