TeleportFling/internal/protocol/jpeg.go

package protocol

/*
#cgo pkg-config: libturbojpeg
#include <turbojpeg.h>
#include <stdlib.h>

static tjhandle new_compressor(void) {
	return tj3Init(TJINIT_COMPRESS);
}
*/
import "C"
import (
	"errors"
	"image"
	"image/color"
	"runtime"
	"unsafe"
)

// JPEGEncoder wraps a TurboJPEG tj3 compressor. Not safe for concurrent
// use; callers must serialise or use one per goroutine.
type JPEGEncoder struct {
	ctx C.tjhandle
}

// NewJPEGEncoder initialises a TJ3 compressor. Must be freed via Close().
func NewJPEGEncoder() (*JPEGEncoder, error) {
	ctx := C.new_compressor()
	if ctx == nil {
		return nil, errors.New("turbojpeg: tj3Init failed")
	}
	// Allow turbojpeg to allocate the output buffer itself (no NOREALLOC).
	// This avoids the need to pin a Go output buffer and simplifies the
	// API: the caller receives a Go-owned copy and the C buffer is freed.
	return &JPEGEncoder{ctx: ctx}, nil
}

// Close destroys the underlying compressor.
func (e *JPEGEncoder) Close() {
	if e.ctx != nil {
		C.tj3Destroy(e.ctx)
		e.ctx = nil
	}
}

// Encode compresses img to JPEG at the given quality (1–100).
//
// Supported source types:
//   - *image.YCbCr – compressed via the YUV path (420/422/444, matching
//     obs-teleport exactly). SubsampleRatio is honoured.
//   - *image.RGBA – compressed as RGB via TJPF_RGBA, 444 subsampling.
//   - any other image.Image – converted to *image.RGBA then encoded as above.
//
// The returned byte slice is owned by the caller and must not be reused
// after the encoder is closed.
func (e *JPEGEncoder) Encode(img image.Image, quality int) ([]byte, error) {
	if quality < 1 {
		quality = 1
	}
	if quality > 100 {
		quality = 100
	}

	C.tj3Set(e.ctx, C.TJPARAM_QUALITY, C.int(quality))

	switch src := img.(type) {
	case *image.YCbCr:
		return e.encodeYCbCr(src)
	case *image.RGBA:
		return e.encodeRGBA(src)
	default:
		return e.encodeGeneric(img)
	}
}

// EncodeBGRA compresses a raw BGRA (blue, green, red, alpha) pixel buffer of
// the given dimensions. This is the fast path for the PipeWire screen-capture
// backend, which delivers frames in BGRA byte order. Subsampling defaults to
// 4:2:0 (a YCbCr JPEG), so the OBS receiver decodes it as I420 — the same
// layout obs-teleport produces for OBS-native YCbCr frames.
func (e *JPEGEncoder) EncodeBGRA(pix []byte, width, height, quality int) ([]byte, error) {
	if quality < 1 {
		quality = 1
	}
	if quality > 100 {
		quality = 100
	}

	want := width * height * 4
	if len(pix) < want {
		return nil, errors.New("turbojpeg: BGRA buffer too small")
	}

	C.tj3Set(e.ctx, C.TJPARAM_QUALITY, C.int(quality))
	C.tj3Set(e.ctx, C.TJPARAM_SUBSAMP, C.TJSAMP_420)
	C.tj3Set(e.ctx, C.TJPARAM_COLORSPACE, C.TJCS_YCbCr)

	size := C.tj3JPEGBufSize(C.int(width), C.int(height), C.TJSAMP_420)
	buf := make([]byte, int(size))

	srcPtr := unsafe.Pointer(&pix[0])
	dstPtr := (*C.uchar)(&buf[0])

	var pin runtime.Pinner
	pin.Pin(srcPtr)
	pin.Pin(dstPtr)
	defer pin.Unpin()

	jpegSize := size
	rc := C.tj3Compress8(e.ctx, (*C.uchar)(srcPtr), C.int(width), 0, C.int(height), C.TJPF_BGRA, &dstPtr, &jpegSize)
	if rc != 0 {
		return nil, errors.New("turbojpeg BGRA compress failed")
	}
	return buf[:int(jpegSize)], nil
}

