TeleportFling/internal/flinger/flinger.go

// Package flinger implements the TeleportFling streaming engine: it captures
// a screen (and optionally system audio), encodes video to JPEG, packetizes
// both into the Teleport protocol, and broadcasts them to connected OBS
// receivers.
//
// The engine is deliberately GUI-free so it can run headless (CLI/daemon) or
// be driven by a desktop app. Callers create an Engine with a Config, call
// Start, poll Status, and call Stop when done.
package flinger

import (
	"errors"
	"fmt"
	"image"
	"image/color"
	"io"
	"log"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"teleportfling/internal/capture"
	"teleportfling/internal/discovery"
	"teleportfling/internal/output"
	"teleportfling/internal/protocol"
)

const (
	// sampleRate and speakers describe the captured/encoded audio stream.
	sampleRate = 48000
	speakers   = 2
	// audioChunk sets how much audio we packetize per WAVE message (~10 ms).
	audioChunk = 10 * time.Millisecond
)

// Config configures the streaming engine.
type Config struct {
	// Name is the announce name advertised to receivers (empty → hostname).
	Name string
	// Port is the TCP listening port for receivers.
	Port int
	// Quality is the JPEG quality, 1–100.
	Quality int
	// FPS is the target video frame rate.
	FPS int
	// Source is "screen" (PipeWire) or "pattern" (synthetic test signal).
	Source string
	// Audio enables system audio capture and streaming.
	Audio bool
	// AudioSource is the PipeWire node serial to capture audio from (0 = system
	// default output). Only applies when Audio is true and Source is "screen".
	AudioSource uint64
	// StreamIndex selects which monitor to capture (screen source only).
	StreamIndex int
	// Scale downscales the captured frame before encoding. 1.0 = native
	// resolution, 0.5 = half width/height, etc. Must be in (0, 1].
	Scale float64
	// Announce controls whether the stream is advertised via UDP multicast.
	// When disabled, receivers must connect by IP manually.
	Announce bool
}

// DefaultConfig returns the recommended defaults.
func DefaultConfig() Config {
	return Config{
		Port:     9756,
		Quality:  80,
		FPS:      30,
		Source:   "screen",
		Audio:    true,
		Scale:    1.0,
		Announce: true,
	}
}

// Validate checks the config for out-of-range or unsupported values.
func (c Config) Validate() error {
	if c.Port < 1 || c.Port > 65535 {
		return errors.New("port must be 1–65535")
	}
	if c.Quality < 1 || c.Quality > 100 {
		return errors.New("quality must be 1–100")
	}
	if c.FPS < 1 || c.FPS > 240 {
		return errors.New("fps must be 1–240")
	}
	if c.Source != "screen" && c.Source != "pattern" {
		return errors.New("source must be screen or pattern")
	}
	if c.StreamIndex < 0 {
		return errors.New("stream index must be >= 0")
	}
	if c.Scale <= 0 || c.Scale > 1 {
		return errors.New("scale must be in (0, 1]")
	}
	return nil
}

// Status is a point-in-time snapshot of the running engine.
type Status struct {
	Running bool
	Frames  int64
	Dropped int64
	Conns   int
	// Bitrate is the measured stream bandwidth in bits per second, averaged
	// over the previous measurement window.
	Bitrate int64
	// Err is the most recent runtime error encountered (capture, encode,
	// packet or audio), or nil if the stream is healthy.
	Err error
}

// Engine owns the capture, encode and send pipeline.
type Engine struct {
	cfgMu sync.RWMutex
	cfg   Config

	sender    *output.Sender
	announcer *discovery.Announcer
	encoder   *protocol.JPEGEncoder
	cam       capture.Capture
	loop      frameSource

	start time.Time
	stop  chan struct{}

	// running gates packet sending: when false the loops consume capture
	// frames but do not transmit. This lets the GUI "stop" without tearing
	// down the portal/PipeWire session (which the portal cannot reliably
	// recreate in-process), so a subsequent Start just flips running on.
	running atomic.Bool

	frames atomic.Int64

	// wg tracks the audio/video/stats goroutines so Stop can wait for them to
	// finish before destroying shared resources (encoder, sender, capture).
	wg sync.WaitGroup

	// scaleBuf is the cached downscaled frame, reallocated only when the
	// configured scale or source dimensions change.
	scaleBuf *capture.VideoFrame

	errMu   sync.RWMutex
	lastErr error

	// bitrate tracking
	bitMu    sync.Mutex
	bitLast  time.Time
	bitBytes int64
	bitrate  int64
}

