diff --git a/core/cmd/dms/commands_common.go b/core/cmd/dms/commands_common.go index d1dc23e0..fabc39a7 100644 --- a/core/cmd/dms/commands_common.go +++ b/core/cmd/dms/commands_common.go @@ -472,5 +472,7 @@ func getCommonCommands() []*cobra.Command { greeterCmd, setupCmd, colorCmd, + screenshotCmd, + notifyActionCmd, } } diff --git a/core/cmd/dms/commands_screenshot.go b/core/cmd/dms/commands_screenshot.go new file mode 100644 index 00000000..5ab58327 --- /dev/null +++ b/core/cmd/dms/commands_screenshot.go @@ -0,0 +1,300 @@ +package main + +import ( + "bytes" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + + "github.com/AvengeMedia/DankMaterialShell/core/internal/screenshot" + "github.com/spf13/cobra" +) + +var ( + ssOutputName string + ssIncludeCursor bool + ssFormat string + ssQuality int + ssOutputDir string + ssFilename string + ssClipboard bool + ssNoFreeze bool + ssNoNotify bool + ssStdout bool +) + +var screenshotCmd = &cobra.Command{ + Use: "screenshot", + Short: "Capture screenshots", + Long: `Capture screenshots from Wayland displays. + +Modes: + region - Select a region interactively (default) + full - Capture the focused output + all - Capture all outputs combined + output - Capture a specific output by name + last - Capture the last selected region + +Output format (--format): + png - PNG format (default) + jpg/jpeg - JPEG format + ppm - PPM format + +Examples: + dms screenshot # Interactive region selection + dms screenshot full # Full screen of focused output + dms screenshot all # All screens combined + dms screenshot output -o DP-1 # Specific output + dms screenshot last # Last region (pre-selected) + dms screenshot --clipboard # Copy to clipboard + dms screenshot --cursor # Include cursor + dms screenshot -f jpg -q 85 # JPEG with quality 85`, +} + +var ssRegionCmd = &cobra.Command{ + Use: "region", + Short: "Select a region interactively", + Run: runScreenshotRegion, +} + +var ssFullCmd = &cobra.Command{ + Use: "full", + Short: "Capture the focused output", + Run: runScreenshotFull, +} + +var ssAllCmd = &cobra.Command{ + Use: "all", + Short: "Capture all outputs combined", + Run: runScreenshotAll, +} + +var ssOutputCmd = &cobra.Command{ + Use: "output", + Short: "Capture a specific output", + Run: runScreenshotOutput, +} + +var ssLastCmd = &cobra.Command{ + Use: "last", + Short: "Capture the last selected region", + Long: `Capture the previously selected region without interactive selection. +If no previous region exists, falls back to interactive selection.`, + Run: runScreenshotLast, +} + +var ssListCmd = &cobra.Command{ + Use: "list", + Short: "List available outputs", + Run: runScreenshotList, +} + +var notifyActionCmd = &cobra.Command{ + Use: "notify-action", + Hidden: true, + Run: func(cmd *cobra.Command, args []string) { + screenshot.RunNotifyActionListener(args) + }, +} + +func init() { + screenshotCmd.PersistentFlags().StringVarP(&ssOutputName, "output", "o", "", "Output name for 'output' mode") + screenshotCmd.PersistentFlags().BoolVar(&ssIncludeCursor, "cursor", false, "Include cursor in screenshot") + screenshotCmd.PersistentFlags().StringVarP(&ssFormat, "format", "f", "png", "Output format (png, jpg, ppm)") + screenshotCmd.PersistentFlags().IntVarP(&ssQuality, "quality", "q", 90, "JPEG quality (1-100)") + screenshotCmd.PersistentFlags().StringVarP(&ssOutputDir, "dir", "d", "", "Output directory") + screenshotCmd.PersistentFlags().StringVar(&ssFilename, "filename", "", "Output filename (auto-generated if empty)") + screenshotCmd.PersistentFlags().BoolVar(&ssClipboard, "clipboard", false, "Copy to clipboard instead of file") + screenshotCmd.PersistentFlags().BoolVar(&ssNoFreeze, "no-freeze", false, "Don't freeze screen during region selection") + screenshotCmd.PersistentFlags().BoolVar(&ssNoNotify, "no-notify", false, "Don't show notification after capture") + screenshotCmd.PersistentFlags().BoolVar(&ssStdout, "stdout", false, "Output image to stdout (for piping to swappy, etc.)") + + screenshotCmd.AddCommand(ssRegionCmd) + screenshotCmd.AddCommand(ssFullCmd) + screenshotCmd.AddCommand(ssAllCmd) + screenshotCmd.AddCommand(ssOutputCmd) + screenshotCmd.AddCommand(ssLastCmd) + screenshotCmd.AddCommand(ssListCmd) + + screenshotCmd.Run = runScreenshotRegion +} + +func getScreenshotConfig(mode screenshot.Mode) screenshot.Config { + config := screenshot.DefaultConfig() + config.Mode = mode + config.OutputName = ssOutputName + config.IncludeCursor = ssIncludeCursor + config.Clipboard = ssClipboard + config.Freeze = !ssNoFreeze + config.Notify = !ssNoNotify + config.Stdout = ssStdout + + if ssOutputDir != "" { + config.OutputDir = ssOutputDir + } + if ssFilename != "" { + config.Filename = ssFilename + } + + switch strings.ToLower(ssFormat) { + case "jpg", "jpeg": + config.Format = screenshot.FormatJPEG + case "ppm": + config.Format = screenshot.FormatPPM + default: + config.Format = screenshot.FormatPNG + } + + if ssQuality < 1 { + ssQuality = 1 + } + if ssQuality > 100 { + ssQuality = 100 + } + config.Quality = ssQuality + + return config +} + +func runScreenshot(config screenshot.Config) { + sc := screenshot.New(config) + result, err := sc.Run() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + if result == nil { + os.Exit(0) + } + + defer result.Buffer.Close() + + if result.YInverted { + result.Buffer.FlipVertical() + } + + if config.Stdout { + if err := writeImageToStdout(result.Buffer, config.Format, config.Quality); err != nil { + fmt.Fprintf(os.Stderr, "Error writing to stdout: %v\n", err) + os.Exit(1) + } + return + } + + if config.Clipboard { + if err := copyImageToClipboard(result.Buffer, config.Format, config.Quality); err != nil { + fmt.Fprintf(os.Stderr, "Error copying to clipboard: %v\n", err) + os.Exit(1) + } + if config.Notify { + screenshot.SendNotification(screenshot.NotifyResult{Clipboard: true}) + } + fmt.Println("Screenshot copied to clipboard") + return + } + + outputDir := config.OutputDir + if outputDir == "" { + outputDir = screenshot.GetOutputDir() + } + + filename := config.Filename + if filename == "" { + filename = screenshot.GenerateFilename(config.Format) + } + + path := filepath.Join(outputDir, filename) + if err := screenshot.WriteToFile(result.Buffer, path, config.Format, config.Quality); err != nil { + fmt.Fprintf(os.Stderr, "Error writing file: %v\n", err) + os.Exit(1) + } + + if config.Notify { + screenshot.SendNotification(screenshot.NotifyResult{FilePath: path}) + } + + fmt.Println(path) +} + +func copyImageToClipboard(buf *screenshot.ShmBuffer, format screenshot.Format, quality int) error { + var mimeType string + var data bytes.Buffer + + img := screenshot.BufferToImage(buf) + + switch format { + case screenshot.FormatJPEG: + mimeType = "image/jpeg" + if err := screenshot.EncodeJPEG(&data, img, quality); err != nil { + return err + } + default: + mimeType = "image/png" + if err := screenshot.EncodePNG(&data, img); err != nil { + return err + } + } + + cmd := exec.Command("wl-copy", "--type", mimeType) + cmd.Stdin = &data + return cmd.Run() +} + +func writeImageToStdout(buf *screenshot.ShmBuffer, format screenshot.Format, quality int) error { + img := screenshot.BufferToImage(buf) + + switch format { + case screenshot.FormatJPEG: + return screenshot.EncodeJPEG(os.Stdout, img, quality) + default: + return screenshot.EncodePNG(os.Stdout, img) + } +} + +func runScreenshotRegion(cmd *cobra.Command, args []string) { + config := getScreenshotConfig(screenshot.ModeRegion) + runScreenshot(config) +} + +func runScreenshotFull(cmd *cobra.Command, args []string) { + config := getScreenshotConfig(screenshot.ModeFullScreen) + runScreenshot(config) +} + +func runScreenshotAll(cmd *cobra.Command, args []string) { + config := getScreenshotConfig(screenshot.ModeAllScreens) + runScreenshot(config) +} + +func runScreenshotOutput(cmd *cobra.Command, args []string) { + if ssOutputName == "" && len(args) > 0 { + ssOutputName = args[0] + } + if ssOutputName == "" { + fmt.Fprintln(os.Stderr, "Error: output name required (use -o or provide as argument)") + os.Exit(1) + } + config := getScreenshotConfig(screenshot.ModeOutput) + runScreenshot(config) +} + +func runScreenshotLast(cmd *cobra.Command, args []string) { + config := getScreenshotConfig(screenshot.ModeLastRegion) + runScreenshot(config) +} + +func runScreenshotList(cmd *cobra.Command, args []string) { + outputs, err := screenshot.ListOutputs() + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + + for _, o := range outputs { + fmt.Printf("%s: %dx%d+%d+%d (scale: %d)\n", + o.Name, o.Width, o.Height, o.X, o.Y, o.Scale) + } +} diff --git a/core/internal/colorpicker/picker.go b/core/internal/colorpicker/picker.go index 7138f95b..0e9e98ca 100644 --- a/core/internal/colorpicker/picker.go +++ b/core/internal/colorpicker/picker.go @@ -10,6 +10,7 @@ import ( "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wlr_layer_shell" "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wlr_screencopy" "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wp_viewporter" + wlhelpers "github.com/AvengeMedia/DankMaterialShell/core/internal/wayland/client" "github.com/AvengeMedia/DankMaterialShell/core/pkg/go-wayland/wayland/client" ) @@ -165,26 +166,7 @@ func (p *Picker) connect() error { } func (p *Picker) roundtrip() error { - callback, err := p.display.Sync() - if err != nil { - return err - } - - done := make(chan struct{}) - callback.SetDoneHandler(func(e client.CallbackDoneEvent) { - close(done) - }) - - for { - select { - case <-done: - return nil - default: - if err := p.ctx.Dispatch(); err != nil { - return err - } - } - } + return wlhelpers.Roundtrip(p.display, p.ctx) } func (p *Picker) setupRegistry() error { diff --git a/core/internal/colorpicker/shm.go b/core/internal/colorpicker/shm.go index 2c94b8dc..6939fcc3 100644 --- a/core/internal/colorpicker/shm.go +++ b/core/internal/colorpicker/shm.go @@ -1,93 +1,28 @@ package colorpicker -import ( - "fmt" +import "github.com/AvengeMedia/DankMaterialShell/core/internal/wayland/shm" - "golang.org/x/sys/unix" -) - -type ShmBuffer struct { - fd int - data []byte - size int - Width int - Height int - Stride int -} +type ShmBuffer = shm.Buffer func CreateShmBuffer(width, height, stride int) (*ShmBuffer, error) { - size := stride * height - - fd, err := unix.MemfdCreate("dms-colorpicker", 0) - if err != nil { - return nil, fmt.Errorf("failed to create memfd: %w", err) - } - - if err := unix.Ftruncate(fd, int64(size)); err != nil { - unix.Close(fd) - return nil, fmt.Errorf("ftruncate failed: %w", err) - } - - data, err := unix.Mmap(fd, 0, size, unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED) - if err != nil { - unix.Close(fd) - return nil, fmt.Errorf("mmap failed: %w", err) - } - - return &ShmBuffer{ - fd: fd, - data: data, - size: size, - Width: width, - Height: height, - Stride: stride, - }, nil + return shm.CreateBuffer(width, height, stride) } -func (s *ShmBuffer) Fd() int { - return s.fd -} - -func (s *ShmBuffer) Size() int { - return s.size -} - -func (s *ShmBuffer) Data() []byte { - return s.data -} - -func (s *ShmBuffer) GetPixel(x, y int) Color { - if x < 0 || x >= s.Width || y < 0 || y >= s.Height { +func GetPixelColor(buf *ShmBuffer, x, y int) Color { + if x < 0 || x >= buf.Width || y < 0 || y >= buf.Height { return Color{} } - offset := y*s.Stride + x*4 - - if offset+3 >= len(s.data) { + data := buf.Data() + offset := y*buf.Stride + x*4 + if offset+3 >= len(data) { return Color{} } return Color{ - B: s.data[offset], - G: s.data[offset+1], - R: s.data[offset+2], - A: s.data[offset+3], + B: data[offset], + G: data[offset+1], + R: data[offset+2], + A: data[offset+3], } } - -func (s *ShmBuffer) Close() error { - var firstErr error - if s.data != nil { - if err := unix.Munmap(s.data); err != nil && firstErr == nil { - firstErr = fmt.Errorf("munmap failed: %w", err) - } - s.data = nil - } - if s.fd >= 0 { - if err := unix.Close(s.fd); err != nil && firstErr == nil { - firstErr = fmt.Errorf("close fd failed: %w", err) - } - s.fd = -1 - } - return firstErr -} diff --git a/core/internal/colorpicker/state.go b/core/internal/colorpicker/state.go index 90d4750c..b2fd6230 100644 --- a/core/internal/colorpicker/state.go +++ b/core/internal/colorpicker/state.go @@ -4,15 +4,17 @@ import ( "math" "strings" "sync" + + "github.com/AvengeMedia/DankMaterialShell/core/internal/wayland/shm" ) -type PixelFormat uint32 +type PixelFormat = shm.PixelFormat const ( - FormatARGB8888 PixelFormat = 0 - FormatXRGB8888 PixelFormat = 1 - FormatABGR8888 PixelFormat = 0x34324241 - FormatXBGR8888 PixelFormat = 0x34324258 + FormatARGB8888 = shm.FormatARGB8888 + FormatXRGB8888 = shm.FormatXRGB8888 + FormatABGR8888 = shm.FormatABGR8888 + FormatXBGR8888 = shm.FormatXBGR8888 ) type SurfaceState struct { @@ -253,7 +255,7 @@ func (s *SurfaceState) Redraw() *ShmBuffer { return nil } - copy(dst.data, s.screenBuf.data) + dst.CopyFrom(s.screenBuf) px := int(math.Round(float64(s.pointerX) * s.scaleX)) py := int(math.Round(float64(s.pointerY) * s.scaleY)) @@ -261,15 +263,15 @@ func (s *SurfaceState) Redraw() *ShmBuffer { px = clamp(px, 0, dst.Width-1) py = clamp(py, 0, dst.Height-1) - picked := s.screenBuf.GetPixel(px, py) + picked := GetPixelColor(s.screenBuf, px, py) drawMagnifier( - dst.data, dst.Stride, dst.Width, dst.Height, - s.screenBuf.data, s.screenBuf.Stride, s.screenBuf.Width, s.screenBuf.Height, + dst.Data(), dst.Stride, dst.Width, dst.Height, + s.screenBuf.Data(), s.screenBuf.Stride, s.screenBuf.Width, s.screenBuf.Height, px, py, picked, ) - drawColorPreview(dst.data, dst.Stride, dst.Width, dst.Height, px, py, picked, s.displayFormat, s.lowercase) + drawColorPreview(dst.Data(), dst.Stride, dst.Width, dst.Height, px, py, picked, s.displayFormat, s.lowercase) return dst } @@ -289,7 +291,7 @@ func (s *SurfaceState) RedrawScreenOnly() *ShmBuffer { return nil } - copy(dst.data, s.screenBuf.data) + dst.CopyFrom(s.screenBuf) return dst } @@ -311,7 +313,7 @@ func (s *SurfaceState) PickColor() (Color, bool) { sy = s.screenBuf.Height - 1 - sy } - return s.screenBuf.GetPixel(sx, sy), true + return GetPixelColor(s.screenBuf, sx, sy), true } func (s *SurfaceState) Destroy() { diff --git a/core/internal/screenshot/encode.go b/core/internal/screenshot/encode.go new file mode 100644 index 00000000..b0dc30a1 --- /dev/null +++ b/core/internal/screenshot/encode.go @@ -0,0 +1,180 @@ +package screenshot + +import ( + "bufio" + "fmt" + "image" + "image/jpeg" + "image/png" + "io" + "os" + "path/filepath" + "time" +) + +func BufferToImage(buf *ShmBuffer) *image.RGBA { + img := image.NewRGBA(image.Rect(0, 0, buf.Width, buf.Height)) + data := buf.Data() + for y := 0; y < buf.Height; y++ { + srcOff := y * buf.Stride + dstOff := y * img.Stride + for x := 0; x < buf.Width; x++ { + si := srcOff + x*4 + di := dstOff + x*4 + if si+3 >= len(data) || di+3 >= len(img.Pix) { + continue + } + img.Pix[di+0] = data[si+2] // R + img.Pix[di+1] = data[si+1] // G + img.Pix[di+2] = data[si+0] // B + img.Pix[di+3] = 255 // A + } + } + return img +} + +func EncodePNG(w io.Writer, img image.Image) error { + enc := png.Encoder{CompressionLevel: png.BestSpeed} + return enc.Encode(w, img) +} + +func EncodeJPEG(w io.Writer, img image.Image, quality int) error { + return jpeg.Encode(w, img, &jpeg.Options{Quality: quality}) +} + +func EncodePPM(w io.Writer, img *image.RGBA) error { + bw := bufio.NewWriter(w) + bounds := img.Bounds() + if _, err := fmt.Fprintf(bw, "P6\n%d %d\n255\n", bounds.Dx(), bounds.Dy()); err != nil { + return err + } + for y := bounds.Min.Y; y < bounds.Max.Y; y++ { + for x := bounds.Min.X; x < bounds.Max.X; x++ { + off := (y-bounds.Min.Y)*img.Stride + (x-bounds.Min.X)*4 + if err := bw.WriteByte(img.Pix[off+0]); err != nil { + return err + } + if err := bw.WriteByte(img.Pix[off+1]); err != nil { + return err + } + if err := bw.WriteByte(img.Pix[off+2]); err != nil { + return err + } + } + } + return bw.Flush() +} + +func GenerateFilename(format Format) string { + t := time.Now() + ext := "png" + switch format { + case FormatJPEG: + ext = "jpg" + case FormatPPM: + ext = "ppm" + } + return fmt.Sprintf("screenshot_%s.%s", t.Format("2006-01-02_15-04-05"), ext) +} + +func GetOutputDir() string { + if dir := os.Getenv("DMS_SCREENSHOT_DIR"); dir != "" { + return dir + } + + if xdgPics := getXDGPicturesDir(); xdgPics != "" { + screenshotDir := filepath.Join(xdgPics, "Screenshots") + if err := os.MkdirAll(screenshotDir, 0755); err == nil { + return screenshotDir + } + return xdgPics + } + + if home := os.Getenv("HOME"); home != "" { + return home + } + return "." +} + +func getXDGPicturesDir() string { + configDir := os.Getenv("XDG_CONFIG_HOME") + if configDir == "" { + home := os.Getenv("HOME") + if home == "" { + return "" + } + configDir = filepath.Join(home, ".config") + } + + userDirsFile := filepath.Join(configDir, "user-dirs.dirs") + data, err := os.ReadFile(userDirsFile) + if err != nil { + return "" + } + + for _, line := range splitLines(string(data)) { + if len(line) == 0 || line[0] == '#' { + continue + } + const prefix = "XDG_PICTURES_DIR=" + if len(line) > len(prefix) && line[:len(prefix)] == prefix { + path := line[len(prefix):] + path = trimQuotes(path) + path = expandHome(path) + return path + } + } + return "" +} + +func splitLines(s string) []string { + var lines []string + start := 0 + for i := 0; i < len(s); i++ { + if s[i] == '\n' { + lines = append(lines, s[start:i]) + start = i + 1 + } + } + if start < len(s) { + lines = append(lines, s[start:]) + } + return lines +} + +func trimQuotes(s string) string { + if len(s) >= 2 && s[0] == '"' && s[len(s)-1] == '"' { + return s[1 : len(s)-1] + } + return s +} + +func expandHome(path string) string { + if len(path) >= 5 && path[:5] == "$HOME" { + home := os.Getenv("HOME") + return home + path[5:] + } + if len(path) >= 1 && path[0] == '~' { + home := os.Getenv("HOME") + return home + path[1:] + } + return path +} + +func WriteToFile(buf *ShmBuffer, path string, format Format, quality int) error { + f, err := os.Create(path) + if err != nil { + return err + } + defer f.Close() + + img := BufferToImage(buf) + switch format { + case FormatJPEG: + return EncodeJPEG(f, img, quality) + case FormatPPM: + return EncodePPM(f, img) + default: + return EncodePNG(f, img) + } +} diff --git a/core/internal/screenshot/notify.go b/core/internal/screenshot/notify.go new file mode 100644 index 00000000..40891b97 --- /dev/null +++ b/core/internal/screenshot/notify.go @@ -0,0 +1,156 @@ +package screenshot + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + "strconv" + "syscall" + + "github.com/AvengeMedia/DankMaterialShell/core/internal/log" + "github.com/godbus/dbus/v5" +) + +const ( + notifyDest = "org.freedesktop.Notifications" + notifyPath = "/org/freedesktop/Notifications" + notifyInterface = "org.freedesktop.Notifications" +) + +type NotifyResult struct { + FilePath string + Clipboard bool +} + +func SendNotification(result NotifyResult) { + conn, err := dbus.SessionBus() + if err != nil { + log.Debug("dbus session failed", "err", err) + return + } + + var actions []string + if !result.Clipboard && result.FilePath != "" { + actions = []string{"default", "Open"} + } + + hints := map[string]dbus.Variant{} + if result.FilePath != "" && !result.Clipboard { + hints["image-path"] = dbus.MakeVariant(result.FilePath) + } + + summary := "Screenshot captured" + body := "" + if result.Clipboard { + body = "Copied to clipboard" + } else { + body = filepath.Base(result.FilePath) + } + + obj := conn.Object(notifyDest, notifyPath) + call := obj.Call( + notifyInterface+".Notify", + 0, + "DMS", + uint32(0), + "", + summary, + body, + actions, + hints, + int32(5000), + ) + + if call.Err != nil { + log.Debug("notify call failed", "err", call.Err) + return + } + + var notificationID uint32 + if err := call.Store(¬ificationID); err != nil { + log.Debug("failed to get notification id", "err", err) + return + } + + if len(actions) == 0 || result.FilePath == "" { + return + } + + spawnActionListener(notificationID, result.FilePath) +} + +func spawnActionListener(notificationID uint32, filePath string) { + exe, err := os.Executable() + if err != nil { + log.Debug("failed to get executable", "err", err) + return + } + + cmd := exec.Command(exe, "notify-action", fmt.Sprintf("%d", notificationID), filePath) + cmd.SysProcAttr = &syscall.SysProcAttr{ + Setsid: true, + } + cmd.Start() +} + +func RunNotifyActionListener(args []string) { + if len(args) < 2 { + return + } + + notificationID, err := strconv.ParseUint(args[0], 10, 32) + if err != nil { + return + } + + filePath := args[1] + + conn, err := dbus.SessionBus() + if err != nil { + return + } + + if err := conn.AddMatchSignal( + dbus.WithMatchObjectPath(notifyPath), + dbus.WithMatchInterface(notifyInterface), + ); err != nil { + return + } + + signals := make(chan *dbus.Signal, 10) + conn.Signal(signals) + + for sig := range signals { + switch sig.Name { + case notifyInterface + ".ActionInvoked": + if len(sig.Body) < 2 { + continue + } + id, ok := sig.Body[0].(uint32) + if !ok || id != uint32(notificationID) { + continue + } + openFile(filePath) + return + + case notifyInterface + ".NotificationClosed": + if len(sig.Body) < 1 { + continue + } + id, ok := sig.Body[0].(uint32) + if !ok || id != uint32(notificationID) { + continue + } + return + } + } +} + +func openFile(filePath string) { + cmd := exec.Command("xdg-open", filePath) + cmd.SysProcAttr = &syscall.SysProcAttr{ + Setsid: true, + } + cmd.Start() +} diff --git a/core/internal/screenshot/region.go b/core/internal/screenshot/region.go new file mode 100644 index 00000000..176e9d9d --- /dev/null +++ b/core/internal/screenshot/region.go @@ -0,0 +1,732 @@ +package screenshot + +import ( + "fmt" + "sync" + + "github.com/AvengeMedia/DankMaterialShell/core/internal/log" + "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/keyboard_shortcuts_inhibit" + "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wlr_layer_shell" + "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wlr_screencopy" + "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wp_viewporter" + wlhelpers "github.com/AvengeMedia/DankMaterialShell/core/internal/wayland/client" + "github.com/AvengeMedia/DankMaterialShell/core/pkg/go-wayland/wayland/client" +) + +type SelectionState struct { + hasSelection bool // There's a selection to display (pre-loaded or user-drawn) + dragging bool // User is actively drawing a new selection + // Surface-local logical coordinates (from pointer events) + anchorX float64 + anchorY float64 + currentX float64 + currentY float64 +} + +type RenderSlot struct { + shm *ShmBuffer + pool *client.ShmPool + wlBuf *client.Buffer + busy bool +} + +type OutputSurface struct { + output *WaylandOutput + wlSurface *client.Surface + layerSurf *wlr_layer_shell.ZwlrLayerSurfaceV1 + viewport *wp_viewporter.WpViewport + screenBuf *ShmBuffer + screenBufNoCursor *ShmBuffer + screenFormat uint32 + logicalW int + logicalH int + configured bool + yInverted bool + + // Triple-buffered render slots + slots [3]*RenderSlot + slotsReady bool +} + +type RegionSelector struct { + screenshoter *Screenshoter + + display *client.Display + registry *client.Registry + ctx *client.Context + + compositor *client.Compositor + shm *client.Shm + seat *client.Seat + pointer *client.Pointer + keyboard *client.Keyboard + layerShell *wlr_layer_shell.ZwlrLayerShellV1 + screencopy *wlr_screencopy.ZwlrScreencopyManagerV1 + viewporter *wp_viewporter.WpViewporter + + shortcutsInhibitMgr *keyboard_shortcuts_inhibit.ZwpKeyboardShortcutsInhibitManagerV1 + shortcutsInhibitor *keyboard_shortcuts_inhibit.ZwpKeyboardShortcutsInhibitorV1 + + outputs map[uint32]*WaylandOutput + outputsMu sync.Mutex + + surfaces []*OutputSurface + activeSurface *OutputSurface + + // Cursor surface for crosshair + cursorSurface *client.Surface + cursorBuffer *ShmBuffer + cursorWlBuf *client.Buffer + cursorPool *client.ShmPool + + selection SelectionState + pointerX float64 + pointerY float64 + preSelect Region + showCapturedCursor bool + + running bool + cancelled bool + result Region +} + +func NewRegionSelector(s *Screenshoter) *RegionSelector { + return &RegionSelector{ + screenshoter: s, + outputs: make(map[uint32]*WaylandOutput), + showCapturedCursor: true, + } +} + +func (r *RegionSelector) Run() (Region, bool, error) { + r.preSelect = GetLastRegion() + + if err := r.connect(); err != nil { + return Region{}, false, fmt.Errorf("wayland connect: %w", err) + } + defer r.cleanup() + + if err := r.setupRegistry(); err != nil { + return Region{}, false, fmt.Errorf("registry setup: %w", err) + } + + if err := r.roundtrip(); err != nil { + return Region{}, false, fmt.Errorf("roundtrip after registry: %w", err) + } + + switch { + case r.screencopy == nil: + return Region{}, false, fmt.Errorf("compositor does not support wlr-screencopy-unstable-v1") + case r.layerShell == nil: + return Region{}, false, fmt.Errorf("compositor does not support wlr-layer-shell-unstable-v1") + case r.seat == nil: + return Region{}, false, fmt.Errorf("no seat available") + case r.compositor == nil: + return Region{}, false, fmt.Errorf("compositor not available") + case r.shm == nil: + return Region{}, false, fmt.Errorf("wl_shm not available") + case len(r.outputs) == 0: + return Region{}, false, fmt.Errorf("no outputs available") + } + + if err := r.roundtrip(); err != nil { + return Region{}, false, fmt.Errorf("roundtrip after protocol check: %w", err) + } + + if err := r.createSurfaces(); err != nil { + return Region{}, false, fmt.Errorf("create surfaces: %w", err) + } + + _ = r.createCursor() + + if err := r.roundtrip(); err != nil { + return Region{}, false, fmt.Errorf("roundtrip after surfaces: %w", err) + } + + r.running = true + for r.running { + if err := r.ctx.Dispatch(); err != nil { + return Region{}, false, fmt.Errorf("dispatch: %w", err) + } + } + + return r.result, r.cancelled, nil +} + +func (r *RegionSelector) connect() error { + display, err := client.Connect("") + if err != nil { + return err + } + r.display = display + r.ctx = display.Context() + return nil +} + +func (r *RegionSelector) roundtrip() error { + return wlhelpers.Roundtrip(r.display, r.ctx) +} + +func (r *RegionSelector) setupRegistry() error { + registry, err := r.display.GetRegistry() + if err != nil { + return err + } + r.registry = registry + + registry.SetGlobalHandler(func(e client.RegistryGlobalEvent) { + r.handleGlobal(e) + }) + + registry.SetGlobalRemoveHandler(func(e client.RegistryGlobalRemoveEvent) { + r.outputsMu.Lock() + delete(r.outputs, e.Name) + r.outputsMu.Unlock() + }) + + return nil +} + +func (r *RegionSelector) handleGlobal(e client.RegistryGlobalEvent) { + switch e.Interface { + case client.CompositorInterfaceName: + comp := client.NewCompositor(r.ctx) + if err := r.registry.Bind(e.Name, e.Interface, e.Version, comp); err == nil { + r.compositor = comp + } + + case client.ShmInterfaceName: + shm := client.NewShm(r.ctx) + if err := r.registry.Bind(e.Name, e.Interface, e.Version, shm); err == nil { + r.shm = shm + } + + case client.SeatInterfaceName: + seat := client.NewSeat(r.ctx) + if err := r.registry.Bind(e.Name, e.Interface, e.Version, seat); err == nil { + r.seat = seat + r.setupInput() + } + + case client.OutputInterfaceName: + output := client.NewOutput(r.ctx) + version := e.Version + if version > 4 { + version = 4 + } + if err := r.registry.Bind(e.Name, e.Interface, version, output); err == nil { + r.outputsMu.Lock() + r.outputs[e.Name] = &WaylandOutput{ + wlOutput: output, + globalName: e.Name, + scale: 1, + } + r.outputsMu.Unlock() + r.setupOutputHandlers(e.Name, output) + } + + case wlr_layer_shell.ZwlrLayerShellV1InterfaceName: + ls := wlr_layer_shell.NewZwlrLayerShellV1(r.ctx) + version := e.Version + if version > 4 { + version = 4 + } + if err := r.registry.Bind(e.Name, e.Interface, version, ls); err == nil { + r.layerShell = ls + } + + case wlr_screencopy.ZwlrScreencopyManagerV1InterfaceName: + sc := wlr_screencopy.NewZwlrScreencopyManagerV1(r.ctx) + version := e.Version + if version > 3 { + version = 3 + } + if err := r.registry.Bind(e.Name, e.Interface, version, sc); err == nil { + r.screencopy = sc + } + + case wp_viewporter.WpViewporterInterfaceName: + vp := wp_viewporter.NewWpViewporter(r.ctx) + if err := r.registry.Bind(e.Name, e.Interface, e.Version, vp); err == nil { + r.viewporter = vp + } + + case keyboard_shortcuts_inhibit.ZwpKeyboardShortcutsInhibitManagerV1InterfaceName: + mgr := keyboard_shortcuts_inhibit.NewZwpKeyboardShortcutsInhibitManagerV1(r.ctx) + if err := r.registry.Bind(e.Name, e.Interface, e.Version, mgr); err == nil { + r.shortcutsInhibitMgr = mgr + } + } +} + +func (r *RegionSelector) setupOutputHandlers(name uint32, output *client.Output) { + output.SetGeometryHandler(func(e client.OutputGeometryEvent) { + r.outputsMu.Lock() + if o, ok := r.outputs[name]; ok { + o.x = e.X + o.y = e.Y + o.transform = int32(e.Transform) + } + r.outputsMu.Unlock() + }) + + output.SetModeHandler(func(e client.OutputModeEvent) { + if e.Flags&uint32(client.OutputModeCurrent) == 0 { + return + } + r.outputsMu.Lock() + if o, ok := r.outputs[name]; ok { + o.width = e.Width + o.height = e.Height + } + r.outputsMu.Unlock() + }) + + output.SetScaleHandler(func(e client.OutputScaleEvent) { + r.outputsMu.Lock() + if o, ok := r.outputs[name]; ok { + o.scale = e.Factor + } + r.outputsMu.Unlock() + }) + + output.SetNameHandler(func(e client.OutputNameEvent) { + r.outputsMu.Lock() + if o, ok := r.outputs[name]; ok { + o.name = e.Name + } + r.outputsMu.Unlock() + }) +} + +func (r *RegionSelector) createSurfaces() error { + r.outputsMu.Lock() + outputs := make([]*WaylandOutput, 0, len(r.outputs)) + for _, o := range r.outputs { + outputs = append(outputs, o) + } + r.outputsMu.Unlock() + + for _, output := range outputs { + os, err := r.createOutputSurface(output) + if err != nil { + return fmt.Errorf("output %s: %w", output.name, err) + } + r.surfaces = append(r.surfaces, os) + } + + return nil +} + +func (r *RegionSelector) createCursor() error { + const size = 24 + const hotspot = size / 2 + + surface, err := r.compositor.CreateSurface() + if err != nil { + return fmt.Errorf("create cursor surface: %w", err) + } + r.cursorSurface = surface + + buf, err := CreateShmBuffer(size, size, size*4) + if err != nil { + return fmt.Errorf("create cursor buffer: %w", err) + } + r.cursorBuffer = buf + + // Draw crosshair + data := buf.Data() + for y := 0; y < size; y++ { + for x := 0; x < size; x++ { + off := (y*size + x) * 4 + // Vertical line + if x >= hotspot-1 && x <= hotspot && y >= 2 && y < size-2 { + data[off+0] = 255 // B + data[off+1] = 255 // G + data[off+2] = 255 // R + data[off+3] = 255 // A + continue + } + // Horizontal line + if y >= hotspot-1 && y <= hotspot && x >= 2 && x < size-2 { + data[off+0] = 255 + data[off+1] = 255 + data[off+2] = 255 + data[off+3] = 255 + continue + } + // Transparent + data[off+0] = 0 + data[off+1] = 0 + data[off+2] = 0 + data[off+3] = 0 + } + } + + pool, err := r.shm.CreatePool(buf.Fd(), int32(buf.Size())) + if err != nil { + return fmt.Errorf("create cursor pool: %w", err) + } + r.cursorPool = pool + + wlBuf, err := pool.CreateBuffer(0, size, size, size*4, uint32(FormatARGB8888)) + if err != nil { + return fmt.Errorf("create cursor wl_buffer: %w", err) + } + r.cursorWlBuf = wlBuf + + if err := surface.Attach(wlBuf, 0, 0); err != nil { + return fmt.Errorf("attach cursor: %w", err) + } + if err := surface.Damage(0, 0, size, size); err != nil { + return fmt.Errorf("damage cursor: %w", err) + } + if err := surface.Commit(); err != nil { + return fmt.Errorf("commit cursor: %w", err) + } + + return nil +} + +func (r *RegionSelector) createOutputSurface(output *WaylandOutput) (*OutputSurface, error) { + surface, err := r.compositor.CreateSurface() + if err != nil { + return nil, fmt.Errorf("create surface: %w", err) + } + + layerSurf, err := r.layerShell.GetLayerSurface( + surface, + output.wlOutput, + uint32(wlr_layer_shell.ZwlrLayerShellV1LayerOverlay), + "dms-screenshot", + ) + if err != nil { + return nil, fmt.Errorf("get layer surface: %w", err) + } + + os := &OutputSurface{ + output: output, + wlSurface: surface, + layerSurf: layerSurf, + } + + if r.viewporter != nil { + vp, err := r.viewporter.GetViewport(surface) + if err == nil { + os.viewport = vp + } + } + + if err := layerSurf.SetAnchor( + uint32(wlr_layer_shell.ZwlrLayerSurfaceV1AnchorTop) | + uint32(wlr_layer_shell.ZwlrLayerSurfaceV1AnchorBottom) | + uint32(wlr_layer_shell.ZwlrLayerSurfaceV1AnchorLeft) | + uint32(wlr_layer_shell.ZwlrLayerSurfaceV1AnchorRight), + ); err != nil { + return nil, fmt.Errorf("set anchor: %w", err) + } + if err := layerSurf.SetExclusiveZone(-1); err != nil { + return nil, fmt.Errorf("set exclusive zone: %w", err) + } + if err := layerSurf.SetKeyboardInteractivity(uint32(wlr_layer_shell.ZwlrLayerSurfaceV1KeyboardInteractivityExclusive)); err != nil { + return nil, fmt.Errorf("set keyboard interactivity: %w", err) + } + + layerSurf.SetConfigureHandler(func(e wlr_layer_shell.ZwlrLayerSurfaceV1ConfigureEvent) { + if err := layerSurf.AckConfigure(e.Serial); err != nil { + log.Error("ack configure failed", "err", err) + return + } + os.logicalW = int(e.Width) + os.logicalH = int(e.Height) + os.configured = true + r.captureForSurface(os) + r.ensureShortcutsInhibitor(os) + }) + + layerSurf.SetClosedHandler(func(e wlr_layer_shell.ZwlrLayerSurfaceV1ClosedEvent) { + r.running = false + r.cancelled = true + }) + + if err := surface.Commit(); err != nil { + return nil, fmt.Errorf("surface commit: %w", err) + } + + return os, nil +} + +func (r *RegionSelector) ensureShortcutsInhibitor(os *OutputSurface) { + if r.shortcutsInhibitMgr == nil || r.seat == nil || r.shortcutsInhibitor != nil { + return + } + inhibitor, err := r.shortcutsInhibitMgr.InhibitShortcuts(os.wlSurface, r.seat) + if err == nil { + r.shortcutsInhibitor = inhibitor + } +} + +func (r *RegionSelector) captureForSurface(os *OutputSurface) { + r.captureWithCursor(os, true, func() { + r.captureWithCursor(os, false, func() { + r.initRenderBuffer(os) + r.applyPreSelection(os) + r.redrawSurface(os) + }) + }) +} + +func (r *RegionSelector) captureWithCursor(os *OutputSurface, withCursor bool, onReady func()) { + cursor := int32(0) + if withCursor { + cursor = 1 + } + + frame, err := r.screencopy.CaptureOutput(cursor, os.output.wlOutput) + if err != nil { + log.Error("screencopy capture failed", "err", err) + return + } + + frame.SetBufferHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1BufferEvent) { + buf, err := CreateShmBuffer(int(e.Width), int(e.Height), int(e.Stride)) + if err != nil { + log.Error("create screen buffer failed", "err", err) + return + } + + if withCursor { + os.screenBuf = buf + os.screenFormat = e.Format + } else { + os.screenBufNoCursor = buf + } + + pool, err := r.shm.CreatePool(buf.Fd(), int32(buf.Size())) + if err != nil { + log.Error("create shm pool failed", "err", err) + return + } + + wlBuf, err := pool.CreateBuffer(0, int32(buf.Width), int32(buf.Height), int32(buf.Stride), e.Format) + if err != nil { + log.Error("create wl_buffer failed", "err", err) + pool.Destroy() + return + } + + if err := frame.Copy(wlBuf); err != nil { + log.Error("frame copy failed", "err", err) + } + pool.Destroy() + }) + + frame.SetFlagsHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1FlagsEvent) { + if withCursor { + os.yInverted = (e.Flags & 1) != 0 + } + }) + + frame.SetReadyHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1ReadyEvent) { + frame.Destroy() + if onReady != nil { + onReady() + } + }) + + frame.SetFailedHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1FailedEvent) { + log.Error("screencopy failed") + frame.Destroy() + }) +} + +func (r *RegionSelector) initRenderBuffer(os *OutputSurface) { + if os.screenBuf == nil { + return + } + + for i := 0; i < 3; i++ { + slot := &RenderSlot{} + + buf, err := CreateShmBuffer(os.screenBuf.Width, os.screenBuf.Height, os.screenBuf.Stride) + if err != nil { + log.Error("create render slot buffer failed", "err", err) + return + } + slot.shm = buf + + pool, err := r.shm.CreatePool(buf.Fd(), int32(buf.Size())) + if err != nil { + log.Error("create render slot pool failed", "err", err) + buf.Close() + return + } + slot.pool = pool + + wlBuf, err := pool.CreateBuffer(0, int32(buf.Width), int32(buf.Height), int32(buf.Stride), os.screenFormat) + if err != nil { + log.Error("create render slot wl_buffer failed", "err", err) + pool.Destroy() + buf.Close() + return + } + slot.wlBuf = wlBuf + + slotRef := slot + wlBuf.SetReleaseHandler(func(e client.BufferReleaseEvent) { + slotRef.busy = false + }) + + os.slots[i] = slot + } + os.slotsReady = true +} + +func (os *OutputSurface) acquireFreeSlot() *RenderSlot { + for _, slot := range os.slots { + if slot != nil && !slot.busy { + return slot + } + } + return nil +} + +func (r *RegionSelector) applyPreSelection(os *OutputSurface) { + if r.preSelect.IsEmpty() || os.screenBuf == nil || r.selection.hasSelection { + return + } + + if r.preSelect.Output != "" && r.preSelect.Output != os.output.name { + return + } + + scaleX := float64(os.logicalW) / float64(os.screenBuf.Width) + scaleY := float64(os.logicalH) / float64(os.screenBuf.Height) + + x1 := float64(r.preSelect.X-os.output.x) * scaleX + y1 := float64(r.preSelect.Y-os.output.y) * scaleY + x2 := float64(r.preSelect.X-os.output.x+r.preSelect.Width) * scaleX + y2 := float64(r.preSelect.Y-os.output.y+r.preSelect.Height) * scaleY + + r.selection.hasSelection = true + r.selection.dragging = false + r.selection.anchorX = x1 + r.selection.anchorY = y1 + r.selection.currentX = x2 + r.selection.currentY = y2 + r.activeSurface = os +} + +func (r *RegionSelector) getSourceBuffer(os *OutputSurface) *ShmBuffer { + if !r.showCapturedCursor && os.screenBufNoCursor != nil { + return os.screenBufNoCursor + } + return os.screenBuf +} + +func (r *RegionSelector) redrawSurface(os *OutputSurface) { + srcBuf := r.getSourceBuffer(os) + if srcBuf == nil || !os.slotsReady { + return + } + + slot := os.acquireFreeSlot() + if slot == nil { + return + } + + slot.shm.CopyFrom(srcBuf) + + // Draw overlay (dimming + selection) into this slot + r.drawOverlay(os, slot.shm) + + // Attach and commit (viewport only needs to be set once, but it's cheap) + scale := os.output.scale + if scale <= 0 { + scale = 1 + } + + if os.viewport != nil { + srcW := float64(slot.shm.Width) / float64(scale) + srcH := float64(slot.shm.Height) / float64(scale) + _ = os.viewport.SetSource(0, 0, srcW, srcH) + _ = os.viewport.SetDestination(int32(os.logicalW), int32(os.logicalH)) + } + _ = os.wlSurface.SetBufferScale(scale) + + _ = os.wlSurface.Attach(slot.wlBuf, 0, 0) + _ = os.wlSurface.Damage(0, 0, int32(os.logicalW), int32(os.logicalH)) + _ = os.wlSurface.Commit() + + // Mark this slot as busy until compositor releases it + slot.busy = true +} + +func (r *RegionSelector) cleanup() { + if r.cursorWlBuf != nil { + r.cursorWlBuf.Destroy() + } + if r.cursorPool != nil { + r.cursorPool.Destroy() + } + if r.cursorSurface != nil { + r.cursorSurface.Destroy() + } + if r.cursorBuffer != nil { + r.cursorBuffer.Close() + } + + for _, os := range r.surfaces { + for _, slot := range os.slots { + if slot == nil { + continue + } + if slot.wlBuf != nil { + slot.wlBuf.Destroy() + } + if slot.pool != nil { + slot.pool.Destroy() + } + if slot.shm != nil { + slot.shm.Close() + } + } + if os.viewport != nil { + os.viewport.Destroy() + } + if os.layerSurf != nil { + os.layerSurf.Destroy() + } + if os.wlSurface != nil { + os.wlSurface.Destroy() + } + if os.screenBuf != nil { + os.screenBuf.Close() + } + if os.screenBufNoCursor != nil { + os.screenBufNoCursor.Close() + } + } + + if r.shortcutsInhibitor != nil { + _ = r.shortcutsInhibitor.Destroy() + } + if r.shortcutsInhibitMgr != nil { + _ = r.shortcutsInhibitMgr.Destroy() + } + if r.viewporter != nil { + r.viewporter.Destroy() + } + if r.screencopy != nil { + r.screencopy.Destroy() + } + if r.pointer != nil { + r.pointer.Release() + } + if r.keyboard != nil { + r.keyboard.Release() + } + if r.display != nil { + r.ctx.Close() + } +} diff --git a/core/internal/screenshot/region_input.go b/core/internal/screenshot/region_input.go new file mode 100644 index 00000000..ba18aa2f --- /dev/null +++ b/core/internal/screenshot/region_input.go @@ -0,0 +1,161 @@ +package screenshot + +import ( + "github.com/AvengeMedia/DankMaterialShell/core/pkg/go-wayland/wayland/client" +) + +func (r *RegionSelector) setupInput() { + if r.seat == nil { + return + } + + r.seat.SetCapabilitiesHandler(func(e client.SeatCapabilitiesEvent) { + if e.Capabilities&uint32(client.SeatCapabilityPointer) != 0 && r.pointer == nil { + if pointer, err := r.seat.GetPointer(); err == nil { + r.pointer = pointer + r.setupPointerHandlers() + } + } + if e.Capabilities&uint32(client.SeatCapabilityKeyboard) != 0 && r.keyboard == nil { + if keyboard, err := r.seat.GetKeyboard(); err == nil { + r.keyboard = keyboard + r.setupKeyboardHandlers() + } + } + }) +} + +func (r *RegionSelector) setupPointerHandlers() { + r.pointer.SetEnterHandler(func(e client.PointerEnterEvent) { + if r.cursorSurface != nil { + _ = r.pointer.SetCursor(e.Serial, r.cursorSurface, 12, 12) + } + + r.activeSurface = nil + for _, os := range r.surfaces { + if os.wlSurface.ID() == e.Surface.ID() { + r.activeSurface = os + break + } + } + + r.pointerX = e.SurfaceX + r.pointerY = e.SurfaceY + }) + + r.pointer.SetMotionHandler(func(e client.PointerMotionEvent) { + if r.activeSurface == nil { + return + } + + r.pointerX = e.SurfaceX + r.pointerY = e.SurfaceY + + if r.selection.dragging { + r.selection.currentX = e.SurfaceX + r.selection.currentY = e.SurfaceY + for _, os := range r.surfaces { + r.redrawSurface(os) + } + } + }) + + r.pointer.SetButtonHandler(func(e client.PointerButtonEvent) { + if r.activeSurface == nil { + return + } + + switch e.Button { + case 0x110: // BTN_LEFT + switch e.State { + case 1: // pressed + r.selection.hasSelection = true + r.selection.dragging = true + r.selection.anchorX = r.pointerX + r.selection.anchorY = r.pointerY + r.selection.currentX = r.pointerX + r.selection.currentY = r.pointerY + for _, os := range r.surfaces { + r.redrawSurface(os) + } + case 0: // released + r.selection.dragging = false + for _, os := range r.surfaces { + r.redrawSurface(os) + } + } + default: + r.cancelled = true + r.running = false + } + }) +} + +func (r *RegionSelector) setupKeyboardHandlers() { + r.keyboard.SetKeyHandler(func(e client.KeyboardKeyEvent) { + if e.State != 1 { + return + } + + switch e.Key { + case 1: // KEY_ESC + r.cancelled = true + r.running = false + case 25: // KEY_P + r.showCapturedCursor = !r.showCapturedCursor + for _, os := range r.surfaces { + r.redrawSurface(os) + } + case 28, 57: // KEY_ENTER, KEY_SPACE + if r.selection.hasSelection { + r.finishSelection() + } + } + }) +} + +func (r *RegionSelector) finishSelection() { + if r.activeSurface == nil { + r.running = false + return + } + + os := r.activeSurface + + x1, y1 := r.selection.anchorX, r.selection.anchorY + x2, y2 := r.selection.currentX, r.selection.currentY + + if x1 > x2 { + x1, x2 = x2, x1 + } + if y1 > y2 { + y1, y2 = y2, y1 + } + + scaleX, scaleY := 1.0, 1.0 + if os.logicalW > 0 && os.screenBuf != nil { + scaleX = float64(os.screenBuf.Width) / float64(os.logicalW) + scaleY = float64(os.screenBuf.Height) / float64(os.logicalH) + } + + bx1, by1 := int32(x1*scaleX), int32(y1*scaleY) + bx2, by2 := int32(x2*scaleX), int32(y2*scaleY) + + w, h := bx2-bx1, by2-by1 + if w < 1 { + w = 1 + } + if h < 1 { + h = 1 + } + + r.result = Region{ + X: bx1 + os.output.x, + Y: by1 + os.output.y, + Width: w, + Height: h, + Output: os.output.name, + } + + r.running = false +} diff --git a/core/internal/screenshot/region_render.go b/core/internal/screenshot/region_render.go new file mode 100644 index 00000000..0056619d --- /dev/null +++ b/core/internal/screenshot/region_render.go @@ -0,0 +1,304 @@ +package screenshot + +import "fmt" + +var fontGlyphs = map[rune][12]uint8{ + '0': {0x3C, 0x66, 0x66, 0x6E, 0x76, 0x66, 0x66, 0x66, 0x66, 0x3C, 0x00, 0x00}, + '1': {0x18, 0x38, 0x78, 0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x7E, 0x00, 0x00}, + '2': {0x3C, 0x66, 0x66, 0x06, 0x0C, 0x18, 0x30, 0x60, 0x66, 0x7E, 0x00, 0x00}, + '3': {0x3C, 0x66, 0x06, 0x06, 0x1C, 0x06, 0x06, 0x06, 0x66, 0x3C, 0x00, 0x00}, + '4': {0x0C, 0x1C, 0x3C, 0x6C, 0xCC, 0xCC, 0xFE, 0x0C, 0x0C, 0x1E, 0x00, 0x00}, + '5': {0x7E, 0x60, 0x60, 0x60, 0x7C, 0x06, 0x06, 0x06, 0x66, 0x3C, 0x00, 0x00}, + '6': {0x1C, 0x30, 0x60, 0x60, 0x7C, 0x66, 0x66, 0x66, 0x66, 0x3C, 0x00, 0x00}, + '7': {0x7E, 0x66, 0x06, 0x06, 0x0C, 0x18, 0x18, 0x18, 0x18, 0x18, 0x00, 0x00}, + '8': {0x3C, 0x66, 0x66, 0x66, 0x3C, 0x66, 0x66, 0x66, 0x66, 0x3C, 0x00, 0x00}, + '9': {0x3C, 0x66, 0x66, 0x66, 0x3E, 0x06, 0x06, 0x06, 0x0C, 0x38, 0x00, 0x00}, + 'x': {0x00, 0x00, 0x00, 0x66, 0x66, 0x3C, 0x18, 0x3C, 0x66, 0x66, 0x00, 0x00}, + 'E': {0x7E, 0x60, 0x60, 0x60, 0x7C, 0x60, 0x60, 0x60, 0x60, 0x7E, 0x00, 0x00}, + 'P': {0x7C, 0x66, 0x66, 0x66, 0x7C, 0x60, 0x60, 0x60, 0x60, 0x60, 0x00, 0x00}, + 'S': {0x3C, 0x66, 0x60, 0x60, 0x3C, 0x06, 0x06, 0x06, 0x66, 0x3C, 0x00, 0x00}, + 'a': {0x00, 0x00, 0x00, 0x3C, 0x06, 0x3E, 0x66, 0x66, 0x66, 0x3E, 0x00, 0x00}, + 'c': {0x00, 0x00, 0x00, 0x3C, 0x66, 0x60, 0x60, 0x60, 0x66, 0x3C, 0x00, 0x00}, + 'd': {0x00, 0x00, 0x06, 0x06, 0x06, 0x3E, 0x66, 0x66, 0x66, 0x3E, 0x00, 0x00}, + 'e': {0x00, 0x00, 0x00, 0x3C, 0x66, 0x66, 0x7E, 0x60, 0x60, 0x3C, 0x00, 0x00}, + 'h': {0x00, 0x60, 0x60, 0x60, 0x7C, 0x66, 0x66, 0x66, 0x66, 0x66, 0x00, 0x00}, + 'i': {0x00, 0x18, 0x00, 0x38, 0x18, 0x18, 0x18, 0x18, 0x18, 0x3C, 0x00, 0x00}, + 'n': {0x00, 0x00, 0x00, 0x7C, 0x66, 0x66, 0x66, 0x66, 0x66, 0x66, 0x00, 0x00}, + 'o': {0x00, 0x00, 0x00, 0x3C, 0x66, 0x66, 0x66, 0x66, 0x66, 0x3C, 0x00, 0x00}, + 'p': {0x00, 0x00, 0x00, 0x7C, 0x66, 0x66, 0x66, 0x7C, 0x60, 0x60, 0x00, 0x00}, + 'r': {0x00, 0x00, 0x00, 0x6E, 0x76, 0x60, 0x60, 0x60, 0x60, 0x60, 0x00, 0x00}, + 's': {0x00, 0x00, 0x00, 0x3E, 0x60, 0x60, 0x3C, 0x06, 0x06, 0x7C, 0x00, 0x00}, + 't': {0x00, 0x18, 0x18, 0x7E, 0x18, 0x18, 0x18, 0x18, 0x18, 0x0E, 0x00, 0x00}, + 'u': {0x00, 0x00, 0x00, 0x66, 0x66, 0x66, 0x66, 0x66, 0x66, 0x3E, 0x00, 0x00}, + 'w': {0x00, 0x00, 0x00, 0x63, 0x63, 0x63, 0x6B, 0x7F, 0x77, 0x63, 0x00, 0x00}, + 'l': {0x38, 0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x3C, 0x00, 0x00}, + ' ': {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}, + ':': {0x00, 0x00, 0x18, 0x18, 0x00, 0x00, 0x00, 0x18, 0x18, 0x00, 0x00, 0x00}, + '/': {0x00, 0x02, 0x06, 0x0C, 0x18, 0x18, 0x30, 0x60, 0x40, 0x00, 0x00, 0x00}, + '[': {0x3C, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x3C, 0x00, 0x00}, + ']': {0x3C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x0C, 0x3C, 0x00, 0x00}, +} + +type OverlayStyle struct { + BackgroundR, BackgroundG, BackgroundB, BackgroundA uint8 + TextR, TextG, TextB uint8 + AccentR, AccentG, AccentB uint8 +} + +var DefaultOverlayStyle = OverlayStyle{ + BackgroundR: 30, BackgroundG: 30, BackgroundB: 30, BackgroundA: 220, + TextR: 255, TextG: 255, TextB: 255, + AccentR: 100, AccentG: 180, AccentB: 255, +} + +func (r *RegionSelector) drawOverlay(os *OutputSurface, renderBuf *ShmBuffer) { + data := renderBuf.Data() + stride := renderBuf.Stride + w, h := renderBuf.Width, renderBuf.Height + + // Dim the entire buffer + for y := 0; y < h; y++ { + off := y * stride + for x := 0; x < w; x++ { + i := off + x*4 + if i+3 >= len(data) { + continue + } + data[i+0] = uint8(int(data[i+0]) * 3 / 5) + data[i+1] = uint8(int(data[i+1]) * 3 / 5) + data[i+2] = uint8(int(data[i+2]) * 3 / 5) + } + } + + r.drawHUD(data, stride, w, h) + + if !r.selection.hasSelection || r.activeSurface != os { + return + } + + scaleX := float64(w) / float64(os.logicalW) + scaleY := float64(h) / float64(os.logicalH) + + bx1 := int(r.selection.anchorX * scaleX) + by1 := int(r.selection.anchorY * scaleY) + bx2 := int(r.selection.currentX * scaleX) + by2 := int(r.selection.currentY * scaleY) + + if bx1 > bx2 { + bx1, bx2 = bx2, bx1 + } + if by1 > by2 { + by1, by2 = by2, by1 + } + + bx1 = clamp(bx1, 0, w-1) + by1 = clamp(by1, 0, h-1) + bx2 = clamp(bx2, 0, w-1) + by2 = clamp(by2, 0, h-1) + + srcBuf := r.getSourceBuffer(os) + srcData := srcBuf.Data() + for y := by1; y <= by2; y++ { + rowOff := y * stride + for x := bx1; x <= bx2; x++ { + si := y*srcBuf.Stride + x*4 + di := rowOff + x*4 + if si+3 >= len(srcData) || di+3 >= len(data) { + continue + } + data[di+0] = srcData[si+0] + data[di+1] = srcData[si+1] + data[di+2] = srcData[si+2] + data[di+3] = srcData[si+3] + } + } + + selW, selH := bx2-bx1+1, by2-by1+1 + r.drawBorder(data, stride, w, h, bx1, by1, selW, selH) + r.drawDimensions(data, stride, w, h, bx1, by1, selW, selH) +} + +func (r *RegionSelector) drawHUD(data []byte, stride, bufW, bufH int) { + if r.selection.dragging { + return + } + + style := LoadOverlayStyle() + const charW, charH, padding, itemSpacing = 8, 12, 12, 24 + + cursorLabel := "hide" + if !r.showCapturedCursor { + cursorLabel = "show" + } + + items := []struct{ key, desc string }{ + {"Space/Enter", "capture"}, + {"P", cursorLabel + " cursor"}, + {"Esc", "cancel"}, + } + + totalW := 0 + for i, item := range items { + totalW += len(item.key)*(charW+1) + 4 + len(item.desc)*(charW+1) + if i < len(items)-1 { + totalW += itemSpacing + } + } + + hudW := totalW + padding*2 + hudH := charH + padding*2 + hudX := (bufW - hudW) / 2 + hudY := bufH - hudH - 20 + + r.fillRect(data, stride, bufW, bufH, hudX, hudY, hudW, hudH, + style.BackgroundR, style.BackgroundG, style.BackgroundB, style.BackgroundA) + + tx, ty := hudX+padding, hudY+padding + for i, item := range items { + r.drawText(data, stride, bufW, bufH, tx, ty, item.key, + style.AccentR, style.AccentG, style.AccentB) + tx += len(item.key) * (charW + 1) + + r.drawText(data, stride, bufW, bufH, tx, ty, " "+item.desc, + style.TextR, style.TextG, style.TextB) + tx += (1 + len(item.desc)) * (charW + 1) + + if i < len(items)-1 { + tx += itemSpacing + } + } +} + +func (r *RegionSelector) drawBorder(data []byte, stride, bufW, bufH, x, y, w, h int) { + const thickness = 2 + for i := 0; i < thickness; i++ { + r.drawHLine(data, stride, bufW, bufH, x-i, y-i, w+2*i) + r.drawHLine(data, stride, bufW, bufH, x-i, y+h+i-1, w+2*i) + r.drawVLine(data, stride, bufW, bufH, x-i, y-i, h+2*i) + r.drawVLine(data, stride, bufW, bufH, x+w+i-1, y-i, h+2*i) + } +} + +func (r *RegionSelector) drawHLine(data []byte, stride, bufW, bufH, x, y, length int) { + if y < 0 || y >= bufH { + return + } + rowOff := y * stride + for i := 0; i < length; i++ { + px := x + i + if px < 0 || px >= bufW { + continue + } + off := rowOff + px*4 + if off+3 >= len(data) { + continue + } + data[off], data[off+1], data[off+2], data[off+3] = 255, 255, 255, 255 + } +} + +func (r *RegionSelector) drawVLine(data []byte, stride, bufW, bufH, x, y, length int) { + if x < 0 || x >= bufW { + return + } + for i := 0; i < length; i++ { + py := y + i + if py < 0 || py >= bufH { + continue + } + off := py*stride + x*4 + if off+3 >= len(data) { + continue + } + data[off], data[off+1], data[off+2], data[off+3] = 255, 255, 255, 255 + } +} + +func (r *RegionSelector) drawDimensions(data []byte, stride, bufW, bufH, x, y, w, h int) { + text := fmt.Sprintf("%dx%d", w, h) + + const charW, charH = 8, 12 + textW := len(text) * (charW + 1) + textH := charH + + tx := x + (w-textW)/2 + ty := y + h + 8 + + if ty+textH > bufH { + ty = y - textH - 8 + } + tx = clamp(tx, 0, bufW-textW) + + r.fillRect(data, stride, bufW, bufH, tx-4, ty-2, textW+8, textH+4, 0, 0, 0, 200) + r.drawText(data, stride, bufW, bufH, tx, ty, text, 255, 255, 255) +} + +func (r *RegionSelector) fillRect(data []byte, stride, bufW, bufH, x, y, w, h int, br, bg, bb, ba uint8) { + alpha := float64(ba) / 255.0 + invAlpha := 1.0 - alpha + + for py := y; py < y+h && py < bufH; py++ { + if py < 0 { + continue + } + for px := x; px < x+w && px < bufW; px++ { + if px < 0 { + continue + } + off := py*stride + px*4 + if off+3 >= len(data) { + continue + } + data[off+0] = uint8(float64(data[off+0])*invAlpha + float64(bb)*alpha) + data[off+1] = uint8(float64(data[off+1])*invAlpha + float64(bg)*alpha) + data[off+2] = uint8(float64(data[off+2])*invAlpha + float64(br)*alpha) + data[off+3] = 255 + } + } +} + +func (r *RegionSelector) drawText(data []byte, stride, bufW, bufH, x, y int, text string, cr, cg, cb uint8) { + for i, ch := range text { + r.drawChar(data, stride, bufW, bufH, x+i*9, y, ch, cr, cg, cb) + } +} + +func (r *RegionSelector) drawChar(data []byte, stride, bufW, bufH, x, y int, ch rune, cr, cg, cb uint8) { + glyph, ok := fontGlyphs[ch] + if !ok { + return + } + + for row := 0; row < 12; row++ { + py := y + row + if py < 0 || py >= bufH { + continue + } + bits := glyph[row] + for col := 0; col < 8; col++ { + if (bits & (1 << (7 - col))) == 0 { + continue + } + px := x + col + if px < 0 || px >= bufW { + continue + } + off := py*stride + px*4 + if off+3 >= len(data) { + continue + } + data[off], data[off+1], data[off+2], data[off+3] = cb, cg, cr, 255 + } + } +} + +func clamp(v, lo, hi int) int { + switch { + case v < lo: + return lo + case v > hi: + return hi + default: + return v + } +} diff --git a/core/internal/screenshot/screenshot.go b/core/internal/screenshot/screenshot.go new file mode 100644 index 00000000..97ecb51b --- /dev/null +++ b/core/internal/screenshot/screenshot.go @@ -0,0 +1,588 @@ +package screenshot + +import ( + "fmt" + "sync" + + "github.com/AvengeMedia/DankMaterialShell/core/internal/log" + "github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wlr_screencopy" + wlhelpers "github.com/AvengeMedia/DankMaterialShell/core/internal/wayland/client" + "github.com/AvengeMedia/DankMaterialShell/core/pkg/go-wayland/wayland/client" +) + +type WaylandOutput struct { + wlOutput *client.Output + globalName uint32 + name string + x, y int32 + width int32 + height int32 + scale int32 + transform int32 +} + +type CaptureResult struct { + Buffer *ShmBuffer + Region Region + YInverted bool +} + +type Screenshoter struct { + config Config + + display *client.Display + registry *client.Registry + ctx *client.Context + + compositor *client.Compositor + shm *client.Shm + screencopy *wlr_screencopy.ZwlrScreencopyManagerV1 + + outputs map[uint32]*WaylandOutput + outputsMu sync.Mutex +} + +func New(config Config) *Screenshoter { + return &Screenshoter{ + config: config, + outputs: make(map[uint32]*WaylandOutput), + } +} + +func (s *Screenshoter) Run() (*CaptureResult, error) { + if err := s.connect(); err != nil { + return nil, fmt.Errorf("wayland connect: %w", err) + } + defer s.cleanup() + + if err := s.setupRegistry(); err != nil { + return nil, fmt.Errorf("registry setup: %w", err) + } + + if err := s.roundtrip(); err != nil { + return nil, fmt.Errorf("roundtrip: %w", err) + } + + if s.screencopy == nil { + return nil, fmt.Errorf("compositor does not support wlr-screencopy-unstable-v1") + } + + if err := s.roundtrip(); err != nil { + return nil, fmt.Errorf("roundtrip: %w", err) + } + + switch s.config.Mode { + case ModeLastRegion: + return s.captureLastRegion() + case ModeRegion: + return s.captureRegion() + case ModeOutput: + return s.captureOutput(s.config.OutputName) + case ModeFullScreen: + return s.captureFullScreen() + case ModeAllScreens: + return s.captureAllScreens() + default: + return s.captureRegion() + } +} + +func (s *Screenshoter) captureLastRegion() (*CaptureResult, error) { + lastRegion := GetLastRegion() + if lastRegion.IsEmpty() { + return s.captureRegion() + } + + output := s.findOutputForRegion(lastRegion) + if output == nil { + return s.captureRegion() + } + + return s.captureRegionOnOutput(output, lastRegion) +} + +func (s *Screenshoter) captureRegion() (*CaptureResult, error) { + selector := NewRegionSelector(s) + region, cancelled, err := selector.Run() + if err != nil { + return nil, fmt.Errorf("region selection: %w", err) + } + if cancelled { + return nil, nil + } + + output := s.findOutputForRegion(region) + if output == nil { + return nil, fmt.Errorf("no output found for region") + } + + if err := SaveLastRegion(region); err != nil { + log.Debug("failed to save last region", "err", err) + } + + return s.captureRegionOnOutput(output, region) +} + +func (s *Screenshoter) captureFullScreen() (*CaptureResult, error) { + output := s.findFocusedOutput() + if output == nil { + s.outputsMu.Lock() + for _, o := range s.outputs { + output = o + break + } + s.outputsMu.Unlock() + } + + if output == nil { + return nil, fmt.Errorf("no output available") + } + + return s.captureWholeOutput(output) +} + +func (s *Screenshoter) captureOutput(name string) (*CaptureResult, error) { + s.outputsMu.Lock() + var output *WaylandOutput + for _, o := range s.outputs { + if o.name == name { + output = o + break + } + } + s.outputsMu.Unlock() + + if output == nil { + return nil, fmt.Errorf("output %q not found", name) + } + + return s.captureWholeOutput(output) +} + +func (s *Screenshoter) captureAllScreens() (*CaptureResult, error) { + s.outputsMu.Lock() + outputs := make([]*WaylandOutput, 0, len(s.outputs)) + var minX, minY, maxX, maxY int32 + first := true + + for _, o := range s.outputs { + outputs = append(outputs, o) + right := o.x + o.width + bottom := o.y + o.height + + if first { + minX, minY = o.x, o.y + maxX, maxY = right, bottom + first = false + continue + } + + if o.x < minX { + minX = o.x + } + if o.y < minY { + minY = o.y + } + if right > maxX { + maxX = right + } + if bottom > maxY { + maxY = bottom + } + } + s.outputsMu.Unlock() + + if len(outputs) == 0 { + return nil, fmt.Errorf("no outputs available") + } + + if len(outputs) == 1 { + return s.captureWholeOutput(outputs[0]) + } + + totalW := maxX - minX + totalH := maxY - minY + + compositeStride := int(totalW) * 4 + composite, err := CreateShmBuffer(int(totalW), int(totalH), compositeStride) + if err != nil { + return nil, fmt.Errorf("create composite buffer: %w", err) + } + + composite.Clear() + + for _, output := range outputs { + result, err := s.captureWholeOutput(output) + if err != nil { + log.Warn("failed to capture output", "name", output.name, "err", err) + continue + } + + s.blitBuffer(composite, result.Buffer, int(output.x-minX), int(output.y-minY), result.YInverted) + result.Buffer.Close() + } + + return &CaptureResult{ + Buffer: composite, + Region: Region{X: minX, Y: minY, Width: totalW, Height: totalH}, + }, nil +} + +func (s *Screenshoter) blitBuffer(dst, src *ShmBuffer, dstX, dstY int, yInverted bool) { + srcData := src.Data() + dstData := dst.Data() + + for srcY := 0; srcY < src.Height; srcY++ { + actualSrcY := srcY + if yInverted { + actualSrcY = src.Height - 1 - srcY + } + + dy := dstY + srcY + if dy < 0 || dy >= dst.Height { + continue + } + + srcRowOff := actualSrcY * src.Stride + dstRowOff := dy * dst.Stride + + for srcX := 0; srcX < src.Width; srcX++ { + dx := dstX + srcX + if dx < 0 || dx >= dst.Width { + continue + } + + si := srcRowOff + srcX*4 + di := dstRowOff + dx*4 + + if si+3 >= len(srcData) || di+3 >= len(dstData) { + continue + } + + dstData[di+0] = srcData[si+0] + dstData[di+1] = srcData[si+1] + dstData[di+2] = srcData[si+2] + dstData[di+3] = srcData[si+3] + } + } +} + +func (s *Screenshoter) captureWholeOutput(output *WaylandOutput) (*CaptureResult, error) { + cursor := int32(0) + if s.config.IncludeCursor { + cursor = 1 + } + + frame, err := s.screencopy.CaptureOutput(cursor, output.wlOutput) + if err != nil { + return nil, fmt.Errorf("capture output: %w", err) + } + + return s.processFrame(frame, Region{ + X: output.x, + Y: output.y, + Width: output.width, + Height: output.height, + Output: output.name, + }) +} + +func (s *Screenshoter) captureRegionOnOutput(output *WaylandOutput, region Region) (*CaptureResult, error) { + localX := region.X - output.x + localY := region.Y - output.y + + cursor := int32(0) + if s.config.IncludeCursor { + cursor = 1 + } + + frame, err := s.screencopy.CaptureOutputRegion( + cursor, + output.wlOutput, + localX, localY, + region.Width, region.Height, + ) + if err != nil { + return nil, fmt.Errorf("capture region: %w", err) + } + + return s.processFrame(frame, region) +} + +func (s *Screenshoter) processFrame(frame *wlr_screencopy.ZwlrScreencopyFrameV1, region Region) (*CaptureResult, error) { + var buf *ShmBuffer + var format PixelFormat + var yInverted bool + ready := false + failed := false + + frame.SetBufferHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1BufferEvent) { + var err error + buf, err = CreateShmBuffer(int(e.Width), int(e.Height), int(e.Stride)) + if err != nil { + log.Error("failed to create buffer", "err", err) + return + } + format = PixelFormat(e.Format) + buf.Format = format + }) + + frame.SetFlagsHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1FlagsEvent) { + yInverted = (e.Flags & 1) != 0 + }) + + frame.SetBufferDoneHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1BufferDoneEvent) { + if buf == nil { + return + } + + pool, err := s.shm.CreatePool(buf.Fd(), int32(buf.Size())) + if err != nil { + log.Error("failed to create pool", "err", err) + return + } + + wlBuf, err := pool.CreateBuffer(0, int32(buf.Width), int32(buf.Height), int32(buf.Stride), uint32(format)) + if err != nil { + pool.Destroy() + log.Error("failed to create wl_buffer", "err", err) + return + } + + if err := frame.Copy(wlBuf); err != nil { + log.Error("failed to copy frame", "err", err) + } + + pool.Destroy() + }) + + frame.SetReadyHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1ReadyEvent) { + ready = true + }) + + frame.SetFailedHandler(func(e wlr_screencopy.ZwlrScreencopyFrameV1FailedEvent) { + failed = true + }) + + for !ready && !failed { + if err := s.ctx.Dispatch(); err != nil { + frame.Destroy() + return nil, fmt.Errorf("dispatch: %w", err) + } + } + + frame.Destroy() + + if failed { + if buf != nil { + buf.Close() + } + return nil, fmt.Errorf("frame capture failed") + } + + return &CaptureResult{ + Buffer: buf, + Region: region, + YInverted: yInverted, + }, nil +} + +func (s *Screenshoter) findOutputForRegion(region Region) *WaylandOutput { + s.outputsMu.Lock() + defer s.outputsMu.Unlock() + + cx := region.X + region.Width/2 + cy := region.Y + region.Height/2 + + for _, o := range s.outputs { + if cx >= o.x && cx < o.x+o.width && cy >= o.y && cy < o.y+o.height { + return o + } + } + + for _, o := range s.outputs { + if region.X >= o.x && region.X < o.x+o.width && + region.Y >= o.y && region.Y < o.y+o.height { + return o + } + } + + return nil +} + +func (s *Screenshoter) findFocusedOutput() *WaylandOutput { + s.outputsMu.Lock() + defer s.outputsMu.Unlock() + for _, o := range s.outputs { + return o + } + return nil +} + +func (s *Screenshoter) connect() error { + display, err := client.Connect("") + if err != nil { + return err + } + s.display = display + s.ctx = display.Context() + return nil +} + +func (s *Screenshoter) roundtrip() error { + return wlhelpers.Roundtrip(s.display, s.ctx) +} + +func (s *Screenshoter) setupRegistry() error { + registry, err := s.display.GetRegistry() + if err != nil { + return err + } + s.registry = registry + + registry.SetGlobalHandler(func(e client.RegistryGlobalEvent) { + s.handleGlobal(e) + }) + + registry.SetGlobalRemoveHandler(func(e client.RegistryGlobalRemoveEvent) { + s.outputsMu.Lock() + delete(s.outputs, e.Name) + s.outputsMu.Unlock() + }) + + return nil +} + +func (s *Screenshoter) handleGlobal(e client.RegistryGlobalEvent) { + switch e.Interface { + case client.CompositorInterfaceName: + comp := client.NewCompositor(s.ctx) + if err := s.registry.Bind(e.Name, e.Interface, e.Version, comp); err == nil { + s.compositor = comp + } + + case client.ShmInterfaceName: + shm := client.NewShm(s.ctx) + if err := s.registry.Bind(e.Name, e.Interface, e.Version, shm); err == nil { + s.shm = shm + } + + case client.OutputInterfaceName: + output := client.NewOutput(s.ctx) + version := e.Version + if version > 4 { + version = 4 + } + if err := s.registry.Bind(e.Name, e.Interface, version, output); err == nil { + s.outputsMu.Lock() + s.outputs[e.Name] = &WaylandOutput{ + wlOutput: output, + globalName: e.Name, + scale: 1, + } + s.outputsMu.Unlock() + s.setupOutputHandlers(e.Name, output) + } + + case wlr_screencopy.ZwlrScreencopyManagerV1InterfaceName: + sc := wlr_screencopy.NewZwlrScreencopyManagerV1(s.ctx) + version := e.Version + if version > 3 { + version = 3 + } + if err := s.registry.Bind(e.Name, e.Interface, version, sc); err == nil { + s.screencopy = sc + } + } +} + +func (s *Screenshoter) setupOutputHandlers(name uint32, output *client.Output) { + output.SetGeometryHandler(func(e client.OutputGeometryEvent) { + s.outputsMu.Lock() + if o, ok := s.outputs[name]; ok { + o.x, o.y = e.X, e.Y + o.transform = int32(e.Transform) + } + s.outputsMu.Unlock() + }) + + output.SetModeHandler(func(e client.OutputModeEvent) { + if e.Flags&uint32(client.OutputModeCurrent) == 0 { + return + } + s.outputsMu.Lock() + if o, ok := s.outputs[name]; ok { + o.width, o.height = e.Width, e.Height + } + s.outputsMu.Unlock() + }) + + output.SetScaleHandler(func(e client.OutputScaleEvent) { + s.outputsMu.Lock() + if o, ok := s.outputs[name]; ok { + o.scale = e.Factor + } + s.outputsMu.Unlock() + }) + + output.SetNameHandler(func(e client.OutputNameEvent) { + s.outputsMu.Lock() + if o, ok := s.outputs[name]; ok { + o.name = e.Name + } + s.outputsMu.Unlock() + }) +} + +func (s *Screenshoter) cleanup() { + if s.screencopy != nil { + s.screencopy.Destroy() + } + if s.display != nil { + s.ctx.Close() + } +} + +func (s *Screenshoter) GetOutputs() []*WaylandOutput { + s.outputsMu.Lock() + defer s.outputsMu.Unlock() + out := make([]*WaylandOutput, 0, len(s.outputs)) + for _, o := range s.outputs { + out = append(out, o) + } + return out +} + +func ListOutputs() ([]Output, error) { + sc := New(DefaultConfig()) + if err := sc.connect(); err != nil { + return nil, err + } + defer sc.cleanup() + + if err := sc.setupRegistry(); err != nil { + return nil, err + } + if err := sc.roundtrip(); err != nil { + return nil, err + } + if err := sc.roundtrip(); err != nil { + return nil, err + } + + sc.outputsMu.Lock() + defer sc.outputsMu.Unlock() + + result := make([]Output, 0, len(sc.outputs)) + for _, o := range sc.outputs { + result = append(result, Output{ + Name: o.name, + X: o.x, + Y: o.y, + Width: o.width, + Height: o.height, + Scale: o.scale, + }) + } + return result, nil +} diff --git a/core/internal/screenshot/shm.go b/core/internal/screenshot/shm.go new file mode 100644 index 00000000..758c828a --- /dev/null +++ b/core/internal/screenshot/shm.go @@ -0,0 +1,18 @@ +package screenshot + +import "github.com/AvengeMedia/DankMaterialShell/core/internal/wayland/shm" + +type PixelFormat = shm.PixelFormat + +const ( + FormatARGB8888 = shm.FormatARGB8888 + FormatXRGB8888 = shm.FormatXRGB8888 + FormatABGR8888 = shm.FormatABGR8888 + FormatXBGR8888 = shm.FormatXBGR8888 +) + +type ShmBuffer = shm.Buffer + +func CreateShmBuffer(width, height, stride int) (*ShmBuffer, error) { + return shm.CreateBuffer(width, height, stride) +} diff --git a/core/internal/screenshot/state.go b/core/internal/screenshot/state.go new file mode 100644 index 00000000..168e27f1 --- /dev/null +++ b/core/internal/screenshot/state.go @@ -0,0 +1,65 @@ +package screenshot + +import ( + "encoding/json" + "os" + "path" + "path/filepath" +) + +type PersistentState struct { + LastRegion Region `json:"last_region"` +} + +func getStateFilePath() string { + cacheDir, err := os.UserCacheDir() + if err != nil { + cacheDir = path.Join(os.Getenv("HOME"), ".cache") + } + return filepath.Join(cacheDir, "dms", "screenshot-state.json") +} + +func LoadState() (*PersistentState, error) { + path := getStateFilePath() + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return &PersistentState{}, nil + } + return nil, err + } + + var state PersistentState + if err := json.Unmarshal(data, &state); err != nil { + return &PersistentState{}, nil + } + return &state, nil +} + +func SaveState(state *PersistentState) error { + path := getStateFilePath() + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return err + } + + data, err := json.MarshalIndent(state, "", " ") + if err != nil { + return err + } + return os.WriteFile(path, data, 0644) +} + +func GetLastRegion() Region { + state, err := LoadState() + if err != nil { + return Region{} + } + return state.LastRegion +} + +func SaveLastRegion(r Region) error { + state, _ := LoadState() + state.LastRegion = r + return SaveState(state) +} diff --git a/core/internal/screenshot/theme.go b/core/internal/screenshot/theme.go new file mode 100644 index 00000000..2398bce7 --- /dev/null +++ b/core/internal/screenshot/theme.go @@ -0,0 +1,127 @@ +package screenshot + +import ( + "encoding/json" + "os" + "os/exec" + "path/filepath" + "strings" +) + +type ThemeColors struct { + Background string `json:"surface"` + OnSurface string `json:"on_surface"` + Primary string `json:"primary"` +} + +type ColorScheme struct { + Dark ThemeColors `json:"dark"` + Light ThemeColors `json:"light"` +} + +type ColorsFile struct { + Colors ColorScheme `json:"colors"` +} + +var cachedStyle *OverlayStyle + +func LoadOverlayStyle() OverlayStyle { + if cachedStyle != nil { + return *cachedStyle + } + + style := DefaultOverlayStyle + colors := loadColorsFile() + if colors == nil { + cachedStyle = &style + return style + } + + theme := &colors.Dark + if isLightMode() { + theme = &colors.Light + } + + if bg, ok := parseHexColor(theme.Background); ok { + style.BackgroundR, style.BackgroundG, style.BackgroundB = bg[0], bg[1], bg[2] + } + if text, ok := parseHexColor(theme.OnSurface); ok { + style.TextR, style.TextG, style.TextB = text[0], text[1], text[2] + } + if accent, ok := parseHexColor(theme.Primary); ok { + style.AccentR, style.AccentG, style.AccentB = accent[0], accent[1], accent[2] + } + + cachedStyle = &style + return style +} + +func loadColorsFile() *ColorScheme { + path := getColorsFilePath() + data, err := os.ReadFile(path) + if err != nil { + return nil + } + + var file ColorsFile + if err := json.Unmarshal(data, &file); err != nil { + return nil + } + + return &file.Colors +} + +func getColorsFilePath() string { + cacheDir := os.Getenv("XDG_CACHE_HOME") + if cacheDir == "" { + home := os.Getenv("HOME") + if home == "" { + return "" + } + cacheDir = filepath.Join(home, ".cache") + } + return filepath.Join(cacheDir, "DankMaterialShell", "dms-colors.json") +} + +func isLightMode() bool { + out, err := exec.Command("gsettings", "get", "org.gnome.desktop.interface", "color-scheme").Output() + if err != nil { + return false + } + + scheme := strings.TrimSpace(string(out)) + switch scheme { + case "'prefer-light'", "'default'": + return true + } + return false +} + +func parseHexColor(hex string) ([3]uint8, bool) { + hex = strings.TrimPrefix(hex, "#") + if len(hex) != 6 { + return [3]uint8{}, false + } + + var r, g, b uint8 + for i, ptr := range []*uint8{&r, &g, &b} { + val := 0 + for j := 0; j < 2; j++ { + c := hex[i*2+j] + val *= 16 + switch { + case c >= '0' && c <= '9': + val += int(c - '0') + case c >= 'a' && c <= 'f': + val += int(c - 'a' + 10) + case c >= 'A' && c <= 'F': + val += int(c - 'A' + 10) + default: + return [3]uint8{}, false + } + } + *ptr = uint8(val) + } + + return [3]uint8{r, g, b}, true +} diff --git a/core/internal/screenshot/types.go b/core/internal/screenshot/types.go new file mode 100644 index 00000000..c1ba1847 --- /dev/null +++ b/core/internal/screenshot/types.go @@ -0,0 +1,68 @@ +package screenshot + +type Mode int + +const ( + ModeRegion Mode = iota + ModeWindow + ModeFullScreen + ModeAllScreens + ModeOutput + ModeLastRegion +) + +type Format int + +const ( + FormatPNG Format = iota + FormatJPEG + FormatPPM +) + +type Region struct { + X int32 `json:"x"` + Y int32 `json:"y"` + Width int32 `json:"width"` + Height int32 `json:"height"` + Output string `json:"output,omitempty"` +} + +func (r Region) IsEmpty() bool { + return r.Width <= 0 || r.Height <= 0 +} + +type Output struct { + Name string + X, Y int32 + Width int32 + Height int32 + Scale int32 +} + +type Config struct { + Mode Mode + OutputName string + IncludeCursor bool + Format Format + Quality int + OutputDir string + Filename string + Clipboard bool + Freeze bool + Notify bool + Stdout bool +} + +func DefaultConfig() Config { + return Config{ + Mode: ModeRegion, + IncludeCursor: false, + Format: FormatPNG, + Quality: 90, + OutputDir: "", + Filename: "", + Clipboard: false, + Freeze: true, + Notify: true, + } +} diff --git a/core/internal/wayland/client/helpers.go b/core/internal/wayland/client/helpers.go new file mode 100644 index 00000000..9075cd38 --- /dev/null +++ b/core/internal/wayland/client/helpers.go @@ -0,0 +1,26 @@ +package client + +import wlclient "github.com/AvengeMedia/DankMaterialShell/core/pkg/go-wayland/wayland/client" + +func Roundtrip(display *wlclient.Display, ctx *wlclient.Context) error { + callback, err := display.Sync() + if err != nil { + return err + } + + done := make(chan struct{}) + callback.SetDoneHandler(func(e wlclient.CallbackDoneEvent) { + close(done) + }) + + for { + select { + case <-done: + return nil + default: + if err := ctx.Dispatch(); err != nil { + return err + } + } + } +} diff --git a/core/internal/wayland/shm/buffer.go b/core/internal/wayland/shm/buffer.go new file mode 100644 index 00000000..d41fe6ad --- /dev/null +++ b/core/internal/wayland/shm/buffer.go @@ -0,0 +1,139 @@ +package shm + +import ( + "fmt" + + "golang.org/x/sys/unix" +) + +type PixelFormat uint32 + +const ( + FormatARGB8888 PixelFormat = 0 + FormatXRGB8888 PixelFormat = 1 + FormatABGR8888 PixelFormat = 0x34324241 + FormatXBGR8888 PixelFormat = 0x34324258 +) + +type Buffer struct { + fd int + data []byte + size int + Width int + Height int + Stride int + Format PixelFormat +} + +func CreateBuffer(width, height, stride int) (*Buffer, error) { + size := stride * height + + fd, err := unix.MemfdCreate("dms-shm", 0) + if err != nil { + return nil, fmt.Errorf("memfd_create: %w", err) + } + + if err := unix.Ftruncate(fd, int64(size)); err != nil { + unix.Close(fd) + return nil, fmt.Errorf("ftruncate: %w", err) + } + + data, err := unix.Mmap(fd, 0, size, unix.PROT_READ|unix.PROT_WRITE, unix.MAP_SHARED) + if err != nil { + unix.Close(fd) + return nil, fmt.Errorf("mmap: %w", err) + } + + return &Buffer{ + fd: fd, + data: data, + size: size, + Width: width, + Height: height, + Stride: stride, + }, nil +} + +func (b *Buffer) Fd() int { return b.fd } +func (b *Buffer) Size() int { return b.size } +func (b *Buffer) Data() []byte { return b.data } + +func (b *Buffer) Close() error { + var firstErr error + + if b.data != nil { + if err := unix.Munmap(b.data); err != nil && firstErr == nil { + firstErr = fmt.Errorf("munmap: %w", err) + } + b.data = nil + } + + if b.fd >= 0 { + if err := unix.Close(b.fd); err != nil && firstErr == nil { + firstErr = fmt.Errorf("close: %w", err) + } + b.fd = -1 + } + + return firstErr +} + +func (b *Buffer) GetPixelRGBA(x, y int) (r, g, b2, a uint8) { + if x < 0 || x >= b.Width || y < 0 || y >= b.Height { + return + } + + off := y*b.Stride + x*4 + if off+3 >= len(b.data) { + return + } + + return b.data[off+2], b.data[off+1], b.data[off], b.data[off+3] +} + +func (b *Buffer) GetPixelBGRA(x, y int) (b2, g, r, a uint8) { + if x < 0 || x >= b.Width || y < 0 || y >= b.Height { + return + } + + off := y*b.Stride + x*4 + if off+3 >= len(b.data) { + return + } + + return b.data[off], b.data[off+1], b.data[off+2], b.data[off+3] +} + +func (b *Buffer) ConvertBGRAtoRGBA() { + for y := 0; y < b.Height; y++ { + rowOff := y * b.Stride + for x := 0; x < b.Width; x++ { + off := rowOff + x*4 + if off+3 >= len(b.data) { + continue + } + b.data[off], b.data[off+2] = b.data[off+2], b.data[off] + } + } +} + +func (b *Buffer) FlipVertical() { + tmp := make([]byte, b.Stride) + for y := 0; y < b.Height/2; y++ { + topOff := y * b.Stride + botOff := (b.Height - 1 - y) * b.Stride + copy(tmp, b.data[topOff:topOff+b.Stride]) + copy(b.data[topOff:topOff+b.Stride], b.data[botOff:botOff+b.Stride]) + copy(b.data[botOff:botOff+b.Stride], tmp) + } +} + +func (b *Buffer) Clear() { + for i := range b.data { + b.data[i] = 0 + } +} + +func (b *Buffer) CopyFrom(src *Buffer) { + copy(b.data, src.data) +}