// encodeRGBA compresses a Go RGBA image (pixel layout [R,G,B,A] per 4 bytes).
// TJPF_RGBA tells turbojpeg the exact layout; colourspace is RGB.
func (e *JPEGEncoder) encodeRGBA(img *image.RGBA) ([]byte, error) {
	w := img.Rect.Dx()
	h := img.Rect.Dy()

	subsamp := C.int(C.TJSAMP_444)
	C.tj3Set(e.ctx, C.TJPARAM_SUBSAMP, subsamp)
	C.tj3Set(e.ctx, C.TJPARAM_COLORSPACE, C.TJCS_RGB)

	size := C.tj3JPEGBufSize(C.int(w), C.int(h), subsamp)
	buf := make([]byte, int(size))
	srcPtr := unsafe.Pointer(&img.Pix[0])
	dstPtr := (*C.uchar)(&buf[0])

	var pin runtime.Pinner
	pin.Pin(srcPtr)
	pin.Pin(dstPtr)
	defer pin.Unpin()

	jpegSize := size
	rc := C.tj3Compress8(e.ctx, (*C.uchar)(srcPtr), C.int(w), 0, C.int(h), C.TJPF_RGBA, &dstPtr, &jpegSize)
	if rc != 0 {
		return nil, errors.New("turbojpeg RGBA compress failed")
	}
	return buf[:int(jpegSize)], nil
}

// encodeYCbCr compresses a YCbCr image via the turbojpeg YUV compressor.
// SubsampleRatio selects the chroma subsampling: 420, 422, or 444.
//
// turbojpeg's tj3CompressFromYUV8 expects the Y, Cb and Cr planes packed
// contiguously in a single buffer (Y, then Cb, then Cr). Go's image.YCbCr
// keeps them in three independent slices, so we copy them into a packed
// scratch buffer first — matching obs-teleport's ToJPEG behaviour.
func (e *JPEGEncoder) encodeYCbCr(img *image.YCbCr) ([]byte, error) {
	w := img.Rect.Dx()
	h := img.Rect.Dy()

	var subsamp C.int
	switch img.SubsampleRatio {
	case image.YCbCrSubsampleRatio420:
		subsamp = C.TJSAMP_420
	case image.YCbCrSubsampleRatio422:
		subsamp = C.TJSAMP_422
	case image.YCbCrSubsampleRatio444:
		subsamp = C.TJSAMP_444
	default:
		// Fall back to 444; this matches obs-teleport's default for non-standard ratios.
		subsamp = C.TJSAMP_444
	}

	C.tj3Set(e.ctx, C.TJPARAM_SUBSAMP, subsamp)
	C.tj3Set(e.ctx, C.TJPARAM_COLORSPACE, C.TJCS_YCbCr)

	size := C.tj3JPEGBufSize(C.int(w), C.int(h), subsamp)
	buf := make([]byte, int(size))

	// Pack the planes contiguously for the compressor.
	yuv := make([]byte, 0, len(img.Y)+len(img.Cb)+len(img.Cr))
	yuv = append(yuv, img.Y...)
	yuv = append(yuv, img.Cb...)
	yuv = append(yuv, img.Cr...)

	srcPtr := unsafe.Pointer(&yuv[0])
	dstPtr := (*C.uchar)(&buf[0])

	var pin runtime.Pinner
	pin.Pin(srcPtr)
	pin.Pin(dstPtr)
	defer pin.Unpin()

	jpegSize := size
	rc := C.tj3CompressFromYUV8(e.ctx, (*C.uchar)(srcPtr), C.int(w), 1, C.int(h), &dstPtr, &jpegSize)
	if rc != 0 {
		return nil, errors.New("turbojpeg YUV compress failed")
	}
	return buf[:int(jpegSize)], nil
}

// encodeGeneric is the fallback for unsupported image types: draw into
// RGBA and encode via the RGBA path.
func (e *JPEGEncoder) encodeGeneric(img image.Image) ([]byte, error) {
	b := img.Bounds()
	rgba := image.NewRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			r, g, b2, a := img.At(x, y).RGBA()
			rgba.SetRGBA(x-b.Min.X, y-b.Min.Y, color.RGBA{
				R: uint8(r >> 8),
				G: uint8(g >> 8),
				B: uint8(b2 >> 8),
				A: uint8(a >> 8),
			})
		}
	}
	return e.encodeRGBA(rgba)
}