// Package thumbs makes the Files tab's previews: one small JPEG for an image or a video,
// by ffmpeg, decoded on a GPU where the machine has one it can use and on the CPU
// otherwise.
//
// Nothing about the machine is assumed. A probe looks ffmpeg and ffprobe up, opens every
// /dev/dri/renderD* and exercises it (a VAAPI upload and download of a synthetic frame),
// so a node the daemon's user cannot open, or a driver that does not speak VAAPI, is
// reported once rather than tried per file. vainfo, where installed, says which codecs
// each node decodes, and a video goes to the first node that lists its codec; without
// vainfo the first working node is tried and ffmpeg falls back to software on its own when
// the stream turns out unsupported. A hardware run that fails outright is retried on the
// CPU before the file is called undecodable. Images are always decoded on the CPU: one
// frame of a JPEG at a reduced DCT size costs less than the trip through a GPU.
package thumbs

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"hamlab/internal/metrics"
	"hamlab/internal/priv"
)

// GPU is one hardware decoder as the probe found it.
type GPU struct {
	API        string   `json:"api"`                  // vaapi | cuda
	Device     string   `json:"device"`               // /dev/dri/renderD128; "0" for cuda (its device index)
	Name       string   `json:"name,omitempty"`       // "Intel Arc A380"
	Codecs     []string `json:"codecs,omitempty"`     // ffmpeg codec names it decodes; nil when nothing could say
	Reason     string   `json:"reason,omitempty"`     // Unusable only: why
	Permission bool     `json:"permission,omitempty"` // Unusable only: the node exists but this user may not open it
}

// Caps is what the machine can make previews with, as probed.
type Caps struct {
	FFmpeg   string    `json:"ffmpeg,omitempty"`
	FFprobe  string    `json:"ffprobe,omitempty"`
	Nice     string    `json:"nice,omitempty"`
	Version  string    `json:"version,omitempty"`
	GPUs     []GPU     `json:"gpus"`               // usable, in preference order
	Unusable []GPU     `json:"unusable,omitempty"` // render nodes that exist but cannot be used, with the reason
	Error    string    `json:"error,omitempty"`    // why nothing can be made at all
	ProbedAt time.Time `json:"probed_at"`
}

// Opts shapes one preview.
type Opts struct {
	Size    int           // the long edge in pixels; 256 when zero
	Timeout time.Duration // for the whole job, ffprobe and every ffmpeg attempt; 30 s when zero
}

func (o Opts) defaults() Opts {
	if o.Size <= 0 {
		o.Size = 256
	}
	if o.Timeout <= 0 {
		o.Timeout = 30 * time.Second
	}
	return o
}

// Result is a finished preview.
type Result struct {
	JPEG   []byte `json:"-"`
	Engine string `json:"engine"` // "cpu", or "vaapi:/dev/dri/renderD129"
	Codec  string `json:"codec"`
	Width  int    `json:"width"`
	Height int    `json:"height"`
}

// ErrUnsupported: not an image or a video by its name.
var ErrUnsupported = errors.New("not an image or video")

// A DecodeError is ffmpeg's verdict on the file itself. It will not change until the
// file does, so the caller may remember it.
type DecodeError struct{ Msg string }

func (e *DecodeError) Error() string { return e.Msg }

// ---- what gets a preview ----

var imageExt = set("png jpg jpeg jpe jfif gif webp bmp tif tiff avif heic heif hif jxl svg ico psd tga thm")
var videoExt = set("mp4 m4v mkv webm mov avi wmv flv mpg mpeg mpe m2v ts m2ts mts 3gp 3g2 ogv ogm mxf vob asf rm rmvb divx f4v mjpeg mjpg 360 osv insv lrf")
var audioExt = set("mp3 flac m4a m4b aac ogg oga opus wav wma aiff aif alac ape wv mka")