// setErr records the most recent runtime error. Pass nil to clear it.
func (e *Engine) setErr(err error) {
	e.errMu.Lock()
	e.lastErr = err
	e.errMu.Unlock()
}

// SetConfig updates engine settings live (FPS, quality, name, etc.) without
// restarting the stream. It validates the new config first; on error the
// engine keeps its current settings.
//
// Not every field is live-applicable mid-stream: source, port and
// stream-index still require a restart (they are ignored if changed).
func (e *Engine) SetConfig(cfg Config) error {
	if err := cfg.Validate(); err != nil {
		return err
	}
	e.cfgMu.Lock()
	defer e.cfgMu.Unlock()

	// Fields that cannot change live keep their current values.
	cfg.Source = e.cfg.Source
	cfg.Port = e.cfg.Port
	cfg.StreamIndex = e.cfg.StreamIndex
	cfg.AudioSource = e.cfg.AudioSource

	e.cfg = cfg
	return nil
}

// getCfg returns a snapshot of the current config.
func (e *Engine) getCfg() Config {
	e.cfgMu.RLock()
	defer e.cfgMu.RUnlock()
	return e.cfg
}

// New creates an engine from cfg. Capture is opened eagerly so that
// misconfiguration (e.g. no screen-share permission) surfaces before Start.
func New(cfg Config) (*Engine, error) {
	if err := cfg.Validate(); err != nil {
		return nil, err
	}

	e := &Engine{cfg: cfg}

	sender := output.New()
	if _, err := sender.Listen(addr(cfg.Port)); err != nil {
		return nil, err
	}

	enc, err := protocol.NewJPEGEncoder()
	if err != nil {
		sender.Close()
		return nil, err
	}

	e.sender = sender
	e.encoder = enc

	switch cfg.Source {
	case "screen":
		cam, err := capture.OpenPipeWire(cfg.StreamIndex, cfg.Audio, cfg.AudioSource)
		if err != nil {
			e.encoder.Close()
			sender.Close()
			return nil, err
		}
		e.cam = cam
		e.loop = captureLoop{cam}
	case "pattern":
		e.loop = &patternLoop{w: 1920, h: 1080}
	default:
		e.encoder.Close()
		sender.Close()
		return nil, errors.New("unknown source " + cfg.Source)
	}

	return e, nil
}

// Start begins the audio, video and stats loops and (unless disabled in the
// config) starts announcing the stream. It is idempotent.
func (e *Engine) Start() {
	if e.stop != nil {
		e.Resume()
		return
	}

	e.start = time.Now()
	e.stop = make(chan struct{})
	e.running.Store(true)

	if e.cfg.Announce {
		e.announcer = discovery.Start(e.cfg.Name, e.sender.Port())
	}

	var src io.ReadCloser
	switch {
	case e.cam == nil:
		src = capture.NewSilenceSource()
	case e.cam.Audio() != nil:
		src = e.cam.Audio()
	default:
		log.Printf("flinger: system audio unavailable, streaming silence")
		src = capture.NewSilenceSource()
	}

	e.wg.Add(3)
	go func() {
		defer e.wg.Done()
		e.audioLoop(src)
	}()
	go func() {
		defer e.wg.Done()
		e.videoLoop()
	}()
	go func() {
		defer e.wg.Done()
		e.statsLoop()
	}()
}

// Stop pauses the stream: loops keep consuming capture frames but stop
// transmitting, and the portal/PipeWire session stays open so the engine can
// be resumed with Start without re-opening the portal (which the portal
// cannot reliably do in-process). Idempotent.
func (e *Engine) Stop() {
	e.running.Store(false)
}

// Pause is an alias for Stop; kept for clarity at call sites.
func (e *Engine) Pause() { e.Stop() }

// Resume restarts transmission on a paused engine. It is a no-op if the
// engine was never started.
func (e *Engine) Resume() {
	if e.stop == nil {
		return
	}
	e.running.Store(true)
}

// Close fully tears down the engine: stops the loops, closes the portal
// session and frees the encoder/sender. After Close the engine must not be
// reused.
func (e *Engine) Close() {
	if e.stop == nil {
		return
	}

	close(e.stop)

	// Closing the capture source unblocks the audio and video loops that may
	// be stuck in a read. Do this before waiting so they can observe stop.
	if e.cam != nil {
		if err := e.cam.Close(); err != nil {
			log.Printf("flinger: capture close: %v", err)
		}
	}

	// Wait for the audio/video/stats goroutines to finish so they no longer
	// touch the encoder or sender before we destroy them.
	e.wg.Wait()

	if e.announcer != nil {
		e.announcer.Stop()
	}
	e.sender.Close()
	e.encoder.Close()
}

