smartbar/internal/pdfassets/pdfassets.go

package pdfassets

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"image"
	stddraw "image/draw"
	"image/jpeg"
	"image/png"
	"io"
	"io/fs"
	"math"
	"os"
	"path/filepath"
	"strings"

	xdraw "golang.org/x/image/draw"
)

type Options struct {
	MaxWidth    int
	MaxHeight   int
	JPEGQuality int
}

type Stats struct {
	ImageFiles      int
	OptimizedImages int
	OriginalBytes   int64
	OutputBytes     int64
}

type PreparedTree struct {
	Root    string
	Stats   Stats
	Cleanup func() error
}

func OptionsForPage(widthIn, heightIn float64, imageDPI, jpegQuality int) Options {
	if imageDPI <= 0 {
		return Options{JPEGQuality: clampJPEGQuality(jpegQuality)}
	}
	return Options{
		MaxWidth:    max(1, int(math.Ceil(widthIn*float64(imageDPI)))),
		MaxHeight:   max(1, int(math.Ceil(heightIn*float64(imageDPI)))),
		JPEGQuality: clampJPEGQuality(jpegQuality),
	}
}

func (o Options) Enabled() bool {
	return o.MaxWidth > 0 && o.MaxHeight > 0
}

func PrepareTree(srcRoot string, opts Options) (PreparedTree, error) {
	if !opts.Enabled() {
		return PreparedTree{
			Root:    srcRoot,
			Cleanup: func() error { return nil },
		}, nil
	}

	absRoot, err := filepath.Abs(srcRoot)
	if err != nil {
		return PreparedTree{}, err
	}

	tmpRoot, err := os.MkdirTemp("", "smartbar_pdf_assets_*")
	if err != nil {
		return PreparedTree{}, err
	}

	prepared := PreparedTree{
		Root:    tmpRoot,
		Cleanup: func() error { return os.RemoveAll(tmpRoot) },
	}

	err = filepath.WalkDir(absRoot, func(path string, d fs.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}

		rel, err := filepath.Rel(absRoot, path)
		if err != nil {
			return err
		}
		if rel == "." {
			return nil
		}

		dst := filepath.Join(tmpRoot, rel)

		if d.IsDir() {
			return os.MkdirAll(dst, 0o755)
		}

		info, err := d.Info()
		if err != nil {
			return err
		}

		if info.Mode()&fs.ModeSymlink != 0 {
			return copySymlink(path, dst)
		}

		if !info.Mode().IsRegular() {
			return nil
		}

		if strings.EqualFold(filepath.Ext(path), ".pdf") {
			return nil
		}

		if isSupportedImage(path) {
			stats, err := optimizeImageFile(path, dst, info.Mode().Perm(), opts)
			if err == nil {
				prepared.Stats.ImageFiles++
				prepared.Stats.OriginalBytes += stats.originalBytes
				prepared.Stats.OutputBytes += stats.outputBytes
				if stats.optimized {
					prepared.Stats.OptimizedImages++
				}
				return nil
			}
		}

		return copyFile(path, dst, info.Mode().Perm())
	})
	if err != nil {
		_ = prepared.Cleanup()
		return PreparedTree{}, err
	}

	return prepared, nil
}

type imageFileStats struct {
	originalBytes int64
	outputBytes   int64
	optimized     bool
}

func optimizeImageFile(src, dst string, mode fs.FileMode, opts Options) (imageFileStats, error) {
	data, err := os.ReadFile(src)
	if err != nil {
		return imageFileStats{}, err
	}

	optimized, changed, err := optimizeImage(data, filepath.Ext(src), opts)
	if err != nil {
		return imageFileStats{}, err
	}

	if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
		return imageFileStats{}, err
	}

	output := optimized
	if !changed {
		output = data
	}
	if err := os.WriteFile(dst, output, mode); err != nil {
		return imageFileStats{}, err
	}

	return imageFileStats{
		originalBytes: int64(len(data)),
		outputBytes:   int64(len(output)),
		optimized:     changed,
	}, nil
}

func optimizeImage(data []byte, ext string, opts Options) ([]byte, bool, error) {
	img, format, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, false, err
	}

	orientation := 1
	if format == "jpeg" || strings.EqualFold(ext, ".jpg") || strings.EqualFold(ext, ".jpeg") {
		orientation = jpegOrientation(data)
		img = applyOrientation(img, orientation)
	}

	srcBounds := img.Bounds()
	dstW, dstH := fitWithin(srcBounds.Dx(), srcBounds.Dy(), opts.MaxWidth, opts.MaxHeight)
	resized := dstW != srcBounds.Dx() || dstH != srcBounds.Dy()
	if resized {
		dst := image.NewRGBA(image.Rect(0, 0, dstW, dstH))
		xdraw.CatmullRom.Scale(dst, dst.Bounds(), img, srcBounds, stddraw.Src, nil)
		img = dst
	}

	var buf bytes.Buffer
	if imageHasTransparency(img) {
		encoder := png.Encoder{CompressionLevel: png.BestCompression}
		err = encoder.Encode(&buf, img)
	} else {
		err = jpeg.Encode(&buf, img, &jpeg.Options{Quality: clampJPEGQuality(opts.JPEGQuality)})
	}
	if err != nil {
		return nil, false, err
	}

	changed := resized || orientation != 1 || buf.Len() < len(data)
	if !changed {
		return data, false, nil
	}
	return buf.Bytes(), true, nil
}