func set(s string) map[string]bool {
	m := map[string]bool{}
	for _, w := range strings.Fields(s) {
		m[w] = true
	}
	return m
}

// Kind is "image", "video", "audio" or "" for a file name. An audio file's preview is its
// cover art, when it carries any.
func Kind(name string) string {
	ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(name), "."))
	switch {
	case imageExt[ext]:
		return "image"
	case videoExt[ext]:
		return "video"
	case audioExt[ext]:
		return "audio"
	}
	return ""
}

// ---- the probe ----

// Probe finds out what this process can make previews with. It runs ffmpeg against every
// render node, so it takes a few hundred milliseconds; callers keep the answer.
func Probe(ctx context.Context) Caps {
	c := Caps{ProbedAt: time.Now(), GPUs: []GPU{}}
	c.FFmpeg, _ = exec.LookPath("ffmpeg")
	c.FFprobe, _ = exec.LookPath("ffprobe")
	c.Nice, _ = exec.LookPath("nice")
	switch {
	case c.FFmpeg == "" && c.FFprobe == "":
		c.Error = "ffmpeg is not installed"
	case c.FFmpeg == "":
		c.Error = "ffmpeg is not installed (ffprobe is)"
	case c.FFprobe == "":
		c.Error = "ffprobe is not installed (ffmpeg is)"
	}
	if c.Error != "" {
		return c
	}
	c.Version = ffmpegVersion(ctx, c.FFmpeg)
	accels := hwaccels(ctx, c.FFmpeg)
	if accels["vaapi"] {
		nodes, _ := filepath.Glob("/dev/dri/renderD*")
		sort.Strings(nodes)
		vainfo, _ := exec.LookPath("vainfo")
		for _, dev := range nodes {
			g := GPU{API: "vaapi", Device: dev, Name: metrics.RenderNodeName(dev)}
			if err := openable(dev); err != nil {
				g.Reason, g.Permission = err.Error(), errors.Is(err, os.ErrPermission)
				c.Unusable = append(c.Unusable, g)
				continue
			}
			if msg := hwWorks(ctx, c.FFmpeg, "vaapi=va:"+dev); msg != "" {
				g.Reason = msg
				c.Unusable = append(c.Unusable, g)
				continue
			}
			if vainfo != "" {
				g.Codecs = vainfoCodecs(ctx, vainfo, dev)
			}
			c.GPUs = append(c.GPUs, g)
		}
	}
	if accels["cuda"] && exists("/dev/nvidiactl") {
		g := GPU{API: "cuda", Device: "0", Name: "NVIDIA GPU"}
		if msg := hwWorks(ctx, c.FFmpeg, "cuda=cu:0"); msg != "" {
			g.Reason = msg
			c.Unusable = append(c.Unusable, g)
		} else {
			c.GPUs = append(c.GPUs, g)
		}
	}
	return c
}

// openable: the node can be opened read-write, which is what a VAAPI driver needs. The
// usual reason it can't be is the owning group: a render node is root:render 0660 on
// most distributions, and a daemon running as an ordinary user is not in render.
func openable(dev string) error {
	f, err := os.OpenFile(dev, os.O_RDWR, 0)
	if err != nil {
		if errors.Is(err, os.ErrPermission) {
			owner := ""
			if st, serr := os.Stat(dev); serr == nil {
				owner = ownerOf(st)
			}
			if owner != "" {
				return fmt.Errorf("%w: %s is %s", os.ErrPermission, dev, owner)
			}
			return fmt.Errorf("%w: cannot open %s", os.ErrPermission, dev)
		}
		return err
	}
	f.Close()
	return nil
}