// Status returns a snapshot of the running engine.
func (e *Engine) Status() Status {
	e.errMu.RLock()
	err := e.lastErr
	e.errMu.RUnlock()

	return Status{
		Running: e.running.Load() && e.stop != nil,
		Frames:  e.frames.Load(),
		Dropped: e.sender.Dropped(),
		Conns:   e.sender.NumConns(),
		Bitrate: e.measureBitrate(),
		Err:     err,
	}
}

// bitrateWindow is the sliding window over which bitrate is averaged.
const bitrateWindow = 2 * time.Second

// measureBitrate computes the current stream bitrate (bits/sec) over a
// sliding window. It is called from Status.
func (e *Engine) measureBitrate() int64 {
	e.bitMu.Lock()
	defer e.bitMu.Unlock()

	now := time.Now()
	bytes := e.sender.BytesSent()

	if e.bitLast.IsZero() {
		e.bitLast = now
		e.bitBytes = bytes
		return 0
	}

	elapsed := now.Sub(e.bitLast)
	if elapsed < bitrateWindow {
		return e.bitrate
	}

	// Bytes accumulated since the previous sample.
	delta := bytes - e.bitBytes
	e.bitrate = int64(float64(delta*8) / elapsed.Seconds())
	e.bitLast = now
	e.bitBytes = bytes
	return e.bitrate
}

// audioLoop reads raw PCM and emits WAVE packets. start is the shared
// reference clock used by the video loop so audio and video timestamps stay
// aligned on the receiver.
//
// The PipeWire capture negotiates interleaved signed 16-bit stereo at 48 kHz
// (the teleportfling stream's negotiated Format is S16LE), which is exactly
// what the WAVE packets carry.
func (e *Engine) audioLoop(src io.ReadCloser) {
	defer func() { _ = src.Close() }()

	framesPerChunk := int(sampleRate) * int(audioChunk) / int(time.Second)
	chunkBytes := framesPerChunk * speakers * 2 // S16
	buf := make([]byte, chunkBytes)

	for {
		select {
		case <-e.stop:
			return
		default:
		}

		n, err := io.ReadFull(src, buf)
		if n > 0 {
			if !e.running.Load() {
				continue
			}
			frames := n / (speakers * 2)
			ts := uint64(time.Since(e.start))
			packet, perr := protocol.BuildWavePacket(ts, protocol.AudioFormatS16, sampleRate, speakers, int32(frames), buf[:n])
			if perr != nil {
				e.setErr(perr)
				log.Printf("flinger: wave: %v", perr)
			} else {
				e.sender.Send(packet)
			}
		}
		if err != nil {
			select {
			case <-e.stop:
				return
			default:
			}
			if !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrClosedPipe) {
				e.setErr(err)
				log.Printf("flinger: audio: %v", err)
			}
		}
	}
}

// scaleFrame downscales frame by the configured factor, returning the original
// when scale is 1.0 (native). The result is a cached buffer reused across
// frames, reallocated only when the dimensions or scale change.
func (e *Engine) scaleFrame(frame *capture.VideoFrame, scale float64) *capture.VideoFrame {
	if scale >= 1.0 || frame == nil {
		return frame
	}

	sw := int(float64(frame.Width) * scale)
	sh := int(float64(frame.Height) * scale)
	if sw < 1 {
		sw = 1
	}
	if sh < 1 {
		sh = 1
	}

	// Reuse the cached buffer if it matches the target size.
	if e.scaleBuf == nil || e.scaleBuf.Width != sw || e.scaleBuf.Height != sh {
		e.scaleBuf = &capture.VideoFrame{
			Pix:    make([]byte, sw*sh*4),
			Width:  sw,
			Height: sh,
			Stride: sw * 4,
		}
	}

	scaleBGRA(frame, e.scaleBuf)
	return e.scaleBuf
}

