auto-attendance/internal/screencapturer/wholescreencapturer.go

// Copyright (c) 2025 Nikolai Papin
//
// This file is part of the Auto Attendance app that looks for
// self-attend QR-codes during lectures and opens their URLs in your
// browser.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See
// the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package screencapturer

import (
	"context"
	"errors"
	"fmt"
	"image"
	"image/png"
	"os"
	"path/filepath"
	"sync"
	"time"

	"git.weirdcat.su/weirdcat/auto-attendance/internal/config"
	"git.weirdcat.su/weirdcat/auto-attendance/internal/constants"
	"git.weirdcat.su/weirdcat/auto-attendance/internal/logger"
	"github.com/kbinani/screenshot"
	"go.uber.org/fx"
)

var (
	ErrNotInitialized = errors.New("wholescreencapturer was not initialized")
)

type wholeScreenCapturer struct {
	config *config.Config
	log    *logger.Logger

	displayIndex  int
	displayBounds image.Rectangle
	bufferCount   int
	tempDirectory string
	initialized   bool

	files []string
	mu    sync.RWMutex
}

// Get implements ScreenCapturer.
func (w *wholeScreenCapturer) Get() (filepath string, err error) {
	w.mu.RLock()
	initialized := w.initialized
	w.mu.RUnlock()
	
	if !initialized {
		return "", ErrNotInitialized
	}

	fp, err := w.captureAndSave()
	if err != nil {
		w.log.Error("screenshot capture failed", "error", err)
		return "", err
	}

	w.addToBuffer(fp)
	return fp, nil
}

// Init implements ScreenCapturer.
func (w *wholeScreenCapturer) Init() (err error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	
	if w.initialized {
		w.log.Debug("wholescreencapturer already initialized, skipping initialization")
		return nil
	}

	displayCount := screenshot.NumActiveDisplays()
	w.displayIndex = w.config.Screenshot.ScreenIndex % displayCount
	w.log.Debug("detected displays", "count", displayCount, "selected", w.displayIndex)

	w.displayBounds = screenshot.GetDisplayBounds(w.displayIndex)
	w.log.Debug(
		"display bounds set",
		"dx", w.displayBounds.Dx(),
		"dy", w.displayBounds.Dy(),
	)

	w.bufferCount = w.config.Screenshot.BufferCount
	w.log.Debug("screenshot buffer count set", "count", w.bufferCount)

	if w.tempDirectory, err = w.createTempDirectory(); err != nil {
		w.log.Error("failed to create temporary directory", "error", err)
		return err
	}
	w.log.Debug("temporary directory created", "path", w.tempDirectory)

	w.files = make([]string, 0, w.bufferCount)
	w.initialized = true
	return nil
}

func (w *wholeScreenCapturer) captureAndSave() (string, error) {
	w.mu.RLock()
	img, err := screenshot.CaptureRect(w.displayBounds)
	if err != nil {
		w.log.Error("failed to capture screenshot", "error", err)
		return "", err
	}

	now := time.Now().UnixMilli()
	filename := fmt.Sprintf("%d.png", now)
	filePath := filepath.Join(w.tempDirectory, filename)
	w.mu.RUnlock()

	file, err := os.Create(filePath)
	if err != nil {
		w.log.Error("failed to create screenshot file", "path", filePath, "error", err)
		return "", err
	}

	defer file.Close()
	err = png.Encode(file, img)
	if err != nil {
		w.log.Error("failed to encode image into file", "path", filePath, "error", err)
		return "", err
	}

	w.log.Debug("Screenshot saved", "path", filePath)
	return filePath, nil
}

func (w *wholeScreenCapturer) addToBuffer(fp string) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.files = append(w.files, fp)

	if len(w.files) > w.bufferCount {
		old := w.files[0]
		if err := os.Remove(old); err != nil {
			w.log.Warn("failed to remove old screenshot", "path", old, "error", err)
		} else {
			w.log.Debug("removed old screenshot", "path", old)
		}
		w.files = w.files[1:]
	}
}

func (w *wholeScreenCapturer) createTempDirectory() (path string, err error) {
	return os.MkdirTemp(w.config.Screenshot.Directory, constants.AppName+"-*")
}

func NewWholeScreenCapturer(
	lc fx.Lifecycle, config *config.Config, log *logger.Logger,
) ScreenCapturer {
	capturer := &wholeScreenCapturer{
		config: config,
		log:    log,
	}
	lc.Append(fx.StopHook(func(ctx context.Context) error {
		capturer.mu.RLock()
		initialized := capturer.initialized
		capturer.mu.RUnlock()
		
		if !initialized {
			log.Debug("wholescreencapturer not initialized, nothing to do")
			return nil
		}

		log.Debug("cleaning up wholescreencapturer")

		// Clean up all screenshot files
		capturer.mu.Lock()
		defer capturer.mu.Unlock()
		for _, file := range capturer.files {
			if err := os.Remove(file); err != nil {
				log.Warn("failed to remove screenshot file during cleanup", "path", file, "error", err)
			}
		}

		err := os.RemoveAll(capturer.tempDirectory)
		if err != nil {
			log.Error("failed to remove temp directory")
			return err
		}
		return nil
	}))
	return capturer
}