// hwWorks initialises a hardware device in ffmpeg and moves one synthetic frame through
// it. "" when it works, else the last thing ffmpeg said.
func hwWorks(ctx context.Context, ffmpeg, device string) string {
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	name, rest, _ := strings.Cut(device, "=") // "vaapi", "va:/dev/…"
	label, _, _ := strings.Cut(rest, ":")     // "va"
	cmd := exec.CommandContext(ctx, ffmpeg, "-nostdin", "-hide_banner", "-v", "error", "-init_hw_device", device, "-filter_hw_device", label,
		"-f", "lavfi", "-i", "color=c=black:s=64x64:r=1:d=1", "-vf", "format=nv12,hwupload,hwdownload,format=nv12", "-frames:v", "1", "-f", "null", "-")
	var errb bytes.Buffer
	cmd.Stderr = &errb
	if err := cmd.Run(); err != nil {
		msg := lastLine(errb.String())
		if msg == "" {
			msg = err.Error()
		}
		return name + ": " + msg
	}
	return ""
}

func ffmpegVersion(ctx context.Context, ffmpeg string) string {
	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	out, _ := exec.CommandContext(ctx, ffmpeg, "-version").Output()
	f := strings.Fields(string(out)) // "ffmpeg version 7.1.5-0+deb13u1 Copyright …"
	if len(f) >= 3 && f[1] == "version" {
		v := f[2]
		if i := strings.IndexAny(v, "-+~"); i > 0 {
			v = v[:i]
		}
		return v
	}
	return ""
}

func hwaccels(ctx context.Context, ffmpeg string) map[string]bool {
	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	out, _ := exec.CommandContext(ctx, ffmpeg, "-hide_banner", "-hwaccels").Output()
	m := map[string]bool{}
	for _, l := range strings.Split(string(out), "\n") {
		l = strings.TrimSpace(l)
		if l != "" && !strings.HasPrefix(l, "Hardware acceleration methods") {
			m[l] = true
		}
	}
	return m
}

var vaProfileRe = regexp.MustCompile(`VAProfile([A-Za-z0-9]+?)(?:Simple|Main|High|Profile|Baseline|Constrained|Scc|Advanced|Version|Stereo|Multiview)\w*\s*:\s*VAEntrypointVLD`)

// vaCodec maps a VAProfile family to the ffmpeg codec it decodes.
var vaCodec = map[string]string{"H264": "h264", "HEVC": "hevc", "VP9": "vp9", "AV1": "av1", "MPEG2": "mpeg2video", "VC1": "vc1", "VP8": "vp8", "JPEG": "mjpeg", "MPEG4": "mpeg4", "H263": "h263"}

func vainfoCodecs(ctx context.Context, vainfo, dev string) []string {
	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	out, _ := exec.CommandContext(ctx, vainfo, "--display", "drm", "--device", dev).Output()
	return parseVainfo(string(out))
}

// parseVainfo keeps the decode (VLD) profiles and names their codecs the way ffmpeg does.
func parseVainfo(text string) []string {
	seen := map[string]bool{}
	for _, m := range vaProfileRe.FindAllStringSubmatch(text, -1) {
		if c, ok := vaCodec[m[1]]; ok {
			seen[c] = true
		}
	}
	out := make([]string, 0, len(seen))
	for c := range seen {
		out = append(out, c)
	}
	sort.Strings(out)
	return out
}

// pick is the GPU for a codec: the first that lists it, else the first that could not
// say (ffmpeg will fall back to software by itself if that guess is wrong), else none.
func (c Caps) pick(codec string) *GPU {
	var unknown *GPU
	for i := range c.GPUs {
		g := &c.GPUs[i]
		if g.Codecs == nil {
			if unknown == nil {
				unknown = g
			}
			continue
		}
		for _, x := range g.Codecs {
			if x == codec {
				return g
			}
		}
	}
	return unknown
}

// ---- making one ----

type fileInfo struct {
	index    string // the stream to take
	codec    string
	w, h     int
	duration float64
}

type attempt struct {
	gpu    *GPU
	lowres int
	seek   float64
}

func (a attempt) engine() string {
	if a.gpu == nil {
		return "cpu"
	}
	return a.gpu.API + ":" + a.gpu.Device
}

