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
	interval      int
	bufferCount   int
	tempDirectory string
	initialized   bool
	running       bool

	files  []string
	latest string
	done   chan struct{}
	sync.WaitGroup
	mu sync.RWMutex
}

// Get implements ScreenCapturer.
func (w *wholeScreenCapturer) Get() (filepath string, err error) {
	w.mu.RLock()
	defer w.mu.Unlock()
    if !w.initialized {
        return "", ErrNotInitialized
    }
    if w.latest == "" {
        return "", errors.New("no screenshot captured yet")
    }
    return w.latest, nil
}

// Init implements ScreenCapturer.
func (w *wholeScreenCapturer) Init() (err error) {

	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.interval = w.config.Screenshot.Interval
	w.log.Debug("screenshot interval set", "interval", w.interval)

	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.latest = ""
	w.initialized = true
	return nil
}

// Start implements ScreenCapturer.
func (w *wholeScreenCapturer) Start() error {
	if !w.initialized {
		return ErrNotInitialized
	}
	if w.running {
		w.log.Debug("wholescreencapturer is already running, ignoring start")
		return nil
	}
	w.running = true
	w.Go(func() {
		w.captureLoop()
	})
	w.log.Debug("wholescreencapturer started")
	return nil
}

// Stop implements ScreenCapturer.
func (w *wholeScreenCapturer) Stop() error {
	if !w.initialized {
		return ErrNotInitialized
	}
	if !w.running {
		w.log.Debug("wholescreencapturer is not running, ignoring stop")
		return nil
	}
	w.log.Debug("stopping wholescreencapturer")
	close(w.done)
	w.Wait()
	w.running = false
	w.log.Debug("wholescreencapturer stopped")
	return nil
}

func (w *wholeScreenCapturer) captureLoop() {
	ticker := time.NewTicker(time.Duration(w.interval) * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-w.done:
			w.log.Debug("capture loop stopped")
			return
		case <-ticker.C:
			fp, err := w.captureAndSave()
			if err != nil {
				w.log.Error("screenshot capture failed", "error", err)
				continue
			}
			w.addToBuffer(fp)
		}
	}
}

func (w *wholeScreenCapturer) captureAndSave() (string, error) {
	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)

	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)
	w.latest = 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,
		done:   make(chan struct{}),
	}
	lc.Append(fx.StopHook(func(ctx context.Context) error {
		if !capturer.initialized {
			log.Debug("wholescreencapturer not initialized, nothing to do")
			return nil
		}

		log.Debug("stopping wholescreencapturer")

		err := capturer.Stop()
		if err != nil {
			log.Error("failed to stop wholescreencapturer gracefully")
			return err
		}

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