Search/cache-images.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"image"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/chai2010/webp"
	"golang.org/x/image/bmp"
	"golang.org/x/image/tiff"
)

var (
	cachingImages    = make(map[string]*sync.Mutex)
	cachingImagesMu  sync.Mutex
	cachingSemaphore = make(chan struct{}, 10) // Limit to 10 concurrent downloads
)

func cacheImage(imageURL, filename string) (string, error) {
	cacheDir := "image_cache"
	cachedImagePath := filepath.Join(cacheDir, filename)

	// Check if the image is already cached
	if _, err := os.Stat(cachedImagePath); err == nil {
		return cachedImagePath, nil
	}

	// Ensure only one goroutine caches the same image
	cachingImagesMu.Lock()
	if _, exists := cachingImages[imageURL]; !exists {
		cachingImages[imageURL] = &sync.Mutex{}
	}
	mu := cachingImages[imageURL]
	cachingImagesMu.Unlock()

	mu.Lock()
	defer mu.Unlock()

	// Double-check if the image was cached while waiting
	if _, err := os.Stat(cachedImagePath); err == nil {
		return cachedImagePath, nil
	}

	cachingSemaphore <- struct{}{}        // Acquire a token
	defer func() { <-cachingSemaphore }() // Release the token

	// Download the image
	resp, err := http.Get(imageURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	// Read the image data into a byte slice
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	// Detect the content type
	contentType := http.DetectContentType(data)

	// If content type is HTML, skip caching
	if strings.HasPrefix(contentType, "text/html") {
		return "", fmt.Errorf("URL returned HTML content instead of an image: %s", imageURL)
	}

	// Handle SVG files directly
	if contentType == "image/svg+xml" {
		// Ensure the cache directory exists
		if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
			os.Mkdir(cacheDir, os.ModePerm)
		}

		// Save the SVG file as-is
		err = os.WriteFile(cachedImagePath, data, 0644)
		if err != nil {
			return "", err
		}

		// Clean up mutex
		cachingImagesMu.Lock()
		delete(cachingImages, imageURL)
		cachingImagesMu.Unlock()

		return cachedImagePath, nil
	}

	// Decode the image based on the content type
	var img image.Image
	switch contentType {
	case "image/jpeg":
		img, err = jpeg.Decode(bytes.NewReader(data))
	case "image/png":
		img, err = png.Decode(bytes.NewReader(data))
	case "image/gif":
		img, err = gif.Decode(bytes.NewReader(data))
	case "image/webp":
		img, err = webp.Decode(bytes.NewReader(data))
	case "image/bmp":
		img, err = bmp.Decode(bytes.NewReader(data))
	case "image/tiff":
		img, err = tiff.Decode(bytes.NewReader(data))
	default:
		return "", fmt.Errorf("unsupported image type: %s", contentType)
	}

	if err != nil {
		return "", fmt.Errorf("failed to decode image: %v", err)
	}

	// Ensure the cache directory exists
	if _, err := os.Stat(cacheDir); os.IsNotExist(err) {
		os.Mkdir(cacheDir, os.ModePerm)
	}

	// Open the cached file for writing
	outFile, err := os.Create(cachedImagePath)
	if err != nil {
		return "", err
	}
	defer outFile.Close()

	// Encode the image to WebP and save
	options := &webp.Options{Lossless: false, Quality: 80}
	err = webp.Encode(outFile, img, options)
	if err != nil {
		return "", err
	}

	// Clean up mutex
	cachingImagesMu.Lock()
	delete(cachingImages, imageURL)
	cachingImagesMu.Unlock()

	return cachedImagePath, nil
}

func handleCachedImages(w http.ResponseWriter, r *http.Request) {
	imageName := filepath.Base(r.URL.Path)
	cacheDir := "image_cache"
	cachedImagePath := filepath.Join(cacheDir, imageName)

	if _, err := os.Stat(cachedImagePath); os.IsNotExist(err) {
		// Serve placeholder image with no-store headers
		placeholderPath := "static/images/placeholder.webp"
		placeholderContentType := "image/webp"

		// You can also check for SVG placeholder if needed
		if strings.HasSuffix(imageName, ".svg") {
			placeholderPath = "static/images/placeholder.svg"
			placeholderContentType = "image/svg+xml"
		}

		w.Header().Set("Content-Type", placeholderContentType)
		w.Header().Set("Cache-Control", "no-store, must-revalidate")
		w.Header().Set("Pragma", "no-cache")
		w.Header().Set("Expires", "0")
		http.ServeFile(w, r, placeholderPath)
		return
	}

	// Determine the content type based on the file extension
	extension := strings.ToLower(filepath.Ext(cachedImagePath))
	var contentType string
	switch extension {
	case ".svg":
		contentType = "image/svg+xml"
	case ".jpg", ".jpeg":
		contentType = "image/jpeg"
	case ".png":
		contentType = "image/png"
	case ".gif":
		contentType = "image/gif"
	case ".webp":
		contentType = "image/webp"
	default:
		// Default to binary stream if unknown
		contentType = "application/octet-stream"
	}

	w.Header().Set("Content-Type", contentType)
	w.Header().Set("Cache-Control", "public, max-age=31536000") // Cache the image for 1 year
	http.ServeFile(w, r, cachedImagePath)
}

func handleImageStatus(w http.ResponseWriter, r *http.Request) {
	imageIDs := r.URL.Query().Get("image_ids")
	ids := strings.Split(imageIDs, ",")

	statusMap := make(map[string]string)

	cacheDir := "image_cache"

	printDebug("Received image status request for IDs: %v", ids)
	printDebug("Status map: %v", statusMap)

	for _, id := range ids {
		// Check for different possible extensions
		extensions := []string{".webp", ".svg"}
		var cachedImagePath string
		var found bool

		for _, ext := range extensions {
			filename := id + ext
			path := filepath.Join(cacheDir, filename)
			if _, err := os.Stat(path); err == nil {
				cachedImagePath = "/image_cache/" + filename
				found = true
				break
			}
		}

		if found {
			statusMap[id] = cachedImagePath
		} else {
			// Image is not ready
			statusMap[id] = ""
		}
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(statusMap)
}