// Make produces the preview of one file.
func Make(ctx context.Context, c Caps, path string, o Opts) (Result, error) {
	kind := Kind(path)
	if kind == "" {
		return Result{}, ErrUnsupported
	}
	if c.Error != "" {
		return Result{}, errors.New(c.Error)
	}
	o = o.defaults()
	ctx, cancel := context.WithTimeout(ctx, o.Timeout)
	defer cancel()
	fi, err := probeFile(ctx, c, path)
	if err != nil {
		return Result{}, err
	}
	var last string
	for _, eng := range engines(c, kind, fi) {
		for _, a := range attempts(eng, kind, fi, o.Size) {
			data, msg, err := run(ctx, c, a, path, fi, o.Size)
			if err == nil && len(data) > 0 {
				return Result{JPEG: data, Engine: a.engine(), Codec: fi.codec, Width: fi.w, Height: fi.h}, nil
			}
			if ctx.Err() != nil {
				return Result{}, ctx.Err() // the clock ran out: not the file's fault
			}
			if err != nil { // the decoder refused: try the next engine, not another seek
				last = msg
				break
			}
			last = "no frame at " + strconv.FormatFloat(a.seek, 'f', 1, 64) + " s"
		}
	}
	if last == "" {
		last = "could not decode"
	}
	return Result{}, &DecodeError{Msg: last}
}

// engines is the order to try: the fitting GPU first for a video, then the CPU.
func engines(c Caps, kind string, fi fileInfo) []*GPU {
	if kind == "video" {
		if g := c.pick(fi.codec); g != nil {
			return []*GPU{g, nil}
		}
	}
	return []*GPU{nil}
}

// attempts for one engine: a seek into a video, then from its start if that found no
// frame (a wrong duration in the header); a reduced-size decode of a JPEG, then full size.
func attempts(gpu *GPU, kind string, fi fileInfo, size int) []attempt {
	if kind == "video" {
		s := seekFor(fi.duration)
		if s > 0 {
			return []attempt{{gpu: gpu, seek: s}, {gpu: gpu}}
		}
		return []attempt{{gpu: gpu}}
	}
	if fi.codec == "mjpeg" && gpu == nil {
		if l := lowresFor(fi.w, fi.h, size); l > 0 {
			return []attempt{{lowres: l}, {}}
		}
	}
	return []attempt{{gpu: gpu}}
}

// seekFor: a tenth of the way in, never more than 30 s, so a title card or a black lead-in
// is skipped and a long file still seeks to a keyframe near its start.
func seekFor(duration float64) float64 {
	if duration <= 0 {
		return 0
	}
	s := duration / 10
	if s > 30 {
		s = 30
	}
	if s < 0.5 {
		return 0
	}
	return s
}

// lowresFor: how many halvings a JPEG decoder may take while the picture stays at least
// twice the preview size. libjpeg does this in the DCT, so a 24 MP photo decodes at 1.5 MP.
func lowresFor(w, h, size int) int {
	m := w
	if h < m {
		m = h
	}
	l := 0
	for l < 3 && m>>(l+1) >= size*2 {
		l++
	}
	return l
}

type probeOut struct {
	Streams []struct {
		Index       int    `json:"index"`
		Codec       string `json:"codec_name"`
		Width       int    `json:"width"`
		Height      int    `json:"height"`
		Disposition struct {
			AttachedPic int `json:"attached_pic"`
		} `json:"disposition"`
	} `json:"streams"`
	Format struct {
		Duration string `json:"duration"`
	} `json:"format"`
}