func fitWithin(width, height, maxWidth, maxHeight int) (int, int) {
	if width <= 0 || height <= 0 || maxWidth <= 0 || maxHeight <= 0 {
		return width, height
	}
	if width <= maxWidth && height <= maxHeight {
		return width, height
	}

	scale := math.Min(float64(maxWidth)/float64(width), float64(maxHeight)/float64(height))
	return max(1, int(math.Round(float64(width)*scale))), max(1, int(math.Round(float64(height)*scale)))
}

func imageHasTransparency(img image.Image) bool {
	switch m := img.(type) {
	case *image.YCbCr:
		return false
	case *image.Gray, *image.Gray16, *image.CMYK:
		return false
	case *image.NRGBA:
		for i := 3; i < len(m.Pix); i += 4 {
			a := m.Pix[i]
			if a != 0xff {
				return true
			}
		}
		return false
	case *image.RGBA:
		for i := 3; i < len(m.Pix); i += 4 {
			a := m.Pix[i]
			if a != 0xff {
				return true
			}
		}
		return false
	case *image.Alpha, *image.Alpha16:
		return true
	}

	b := img.Bounds()
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			_, _, _, a := img.At(x, y).RGBA()
			if a != 0xffff {
				return true
			}
		}
	}
	return false
}

func applyOrientation(img image.Image, orientation int) image.Image {
	if orientation < 2 || orientation > 8 {
		return img
	}

	b := img.Bounds()
	w, h := b.Dx(), b.Dy()
	dstW, dstH := w, h
	if orientation >= 5 {
		dstW, dstH = h, w
	}

	dst := image.NewRGBA(image.Rect(0, 0, dstW, dstH))
	for y := 0; y < dstH; y++ {
		for x := 0; x < dstW; x++ {
			sx, sy := orientedSourcePoint(x, y, w, h, orientation)
			dst.Set(x, y, img.At(b.Min.X+sx, b.Min.Y+sy))
		}
	}
	return dst
}

func orientedSourcePoint(x, y, width, height, orientation int) (int, int) {
	switch orientation {
	case 2:
		return width - 1 - x, y
	case 3:
		return width - 1 - x, height - 1 - y
	case 4:
		return x, height - 1 - y
	case 5:
		return y, x
	case 6:
		return y, height - 1 - x
	case 7:
		return width - 1 - y, height - 1 - x
	case 8:
		return width - 1 - y, x
	default:
		return x, y
	}
}

func jpegOrientation(data []byte) int {
	if len(data) < 4 || data[0] != 0xff || data[1] != 0xd8 {
		return 1
	}

	for i := 2; i+4 <= len(data); {
		if data[i] != 0xff {
			i++
			continue
		}
		for i < len(data) && data[i] == 0xff {
			i++
		}
		if i >= len(data) {
			return 1
		}

		marker := data[i]
		i++
		if marker == 0xd9 || marker == 0xda {
			return 1
		}
		if marker >= 0xd0 && marker <= 0xd7 {
			continue
		}
		if i+2 > len(data) {
			return 1
		}

		segmentLen := int(binary.BigEndian.Uint16(data[i : i+2]))
		i += 2
		if segmentLen < 2 || i+segmentLen-2 > len(data) {
			return 1
		}

		payload := data[i : i+segmentLen-2]
		if marker == 0xe1 && bytes.HasPrefix(payload, []byte("Exif\x00\x00")) {
			return tiffOrientation(payload[6:])
		}
		i += segmentLen - 2
	}

	return 1
}

func tiffOrientation(data []byte) int {
	if len(data) < 8 {
		return 1
	}

	var order binary.ByteOrder
	switch string(data[:2]) {
	case "II":
		order = binary.LittleEndian
	case "MM":
		order = binary.BigEndian
	default:
		return 1
	}
	if order.Uint16(data[2:4]) != 42 {
		return 1
	}

	ifdOffset := int(order.Uint32(data[4:8]))
	if ifdOffset < 0 || ifdOffset+2 > len(data) {
		return 1
	}

	entryCount := int(order.Uint16(data[ifdOffset : ifdOffset+2]))
	pos := ifdOffset + 2
	for i := 0; i < entryCount; i++ {
		if pos+12 > len(data) {
			return 1
		}

		tag := order.Uint16(data[pos : pos+2])
		fieldType := order.Uint16(data[pos+2 : pos+4])
		count := order.Uint32(data[pos+4 : pos+8])
		if tag == 0x0112 && fieldType == 3 && count == 1 {
			value := int(order.Uint16(data[pos+8 : pos+10]))
			if value >= 1 && value <= 8 {
				return value
			}
			return 1
		}
		pos += 12
	}

	return 1
}

func isSupportedImage(path string) bool {
	switch strings.ToLower(filepath.Ext(path)) {
	case ".jpg", ".jpeg", ".png":
		return true
	default:
		return false
	}
}

func copyFile(src, dst string, mode fs.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()

	if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
		return err
	}

	out, err := os.OpenFile(dst, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, mode)
	if err != nil {
		return err
	}

	if _, err := io.Copy(out, in); err != nil {
		_ = out.Close()
		return err
	}
	return out.Close()
}

func copySymlink(src, dst string) error {
	target, err := os.Readlink(src)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
		return err
	}
	return os.Symlink(target, dst)
}

func clampJPEGQuality(q int) int {
	if q <= 0 {
		return 82
	}
	if q < 1 {
		return 1
	}
	if q > 100 {
		return 100
	}
	return q
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func FormatBytes(n int64) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := int64(unit), 0
	for n >= unit*div && exp < 4 {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(n)/float64(div), "KMGTPE"[exp])
}