Search/cache-images.go

package main

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

	"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{}, 30) // Limit to concurrent downloads

	invalidImageIDs   = make(map[string]struct{})
	invalidImageIDsMu sync.Mutex
)

func cacheImage(imageURL, filename, imageID string) (string, bool, error) {
	cacheDir := "image_cache"
	cachedImagePath := filepath.Join(cacheDir, filename)
	tempImagePath := cachedImagePath + ".tmp"

	// Check if the image is already cached
	if _, err := os.Stat(cachedImagePath); err == nil {
		return cachedImagePath, true, 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, true, nil
	}

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

	// Create a custom http.Client that skips SSL certificate verification
	client := &http.Client{
		Timeout: 15 * time.Second,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		},
	}

	// Download the image using the custom client
	resp, err := client.Get(imageURL)
	if err != nil {
		recordInvalidImageID(imageID)
		return "", false, err
	}
	defer resp.Body.Close()

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

	// Check if the response is actually an image
	contentType := http.DetectContentType(data)
	if !strings.HasPrefix(contentType, "image/") {
		recordInvalidImageID(imageID)
		return "", false, fmt.Errorf("URL did not return 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 to the temp path
		err = os.WriteFile(tempImagePath, data, 0644)
		if err != nil {
			recordInvalidImageID(imageID)
			return "", false, err
		}

		// Atomically rename the temp file to the final cached image path
		err = os.Rename(tempImagePath, cachedImagePath)
		if err != nil {
			recordInvalidImageID(imageID)
			return "", false, err
		}

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

		return cachedImagePath, true, 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:
		recordInvalidImageID(imageID)
		return "", false, fmt.Errorf("unsupported image type: %s", contentType)
	}

	if err != nil {
		recordInvalidImageID(imageID)
		return "", false, 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 temp file for writing
	outFile, err := os.Create(tempImagePath)
	if err != nil {
		recordInvalidImageID(imageID)
		return "", false, err
	}

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

	// Atomically rename the temp file to the final cached image path
	err = os.Rename(tempImagePath, cachedImagePath)
	if err != nil {
		recordInvalidImageID(imageID)
		return "", false, err
	}

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

	return cachedImagePath, true, 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) {
		printDebug("Cached image not found: %s, serving missing.svg", cachedImagePath)
		// Serve missing image
		missingImagePath := filepath.Join("static", "images", "missing.svg")
		w.Header().Set("Content-Type", "image/svg+xml")
		http.ServeFile(w, r, missingImagePath)
		return
	} else if err != nil {
		printWarn("Error checking image file: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		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)

	invalidImageIDsMu.Lock()
	defer invalidImageIDsMu.Unlock()

	for _, id := range ids {
		// Check if the image ID is in the invalidImageIDs map
		if _, invalid := invalidImageIDs[id]; invalid {
			// Image is invalid, set status to "missing"
			statusMap[id] = "/static/images/missing.svg"
			continue
		}

		// 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 yet
			statusMap[id] = ""
		}
	}

	printDebug("Status map: %v", statusMap)

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

func recordInvalidImageID(imageID string) {
	invalidImageIDsMu.Lock()
	defer invalidImageIDsMu.Unlock()
	invalidImageIDs[imageID] = struct{}{}
	printDebug("Recorded invalid image ID: %s", imageID)
}

func filterValidImages(imageResults []ImageSearchResult) []ImageSearchResult {
	invalidImageIDsMu.Lock()
	defer invalidImageIDsMu.Unlock()

	var filteredResults []ImageSearchResult
	for _, img := range imageResults {
		if _, invalid := invalidImageIDs[img.ID]; !invalid {
			filteredResults = append(filteredResults, img)
		} else {
			printDebug("Filtering out invalid image ID: %s", img.ID)
		}
	}
	return filteredResults
}

func removeImageResultFromCache(query string, page int, safe bool, lang string, imageID string) {
	cacheKey := CacheKey{
		Query: query,
		Page:  page,
		Safe:  safe,
		Lang:  lang,
		Type:  "image",
	}

	rc := resultsCache

	rc.mu.Lock()
	defer rc.mu.Unlock()

	keyStr := rc.keyToString(cacheKey)
	item, exists := rc.results[keyStr]
	if !exists {
		return
	}

	// Filter out the image with the given ID
	var newResults []SearchResult
	for _, r := range item.Results {
		if imgResult, ok := r.(ImageSearchResult); ok {
			if imgResult.ID != imageID {
				newResults = append(newResults, r)
			} else {
				printDebug("Removing invalid image ID from cache: %s", imageID)
			}
		} else {
			newResults = append(newResults, r)
		}
	}

	// Update or delete the cache entry
	if len(newResults) > 0 {
		rc.results[keyStr] = CachedItem{
			Results:    newResults,
			StoredTime: item.StoredTime,
		}
	} else {
		delete(rc.results, keyStr)
	}
}