// probeFile asks ffprobe which stream to take and what is in it.
func probeFile(ctx context.Context, c Caps, path string) (fileInfo, error) {
	cmd := exec.CommandContext(ctx, c.FFprobe, "-v", "error", "-select_streams", "v", "-show_entries", "stream=index,codec_name,width,height,disposition:format=duration", "-of", "json", path)
	var errb bytes.Buffer
	cmd.Stderr = &errb
	out, err := cmd.Output()
	if err != nil {
		if ctx.Err() != nil {
			return fileInfo{}, ctx.Err()
		}
		msg := lastLine(errb.String())
		if msg == "" {
			msg = "ffprobe: " + err.Error()
		}
		return fileInfo{}, &DecodeError{Msg: msg}
	}
	var po probeOut
	if err := json.Unmarshal(out, &po); err != nil {
		return fileInfo{}, &DecodeError{Msg: "ffprobe: unreadable answer"}
	}
	if len(po.Streams) == 0 {
		return fileInfo{}, &DecodeError{Msg: "no video stream"}
	}
	s := po.Streams[0]
	for _, x := range po.Streams { // cover art rides along as a video stream; the picture is not the film
		if x.Disposition.AttachedPic == 0 {
			s = x
			break
		}
	}
	fi := fileInfo{index: strconv.Itoa(s.Index), codec: s.Codec, w: s.Width, h: s.Height}
	fi.duration, _ = strconv.ParseFloat(po.Format.Duration, 64)
	return fi, nil
}

// run is one ffmpeg: the JPEG on success; on failure what ffmpeg said and the error.
// Empty output with a clean exit means the seek found nothing.
func run(ctx context.Context, c Caps, a attempt, path string, fi fileInfo, size int) ([]byte, string, error) {
	args := []string{"-nostdin", "-hide_banner", "-loglevel", "error", "-y"}
	if a.gpu != nil {
		args = append(args, "-hwaccel", a.gpu.API, "-hwaccel_device", a.gpu.Device)
	}
	if a.lowres > 0 {
		args = append(args, "-lowres", strconv.Itoa(a.lowres))
	}
	if a.seek > 0 {
		args = append(args, "-noaccurate_seek", "-ss", strconv.FormatFloat(a.seek, 'f', 3, 64))
	}
	args = append(args, "-i", path, "-map", "0:"+fi.index, "-frames:v", "1",
		"-vf", fmt.Sprintf("scale=w=%d:h=%d:force_original_aspect_ratio=decrease:flags=bicubic", size, size),
		"-f", "image2pipe", "-c:v", "mjpeg", "-q:v", "4", "pipe:1")
	name, argv := c.FFmpeg, args
	if c.Nice != "" { // a burst of previews must not crowd out what the machine is for
		name, argv = c.Nice, append([]string{"-n", "10", c.FFmpeg}, args...)
	}
	cmd := exec.CommandContext(ctx, name, argv...)
	var out, errb bytes.Buffer
	cmd.Stdout, cmd.Stderr = &out, &errb
	if err := cmd.Run(); err != nil {
		msg := lastLine(errb.String())
		if msg == "" {
			msg = err.Error()
		}
		return nil, msg, err
	}
	return out.Bytes(), "", nil
}

// ---- the same, as root ----

// ProbeRoot is Probe run as root through this binary, for the Files tab's root mode: root
// opens every render node, whatever group the daemon's user is in.
func ProbeRoot(ctx context.Context, self string) (Caps, error) {
	full := priv.Wrap(self, "thumb", "--probe")
	cmd := exec.CommandContext(ctx, full[0], full[1:]...)
	var errb bytes.Buffer
	cmd.Stderr = &errb
	out, err := cmd.Output()
	if err != nil {
		if msg := lastLine(errb.String()); msg != "" {
			return Caps{}, errors.New(msg)
		}
		return Caps{}, err
	}
	var c Caps
	if err := json.Unmarshal(out, &c); err != nil {
		return Caps{}, err
	}
	return c, nil
}

