vibedit/media.go

package main

import (
	"context"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Stream struct {
	Index         int     `json:"index"`
	Type          string  `json:"type"`
	Codec         string  `json:"codec"`
	Profile       string  `json:"profile,omitempty"`
	Width         int     `json:"width,omitempty"`
	Height        int     `json:"height,omitempty"`
	DisplayWidth  int     `json:"displayWidth,omitempty"` // after pixel aspect ratio and rotation
	DisplayHeight int     `json:"displayHeight,omitempty"`
	SquarePixels  bool    `json:"squarePixels,omitempty"`
	Rotation      int     `json:"rotation,omitempty"`
	PixFmt        string  `json:"pixFmt,omitempty"`
	FPS           float64 `json:"fps,omitempty"`
	FrameRate     string  `json:"frameRate,omitempty"` // rational, e.g. 30000/1001
	SampleRate    int     `json:"sampleRate,omitempty"`
	Channels      int     `json:"channels,omitempty"`
	ChannelLayout string  `json:"channelLayout,omitempty"`
	Language      string  `json:"language,omitempty"`
	Title         string  `json:"title,omitempty"`
}

type Media struct {
	ID         string    `json:"id"`
	Path       string    `json:"path"`
	Name       string    `json:"name"`
	Size       int64     `json:"size"`
	ModTime    time.Time `json:"modTime"`
	Duration   float64   `json:"duration"`
	StartTime  float64   `json:"startTime"`
	Format     string    `json:"format"`
	FormatLong string    `json:"formatLong"`
	BitRate    int64     `json:"bitRate"`
	Streams    []Stream  `json:"streams"`
	Video      *Stream   `json:"video"` // primary video stream, nil if none
	Audio      []Stream  `json:"audio"`
}

func (m *Media) FirstAudio() *Stream {
	if len(m.Audio) == 0 {
		return nil
	}
	return &m.Audio[0]
}

func mediaID(path string) string {
	h := sha1.Sum([]byte(path))
	return hex.EncodeToString(h[:6])
}

type ffprobeOutput struct {
	Format struct {
		FormatName     string `json:"format_name"`
		FormatLongName string `json:"format_long_name"`
		Duration       string `json:"duration"`
		StartTime      string `json:"start_time"`
		BitRate        string `json:"bit_rate"`
	} `json:"format"`
	Streams []struct {
		Index         int               `json:"index"`
		CodecType     string            `json:"codec_type"`
		CodecName     string            `json:"codec_name"`
		Profile       string            `json:"profile"`
		Width         int               `json:"width"`
		Height        int               `json:"height"`
		PixFmt        string            `json:"pix_fmt"`
		SAR           string            `json:"sample_aspect_ratio"`
		AvgFrameRate  string            `json:"avg_frame_rate"`
		RFrameRate    string            `json:"r_frame_rate"`
		SampleRate    string            `json:"sample_rate"`
		Channels      int               `json:"channels"`
		ChannelLayout string            `json:"channel_layout"`
		Duration      string            `json:"duration"`
		Tags          map[string]string `json:"tags"`
		Disposition   map[string]int    `json:"disposition"`
		SideData      []map[string]any  `json:"side_data_list"`
	} `json:"streams"`
}

func (t *Tools) Probe(ctx context.Context, path string) (*Media, error) {
	abs, err := filepath.Abs(path)
	if err != nil {
		return nil, err
	}
	st, err := os.Stat(abs)
	if err != nil {
		return nil, fmt.Errorf("file not found: %s", abs)
	}
	if st.IsDir() {
		return nil, fmt.Errorf("%s is a folder", abs)
	}
	out, err := command(ctx, t.FFprobe, "-v", "error", "-print_format", "json", "-show_format", "-show_streams", abs).Output()
	if err != nil {
		var ee interface{ ExitCode() int }
		if errors.As(err, &ee) {
			return nil, fmt.Errorf("ffprobe could not read %s - is it a video file?", filepath.Base(abs))
		}
		return nil, err
	}
	var p ffprobeOutput
	if err := json.Unmarshal(out, &p); err != nil {
		return nil, fmt.Errorf("unexpected ffprobe output: %w", err)
	}
	m := &Media{
		ID:         mediaID(abs),
		Path:       abs,
		Name:       filepath.Base(abs),
		Size:       st.Size(),
		ModTime:    st.ModTime(),
		Duration:   parseFloat(p.Format.Duration),
		StartTime:  parseFloat(p.Format.StartTime),
		Format:     p.Format.FormatName,
		FormatLong: p.Format.FormatLongName,
		BitRate:    int64(parseFloat(p.Format.BitRate)),
		Streams:    []Stream{},
		Audio:      []Stream{},
	}
	var maxStreamDur float64
	for _, s := range p.Streams {
		maxStreamDur = math.Max(maxStreamDur, parseFloat(s.Duration))
		ms := Stream{Index: s.Index, Type: s.CodecType, Codec: s.CodecName, Profile: s.Profile,
			Language: s.Tags["language"], Title: s.Tags["title"]}
		switch s.CodecType {
		case "video":
			ms.Width, ms.Height, ms.PixFmt = s.Width, s.Height, s.PixFmt
			ms.FrameRate, ms.FPS = pickFrameRate(s.AvgFrameRate, s.RFrameRate)
			ms.Rotation = streamRotation(s.Tags, s.SideData)
			sar := parseRatio(s.SAR, ":")
			ms.SquarePixels = sar <= 0 || math.Abs(sar-1) < 0.001
			dw, dh := float64(s.Width), float64(s.Height)
			if !ms.SquarePixels {
				dw *= sar
			}
			if ms.Rotation%180 != 0 {
				dw, dh = dh, dw
			}
			ms.DisplayWidth, ms.DisplayHeight = int(math.Round(dw)), int(math.Round(dh))
			if m.Video == nil && s.Disposition["attached_pic"] == 0 && s.Width > 0 {
				v := ms
				m.Video = &v
			}
		case "audio":
			ms.SampleRate = int(parseFloat(s.SampleRate))
			ms.Channels, ms.ChannelLayout = s.Channels, s.ChannelLayout
			m.Audio = append(m.Audio, ms)
		}
		m.Streams = append(m.Streams, ms)
	}
	if m.Duration <= 0 {
		m.Duration = maxStreamDur
	}
	if m.Video == nil && len(m.Audio) == 0 {
		return nil, fmt.Errorf("%s has no video or audio streams", m.Name)
	}
	if m.Duration <= 0 {
		return nil, fmt.Errorf("could not determine the duration of %s", m.Name)
	}
	return m, nil
}

func parseFloat(s string) float64 {
	f, err := strconv.ParseFloat(strings.TrimSpace(s), 64)
	if err != nil || math.IsNaN(f) || math.IsInf(f, 0) {
		return 0
	}
	return f
}

func parseRatio(s, sep string) float64 {
	a, b, ok := strings.Cut(s, sep)
	if !ok {
		return parseFloat(s)
	}
	num, den := parseFloat(a), parseFloat(b)
	if den == 0 {
		return 0
	}
	return num / den
}

// pickFrameRate prefers the nominal rate (r_frame_rate) when it agrees with the
// average, since containers like MKV report odd r_frame_rate values.
func pickFrameRate(avg, r string) (string, float64) {
	a, rr := parseRatio(avg, "/"), parseRatio(r, "/")
	switch {
	case rr >= 1 && rr <= 240 && (a <= 0 || math.Abs(rr-a)/rr < 0.1):
		return r, rr
	case a >= 1 && a <= 240:
		return avg, a
	case rr > 0:
		return r, rr
	}
	return "30", 30
}

func streamRotation(tags map[string]string, side []map[string]any) int {
	rot := 0.0
	if v, ok := tags["rotate"]; ok {
		rot = parseFloat(v)
	}
	for _, sd := range side {
		if v, ok := sd["rotation"].(float64); ok {
			rot = v
		}
	}
	r := int(math.Round(rot)) % 360
	if r < 0 {
		r += 360
	}
	return r
}

// --- MediaStore: probed sources, thumbnails, keyframes and preview proxies ---

type kfEntry struct {
	status string // pending, done, error
	times  []float64
	err    string
	done   chan struct{}
}

type MediaStore struct {
	tools *Tools
	jobs  *JobManager

	mu    sync.Mutex
	items map[string]*Media
	kf    map[string]*kfEntry

	kfSem    chan struct{}
	thumbSem chan struct{}
	thumbs   *byteCache
}

func newMediaStore(t *Tools, jobs *JobManager) *MediaStore {
	return &MediaStore{
		tools:    t,
		jobs:     jobs,
		items:    map[string]*Media{},
		kf:       map[string]*kfEntry{},
		kfSem:    make(chan struct{}, 2),
		thumbSem: make(chan struct{}, 3),
		thumbs:   newByteCache(96 << 20),
	}
}

func (s *MediaStore) Get(id string) *Media {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.items[id]
}

func (a *App) handleProbe(w http.ResponseWriter, r *http.Request) {
	if !a.tools.OK() {
		writeError(w, http.StatusServiceUnavailable, "ffmpeg/ffprobe not found")
		return
	}
	var req struct {
		Path string `json:"path"`
	}
	if err := readJSON(r, &req); err != nil || strings.TrimSpace(req.Path) == "" {
		writeError(w, http.StatusBadRequest, "missing path")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
	defer cancel()
	m, err := a.tools.Probe(ctx, strings.TrimSpace(req.Path))
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	a.media.mu.Lock()
	a.media.items[m.ID] = m
	a.media.mu.Unlock()
	writeJSON(w, m)
}

func (a *App) handleMedia(w http.ResponseWriter, r *http.Request) {
	m := a.media.Get(r.PathValue("id"))
	if m == nil {
		http.NotFound(w, r)
		return
	}
	path, name := m.Path, m.Name
	if r.URL.Query().Get("proxy") == "1" {
		path, name = a.media.proxyPath(m), "proxy.mp4"
	}
	f, err := os.Open(path)
	if err != nil {
		http.NotFound(w, r)
		return
	}
	defer f.Close()
	st, err := f.Stat()
	if err != nil {
		http.NotFound(w, r)
		return
	}
	// Browsers sniff media content; these types just keep them from refusing.
	switch strings.ToLower(filepath.Ext(name)) {
	case ".mp4", ".m4v", ".mov", ".3gp":
		w.Header().Set("Content-Type", "video/mp4")
	case ".mkv", ".webm":
		w.Header().Set("Content-Type", "video/webm")
	case ".ogv":
		w.Header().Set("Content-Type", "video/ogg")
	default:
		w.Header().Set("Content-Type", "application/octet-stream")
	}
	http.ServeContent(w, r, name, st.ModTime(), f)
}

// handleThumb renders a single frame as JPEG.
func (a *App) handleThumb(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	m := a.media.Get(q.Get("id"))
	if m == nil || m.Video == nil {
		http.NotFound(w, r)
		return
	}
	t := math.Max(0, math.Min(parseFloat(q.Get("t")), m.Duration-0.05))
	h, _ := strconv.Atoi(q.Get("h"))
	if h <= 0 || h > 1080 {
		h = 90
	}
	h -= h % 2
	key := fmt.Sprintf("%s|%.3f|%d", m.ID, t, h)
	img := a.media.thumbs.Get(key)
	if img == nil {
		select {
		case a.media.thumbSem <- struct{}{}:
		case <-r.Context().Done():
			return
		}
		ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
		out, err := command(ctx, a.tools.FFmpeg, "-hide_banner", "-loglevel", "error", "-nostdin",
			"-ss", ffTime(t), "-i", m.Path, "-map", fmt.Sprintf("0:%d", m.Video.Index), "-frames:v", "1",
			"-vf", fmt.Sprintf("scale=-2:%d", h), "-f", "image2pipe", "-c:v", "mjpeg", "-q:v", "5", "-").Output()
		cancel()
		<-a.media.thumbSem
		if err != nil || len(out) == 0 {
			http.Error(w, "could not render frame", http.StatusInternalServerError)
			return
		}
		img = out
		a.media.thumbs.Put(key, img)
	}
	w.Header().Set("Content-Type", "image/jpeg")
	w.Header().Set("Cache-Control", "private, max-age=3600")
	w.Write(img)
}

// handleKeyframes returns the keyframe times of the primary video stream,
// scanning the file in the background on first request.
func (a *App) handleKeyframes(w http.ResponseWriter, r *http.Request) {
	m := a.media.Get(r.URL.Query().Get("id"))
	if m == nil {
		writeError(w, http.StatusNotFound, "unknown source")
		return
	}
	if m.Video == nil {
		writeJSON(w, map[string]any{"status": "done", "times": []float64{}})
		return
	}
	s := a.media
	e := s.keyframeEntry(m)
	s.mu.Lock()
	resp := map[string]any{"status": e.status, "error": e.err}
	if e.status == "done" {
		resp["times"] = e.times
	}
	s.mu.Unlock()
	writeJSON(w, resp)
}

// keyframeEntry returns the keyframe scan for a source, starting it if needed.
func (s *MediaStore) keyframeEntry(m *Media) *kfEntry {
	s.mu.Lock()
	defer s.mu.Unlock()
	e := s.kf[m.ID]
	if e == nil {
		e = &kfEntry{status: "pending", done: make(chan struct{})}
		s.kf[m.ID] = e
		go s.scanKeyframes(m, e)
	}
	return e
}

// Keyframes waits for the keyframe scan of a source to finish.
func (s *MediaStore) Keyframes(ctx context.Context, m *Media) ([]float64, error) {
	if m.Video == nil {
		return nil, nil
	}
	e := s.keyframeEntry(m)
	select {
	case <-e.done:
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if e.status != "done" {
		return nil, errors.New(e.err)
	}
	return e.times, nil
}

func (s *MediaStore) scanKeyframes(m *Media, e *kfEntry) {
	s.kfSem <- struct{}{}
	defer func() { <-s.kfSem }()
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
	defer cancel()
	out, err := command(ctx, s.tools.FFprobe, "-v", "error", "-select_streams", strconv.Itoa(m.Video.Index),
		"-show_entries", "packet=pts_time,flags", "-of", "csv=p=0", m.Path).Output()
	var times []float64
	if err == nil {
		for _, line := range strings.Split(string(out), "\n") {
			pts, flags, ok := strings.Cut(strings.TrimSpace(line), ",")
			if !ok || !strings.HasPrefix(flags, "K") || pts == "N/A" {
				continue
			}
			t, perr := strconv.ParseFloat(pts, 64)
			if perr != nil {
				continue
			}
			times = append(times, math.Max(0, t-m.StartTime))
		}
		sort.Float64s(times)
		times = dedupe(times)
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	defer close(e.done)
	if err != nil {
		e.status, e.err = "error", err.Error()
		return
	}
	e.status, e.times = "done", times
}

func dedupe(sorted []float64) []float64 {
	out := sorted[:0]
	for i, t := range sorted {
		if i == 0 || t-out[len(out)-1] > 1e-4 {
			out = append(out, t)
		}
	}
	return out
}

// proxyPath is where the browser-friendly preview copy of a source lives.
func (s *MediaStore) proxyPath(m *Media) string {
	dir, err := os.UserCacheDir()
	if err != nil {
		dir = os.TempDir()
	}
	key := fmt.Sprintf("%s|%d|%d", m.Path, m.Size, m.ModTime.UnixNano())
	h := sha1.Sum([]byte(key))
	return filepath.Join(dir, "vibedit", "proxies", hex.EncodeToString(h[:10])+".mp4")
}

// handleProxy creates a small H.264 preview copy for sources the browser
// cannot play (AVI, WMV, unsupported codecs, ...).
func (a *App) handleProxy(w http.ResponseWriter, r *http.Request) {
	var req struct {
		ID string `json:"id"`
	}
	readJSON(r, &req)
	m := a.media.Get(req.ID)
	if m == nil {
		writeError(w, http.StatusNotFound, "unknown source")
		return
	}
	out := a.media.proxyPath(m)
	if _, err := os.Stat(out); err == nil {
		writeJSON(w, map[string]any{"ready": true})
		return
	}
	if j := a.jobs.Find("proxy", m.ID); j != nil {
		writeJSON(w, map[string]any{"jobId": j.ID})
		return
	}
	if !a.tools.HasEncoder("libx264") && a.tools.EncodersReady() {
		writeError(w, http.StatusInternalServerError, "this ffmpeg build has no H.264 encoder (libx264) for previews")
		return
	}
	j := a.jobs.Start("proxy", "Preview of "+m.Name, m.ID, func(ctx context.Context, c *JobCtl) error {
		if err := os.MkdirAll(filepath.Dir(out), 0o755); err != nil {
			return err
		}
		tmp := strings.TrimSuffix(out, ".mp4") + ".part.mp4"
		args := ffmpegBaseArgs()
		args = append(args, "-i", m.Path)
		if m.Video != nil {
			args = append(args, "-map", fmt.Sprintf("0:%d", m.Video.Index),
				"-vf", "scale=-2:'2*trunc(min(540,ih)/2)',format=yuv420p",
				"-c:v", "libx264", "-preset", "veryfast", "-crf", "26", "-g", "12", "-bf", "0")
		}
		if a := m.FirstAudio(); a != nil {
			args = append(args, "-map", fmt.Sprintf("0:%d", a.Index), "-c:a", "aac", "-b:a", "128k", "-ac", "2")
		}
		args = append(args, "-sn", "-dn", "-map_chapters", "-1", "-movflags", "+faststart", tmp)
		err := a.tools.runFFmpeg(ctx, args, func(sec float64, speed string) {
			c.Progress(sec/m.Duration, speed)
		})
		if err != nil {
			os.Remove(tmp)
			return err
		}
		return os.Rename(tmp, out)
	})
	writeJSON(w, map[string]any{"jobId": j.ID})
}

// byteCache is a size-bounded FIFO cache for rendered thumbnails.
type byteCache struct {
	mu    sync.Mutex
	max   int
	size  int
	order []string
	items map[string][]byte
}

func newByteCache(max int) *byteCache { return &byteCache{max: max, items: map[string][]byte{}} }

func (c *byteCache) Get(k string) []byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.items[k]
}

func (c *byteCache) Put(k string, v []byte) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if _, ok := c.items[k]; ok {
		return
	}
	c.items[k] = v
	c.order = append(c.order, k)
	c.size += len(v)
	for c.size > c.max && len(c.order) > 0 {
		old := c.order[0]
		c.order = c.order[1:]
		c.size -= len(c.items[old])
		delete(c.items, old)
	}
}