visionA/local-agent/server/internal/camera/multi_image_source.go

package camera

import (
	"bytes"
	"fmt"
	"image"
	"image/jpeg"
	"image/png"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// BatchImageEntry holds metadata and decoded JPEG data for a single image in a batch.
type BatchImageEntry struct {
	Filename string
	FilePath string
	JpegData []byte
	Width    int
	Height   int
}

// MultiImageSource provides sequential JPEG frames from multiple uploaded images.
// It implements FrameSource. The pipeline calls ReadFrame for the current image,
// then Advance to move to the next one. After the last image, ReadFrame returns an error.
type MultiImageSource struct {
	images     []BatchImageEntry
	currentIdx int
	mu         sync.Mutex
}

// NewMultiImageSource creates a source from multiple file paths.
// Each file is decoded (JPG/PNG) and converted to JPEG in memory.
func NewMultiImageSource(filePaths []string, filenames []string) (*MultiImageSource, error) {
	if len(filePaths) != len(filenames) {
		return nil, fmt.Errorf("filePaths and filenames length mismatch")
	}
	entries := make([]BatchImageEntry, 0, len(filePaths))
	for i, fp := range filePaths {
		entry, err := loadBatchImageEntry(fp, filenames[i])
		if err != nil {
			// Clean up already-loaded temp files
			for _, e := range entries {
				os.Remove(e.FilePath)
			}
			return nil, fmt.Errorf("image %d (%s): %w", i, filenames[i], err)
		}
		entries = append(entries, entry)
	}
	return &MultiImageSource{images: entries}, nil
}

func loadBatchImageEntry(filePath, filename string) (BatchImageEntry, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return BatchImageEntry{}, fmt.Errorf("failed to open: %w", err)
	}
	defer f.Close()

	ext := strings.ToLower(filepath.Ext(filePath))
	var img image.Image
	switch ext {
	case ".jpg", ".jpeg":
		img, err = jpeg.Decode(f)
	case ".png":
		img, err = png.Decode(f)
	default:
		return BatchImageEntry{}, fmt.Errorf("unsupported format: %s", ext)
	}
	if err != nil {
		return BatchImageEntry{}, fmt.Errorf("failed to decode: %w", err)
	}

	bounds := img.Bounds()
	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: 90}); err != nil {
		return BatchImageEntry{}, fmt.Errorf("failed to encode JPEG: %w", err)
	}

	return BatchImageEntry{
		Filename: filename,
		FilePath: filePath,
		JpegData: buf.Bytes(),
		Width:    bounds.Dx(),
		Height:   bounds.Dy(),
	}, nil
}

// ReadFrame returns the current image's JPEG data.
func (s *MultiImageSource) ReadFrame() ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.currentIdx >= len(s.images) {
		return nil, fmt.Errorf("all images consumed")
	}
	return s.images[s.currentIdx].JpegData, nil
}

// Advance moves to the next image. Returns false if no more images remain.
func (s *MultiImageSource) Advance() bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.currentIdx++
	return s.currentIdx < len(s.images)
}

// CurrentIndex returns the 0-based index of the current image.
func (s *MultiImageSource) CurrentIndex() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.currentIdx
}

// CurrentEntry returns metadata for the current image.
func (s *MultiImageSource) CurrentEntry() BatchImageEntry {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.images[s.currentIdx]
}

// TotalImages returns the number of images in the batch.
func (s *MultiImageSource) TotalImages() int {
	return len(s.images)
}

// GetImageByIndex returns JPEG data for a specific image by index.
func (s *MultiImageSource) GetImageByIndex(index int) ([]byte, error) {
	if index < 0 || index >= len(s.images) {
		return nil, fmt.Errorf("image index %d out of range [0, %d)", index, len(s.images))
	}
	return s.images[index].JpegData, nil
}

// Images returns all batch entries.
func (s *MultiImageSource) Images() []BatchImageEntry {
	return s.images
}

// Close removes all temporary files.
func (s *MultiImageSource) Close() error {
	for _, entry := range s.images {
		os.Remove(entry.FilePath)
	}
	return nil
}