// scaleBGRA bilinearly downsamples an interleaved BGRA frame into dst.
func scaleBGRA(src, dst *capture.VideoFrame) {
	sw, sh := float64(src.Width), float64(src.Height)
	for y := 0; y < dst.Height; y++ {
		srcY := (float64(y) + 0.5) * sh / float64(dst.Height)
		y0 := int(srcY)
		if y0 >= src.Height-1 {
			y0 = src.Height - 2
		}
		yFrac := srcY - float64(y0)
		row0 := y0 * src.Stride
		row1 := (y0 + 1) * src.Stride
		di := y * dst.Stride

		for x := 0; x < dst.Width; x++ {
			srcX := (float64(x) + 0.5) * sw / float64(dst.Width)
			x0 := int(srcX)
			if x0 >= src.Width-1 {
				x0 = src.Width - 2
			}
			xFrac := srcX - float64(x0)
			p00 := row0 + x0*4
			p01 := row0 + (x0+1)*4
			p10 := row1 + x0*4
			p11 := row1 + (x0+1)*4

			for c := 0; c < 4; c++ {
				top := float64(src.Pix[p00+c])*(1-xFrac) + float64(src.Pix[p01+c])*xFrac
				bot := float64(src.Pix[p10+c])*(1-xFrac) + float64(src.Pix[p11+c])*xFrac
				dst.Pix[di+x*4+c] = uint8(top*(1-yFrac) + bot*yFrac)
			}
		}
	}
}

// videoLoop pulls frames and sends them at the configured fps. The frame rate
// and JPEG quality are read from the live config so SetConfig takes effect
// without restarting.
func (e *Engine) videoLoop() {
	frameInterval := time.Second / time.Duration(e.getCfg().FPS)
	next := e.start

	for {
		select {
		case <-e.stop:
			return
		default:
		}

		frame, err := e.loop.Next()
		if err != nil {
			select {
			case <-e.stop:
				return
			default:
			}
			if !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrClosedPipe) {
				e.setErr(err)
				log.Printf("flinger: capture: %v", err)
			}
			continue
		}

		// Re-read the live config each frame so FPS/quality changes apply
		// immediately. When the interval changes, resync `next` to now.
		cfg := e.getCfg()
		interval := time.Second / time.Duration(cfg.FPS)
		if interval != frameInterval {
			frameInterval = interval
			next = time.Now()
		}

		// Drop frames when running ahead of the target fps to keep
		// timestamps monotonic (e.g. a 60 Hz monitor captured at 30 fps).
		now := time.Now()
		if now.Before(next) {
			continue
		}
		next = now.Add(frameInterval)

		// When paused (Stop), keep consuming frames to stay alive but do not
		// encode or transmit. The portal session is left open so a subsequent
		// Start can resume without re-opening the portal.
		if !e.running.Load() {
			continue
		}

		ts := uint64(now.Sub(e.start))
		enc := e.scaleFrame(frame, cfg.Scale)
		buf, err := e.encoder.EncodeBGRA(enc.Pix, enc.Width, enc.Height, cfg.Quality)
		if err != nil {
			e.setErr(err)
			log.Printf("flinger: jpeg: %v", err)
			continue
		}

		packet, err := protocol.WritePacket(
			protocol.Header{Type: protocol.VideoType, Timestamp: ts, Size: int32(len(buf))},
			ptr(protocol.DefaultBT709Full()),
			nil,
			buf,
		)
		if err != nil {
			e.setErr(err)
			log.Printf("flinger: packet: %v", err)
			continue
		}
		e.sender.Send(packet)
		e.frames.Add(1)
	}
}

// statsLoop logs a periodic summary.
func (e *Engine) statsLoop() {
	tick := time.NewTicker(5 * time.Second)
	defer tick.Stop()
	for {
		select {
		case <-tick.C:
			st := e.Status()
			rate := formatBitrate(st.Bitrate)
			if st.Dropped > 0 {
				log.Printf("flinger: %d frames, %d dropped, %d conns, %s", st.Frames, st.Dropped, st.Conns, rate)
			} else {
				log.Printf("flinger: %d frames, %d conns, %s", st.Frames, st.Conns, rate)
			}
		case <-e.stop:
			return
		}
	}
}

// formatBitrate renders a bits/sec value in a human-readable form.
func formatBitrate(bps int64) string {
	switch {
	case bps >= 1_000_000:
		return fmt.Sprintf("%.1f Mbps", float64(bps)/1_000_000)
	case bps >= 1_000:
		return fmt.Sprintf("%.0f kbps", float64(bps)/1_000)
	default:
		return fmt.Sprintf("%d bps", bps)
	}
}

// frameSource abstracts the frame source: real capture or the test pattern.
type frameSource interface {
	Next() (*capture.VideoFrame, error)
}

// captureLoop wraps the PipeWire capture backend.
type captureLoop struct {
	cam capture.Capture
}

func (c captureLoop) Next() (*capture.VideoFrame, error) {
	return c.cam.Video().NextFrame()
}