type helperOut struct {
	Engine    string `json:"engine,omitempty"`
	Codec     string `json:"codec,omitempty"`
	Width     int    `json:"width,omitempty"`
	Height    int    `json:"height,omitempty"`
	JPEG      []byte `json:"jpeg,omitempty"`
	Error     string `json:"error,omitempty"`
	Permanent bool   `json:"permanent,omitempty"` // a DecodeError: the file, not the moment
}

// MakeRoot is Make run as root through this binary, with the caps ProbeRoot found.
func MakeRoot(ctx context.Context, self string, c Caps, path string, o Opts) (Result, error) {
	o = o.defaults()
	caps, _ := json.Marshal(c)
	full := priv.Wrap(self, "thumb", "--size", strconv.Itoa(o.Size), "--timeout", o.Timeout.String(), path)
	cmd := exec.CommandContext(ctx, full[0], full[1:]...)
	cmd.Stdin = bytes.NewReader(caps)
	var errb bytes.Buffer
	cmd.Stderr = &errb
	out, err := cmd.Output()
	if err != nil {
		if ctx.Err() != nil {
			return Result{}, ctx.Err()
		}
		if msg := lastLine(errb.String()); msg != "" {
			return Result{}, errors.New(msg)
		}
		return Result{}, err
	}
	var h helperOut
	if err := json.Unmarshal(out, &h); err != nil {
		return Result{}, err
	}
	if h.Error != "" {
		if h.Permanent {
			return Result{}, &DecodeError{Msg: h.Error}
		}
		return Result{}, errors.New(h.Error)
	}
	return Result{JPEG: h.JPEG, Engine: h.Engine, Codec: h.Codec, Width: h.Width, Height: h.Height}, nil
}

// Helper implements `hamlab thumb …`, the root side of the above. `--probe` prints Caps;
// `[--size N] [--timeout D] <path>` reads Caps from stdin and prints the preview as JSON.
// Returns the exit code.
func Helper(args []string, stdin io.Reader, stdout, stderr io.Writer) int {
	fail := func(err error) int { fmt.Fprintln(stderr, err.Error()); return 1 }
	ctx := context.Background()
	o := Opts{}
	path := ""
	for i := 0; i < len(args); i++ {
		switch a := args[i]; {
		case a == "--probe":
			return emit(stdout, Probe(ctx), stderr)
		case a == "--size" && i+1 < len(args):
			o.Size, _ = strconv.Atoi(args[i+1])
			i++
		case a == "--timeout" && i+1 < len(args):
			o.Timeout, _ = time.ParseDuration(args[i+1])
			i++
		case strings.HasPrefix(a, "-"):
			return fail(errors.New("usage: hamlab thumb --probe | [--size N] [--timeout D] <path>"))
		default:
			path = a
		}
	}
	if path == "" || !filepath.IsAbs(path) {
		return fail(errors.New("usage: hamlab thumb --probe | [--size N] [--timeout D] <absolute path>"))
	}
	var c Caps
	if err := json.NewDecoder(io.LimitReader(stdin, 1<<20)).Decode(&c); err != nil {
		return fail(fmt.Errorf("caps on stdin: %w", err))
	}
	res, err := Make(ctx, c, path, o)
	if err != nil {
		var de *DecodeError
		return emit(stdout, helperOut{Error: err.Error(), Permanent: errors.As(err, &de)}, stderr)
	}
	return emit(stdout, helperOut{Engine: res.Engine, Codec: res.Codec, Width: res.Width, Height: res.Height, JPEG: res.JPEG}, stderr)
}

func emit(w io.Writer, v any, stderr io.Writer) int {
	if err := json.NewEncoder(w).Encode(v); err != nil {
		fmt.Fprintln(stderr, err.Error())
		return 1
	}
	return 0
}

// ---- small helpers ----

func lastLine(s string) string {
	lines := strings.Split(strings.TrimSpace(s), "\n")
	return strings.TrimSpace(lines[len(lines)-1])
}

func exists(p string) bool { _, err := os.Stat(p); return err == nil }