// patternLoop synthesizes the M1 test pattern (colour bars + moving box).
type patternLoop struct {
	w, h int
	seq  int64
}

func (p *patternLoop) Next() (*capture.VideoFrame, error) {
	img := testPattern(p.w, p.h, int(p.seq))
	p.seq++
	return ycrcbToBGRA(img), nil
}

// ycrcbToBGRA converts a YCbCr image to a BGRA VideoFrame so both sources
// share the encode path (EncodeBGRA).
func ycrcbToBGRA(img *image.YCbCr) *capture.VideoFrame {
	w, h := img.Rect.Dx(), img.Rect.Dy()
	frame := &capture.VideoFrame{
		Pix:    make([]byte, w*h*4),
		Width:  w,
		Height: h,
		Stride: w * 4,
	}
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			yi := y*img.YStride + x
			ci := (y/2)*img.CStride + x/2
			r, g, b := color.YCbCrToRGB(img.Y[yi], img.Cb[ci], img.Cr[ci])
			off := (y*w + x) * 4
			frame.Pix[off], frame.Pix[off+1], frame.Pix[off+2], frame.Pix[off+3] = b, g, r, 255
		}
	}
	return frame
}

// ptr returns a pointer to v, for passing headers to WritePacket.
func ptr[T any](v T) *T { return &v }

// addr formats a port as a listen address.
func addr(port int) string {
	return ":" + strconv.Itoa(port)
}

// testPattern renders a standard SMPTE colour bar with a moving white box at
// the given frame index.
func testPattern(w, h, frame int) *image.YCbCr {
	img := image.NewYCbCr(image.Rect(0, 0, w, h), image.YCbCrSubsampleRatio420)

	// 7 vertical colour bars (grey, yellow, cyan, green, magenta, red, blue).
	bars := []color.RGBA{
		{R: 191, G: 191, B: 191}, // 75% grey
		{R: 191, G: 191, B: 0},   // yellow
		{R: 0, G: 191, B: 191},   // cyan
		{R: 0, G: 191, B: 0},     // green
		{R: 191, G: 0, B: 191},   // magenta
		{R: 191, G: 0, B: 0},     // red
		{R: 0, G: 0, B: 191},     // blue
	}

	const barCount = 7
	barW := w / barCount
	const boxSize = 80

	// Moving white box sweeps left→right across the lower black block.
	boxMinX := (frame*(w+boxSize)/120)%(w+boxSize) - boxSize/2

	buf := make([]color.RGBA, w*h)
	for by := 0; by < h; by++ {
		rowIsBars := by < h*2/3
		for bx := 0; bx < w; bx++ {
			var c color.RGBA
			switch {
			case rowIsBars:
				idx := bx / barW
				if idx >= barCount {
					idx = barCount - 1
				}
				c = bars[idx]
			case by%8 < 4 && bx > w/3 && bx < w*2/3:
				c = color.RGBA{R: 255, G: 255, B: 255, A: 255}
			default:
				c = color.RGBA{}
			}
			if bx >= boxMinX && bx < boxMinX+boxSize && by >= h*2/3 {
				c = color.RGBA{R: 255, G: 255, B: 255, A: 255}
			}
			buf[by*w+bx] = c
		}
	}

	// Chroma planes: average each 2x2 RGB block, then convert to Cb/Cr.
	for by := 0; by < h; by += 2 {
		for bx := 0; bx < w; bx += 2 {
			var rSum, gSum, bSum uint32
			n := uint32(0)
			for dy := 0; dy < 2; dy++ {
				for dx := 0; dx < 2; dx++ {
					xx, yy := bx+dx, by+dy
					if xx >= w || yy >= h {
						continue
					}
					px := buf[yy*w+xx]
					rSum += uint32(px.R)
					gSum += uint32(px.G)
					bSum += uint32(px.B)
					n++
				}
			}
			_, cb, cr := color.RGBToYCbCr(uint8(rSum/n), uint8(gSum/n), uint8(bSum/n))
			img.Cb[(by/2)*img.CStride+bx/2] = cb
			img.Cr[(by/2)*img.CStride+bx/2] = cr
		}
	}

	// Luma plane: Y = YCbCr luma of every pixel.
	for by := 0; by < h; by++ {
		for bx := 0; bx < w; bx++ {
			px := buf[by*w+bx]
			y, _, _ := color.RGBToYCbCr(px.R, px.G, px.B)
			img.Y[by*img.YStride+bx] = y
		}
	}

	return img
}