food-ai/backend/internal/recognition/handler.go

package recognition

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"strings"
	"sync"

	"github.com/food-ai/backend/internal/gemini"
	"github.com/food-ai/backend/internal/ingredient"
	"github.com/food-ai/backend/internal/middleware"
)

// ingredientRepo is the subset of ingredient.Repository used by this handler.
type ingredientRepo interface {
	FuzzyMatch(ctx context.Context, name string) (*ingredient.IngredientMapping, error)
	Upsert(ctx context.Context, m *ingredient.IngredientMapping) (*ingredient.IngredientMapping, error)
}

// Handler handles POST /ai/* recognition endpoints.
type Handler struct {
	gemini         *gemini.Client
	ingredientRepo ingredientRepo
}

// NewHandler creates a new Handler.
func NewHandler(geminiClient *gemini.Client, repo ingredientRepo) *Handler {
	return &Handler{gemini: geminiClient, ingredientRepo: repo}
}

// ---------------------------------------------------------------------------
// Request / Response types
// ---------------------------------------------------------------------------

// imageRequest is the common request body containing a single base64-encoded image.
type imageRequest struct {
	ImageBase64 string `json:"image_base64"`
	MimeType    string `json:"mime_type"`
}

// imagesRequest is the request body for multi-image endpoints.
type imagesRequest struct {
	Images []imageRequest `json:"images"`
}

// EnrichedItem is a recognized food item enriched with ingredient_mappings data.
type EnrichedItem struct {
	Name        string  `json:"name"`
	Quantity    float64 `json:"quantity"`
	Unit        string  `json:"unit"`
	Category    string  `json:"category"`
	Confidence  float64 `json:"confidence"`
	MappingID   *string `json:"mapping_id"`
	StorageDays int     `json:"storage_days"`
}

// ReceiptResponse is the response for POST /ai/recognize-receipt.
type ReceiptResponse struct {
	Items        []EnrichedItem          `json:"items"`
	Unrecognized []gemini.UnrecognizedItem `json:"unrecognized"`
}

// DishResponse is the response for POST /ai/recognize-dish.
type DishResponse = gemini.DishResult

// ---------------------------------------------------------------------------
// Handlers
// ---------------------------------------------------------------------------

// RecognizeReceipt handles POST /ai/recognize-receipt.
// Body: {"image_base64": "...", "mime_type": "image/jpeg"}
func (h *Handler) RecognizeReceipt(w http.ResponseWriter, r *http.Request) {
	userID := middleware.UserIDFromCtx(r.Context())
	_ = userID // logged for tracing

	var req imageRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.ImageBase64 == "" {
		writeErrorJSON(w, http.StatusBadRequest, "image_base64 is required")
		return
	}

	result, err := h.gemini.RecognizeReceipt(r.Context(), req.ImageBase64, req.MimeType)
	if err != nil {
		slog.Error("recognize receipt", "err", err)
		writeErrorJSON(w, http.StatusServiceUnavailable, "recognition failed, please try again")
		return
	}

	enriched := h.enrichItems(r.Context(), result.Items)
	writeJSON(w, http.StatusOK, ReceiptResponse{
		Items:        enriched,
		Unrecognized: result.Unrecognized,
	})
}

// RecognizeProducts handles POST /ai/recognize-products.
// Body: {"images": [{"image_base64": "...", "mime_type": "image/jpeg"}, ...]}
func (h *Handler) RecognizeProducts(w http.ResponseWriter, r *http.Request) {
	var req imagesRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || len(req.Images) == 0 {
		writeErrorJSON(w, http.StatusBadRequest, "at least one image is required")
		return
	}
	if len(req.Images) > 3 {
		req.Images = req.Images[:3] // cap at 3 photos as per spec
	}

	// Process each image in parallel.
	allItems := make([][]gemini.RecognizedItem, len(req.Images))
	var wg sync.WaitGroup
	for i, img := range req.Images {
		wg.Add(1)
		go func(i int, img imageRequest) {
			defer wg.Done()
			items, err := h.gemini.RecognizeProducts(r.Context(), img.ImageBase64, img.MimeType)
			if err != nil {
				slog.Warn("recognize products from image", "index", i, "err", err)
				return
			}
			allItems[i] = items
		}(i, img)
	}
	wg.Wait()

	merged := mergeAndDeduplicate(allItems)
	enriched := h.enrichItems(r.Context(), merged)
	writeJSON(w, http.StatusOK, map[string]any{"items": enriched})
}

// RecognizeDish handles POST /ai/recognize-dish.
// Body: {"image_base64": "...", "mime_type": "image/jpeg"}
func (h *Handler) RecognizeDish(w http.ResponseWriter, r *http.Request) {
	var req imageRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.ImageBase64 == "" {
		writeErrorJSON(w, http.StatusBadRequest, "image_base64 is required")
		return
	}

	result, err := h.gemini.RecognizeDish(r.Context(), req.ImageBase64, req.MimeType)
	if err != nil {
		slog.Error("recognize dish", "err", err)
		writeErrorJSON(w, http.StatusServiceUnavailable, "recognition failed, please try again")
		return
	}

	writeJSON(w, http.StatusOK, result)
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

// enrichItems matches each recognized item against ingredient_mappings.
// Items without a match trigger a Gemini classification call and upsert into the DB.
func (h *Handler) enrichItems(ctx context.Context, items []gemini.RecognizedItem) []EnrichedItem {
	result := make([]EnrichedItem, 0, len(items))
	for _, item := range items {
		enriched := EnrichedItem{
			Name:        item.Name,
			Quantity:    item.Quantity,
			Unit:        item.Unit,
			Category:    item.Category,
			Confidence:  item.Confidence,
			StorageDays: 7, // sensible default
		}

		mapping, err := h.ingredientRepo.FuzzyMatch(ctx, item.Name)
		if err != nil {
			slog.Warn("fuzzy match ingredient", "name", item.Name, "err", err)
		}

		if mapping != nil {
			// Found existing mapping — use its canonical data.
			id := mapping.ID
			enriched.MappingID = &id
			if mapping.DefaultUnit != nil {
				enriched.Unit = *mapping.DefaultUnit
			}
			if mapping.StorageDays != nil {
				enriched.StorageDays = *mapping.StorageDays
			}
			if mapping.Category != nil {
				enriched.Category = *mapping.Category
			}
		} else {
			// No mapping — ask AI to classify and save for future reuse.
			classification, err := h.gemini.ClassifyIngredient(ctx, item.Name)
			if err != nil {
				slog.Warn("classify unknown ingredient", "name", item.Name, "err", err)
			} else {
				saved := h.saveClassification(ctx, classification)
				if saved != nil {
					id := saved.ID
					enriched.MappingID = &id
				}
				enriched.Category = classification.Category
				enriched.Unit = classification.DefaultUnit
				enriched.StorageDays = classification.StorageDays
			}
		}
		result = append(result, enriched)
	}
	return result
}

// saveClassification upserts an AI-produced ingredient classification into the DB.
func (h *Handler) saveClassification(ctx context.Context, c *gemini.IngredientClassification) *ingredient.IngredientMapping {
	if c == nil || c.CanonicalName == "" {
		return nil
	}

	aliasesJSON, err := json.Marshal(c.Aliases)
	if err != nil {
		return nil
	}

	m := &ingredient.IngredientMapping{
		CanonicalName:   c.CanonicalName,
		CanonicalNameRu: &c.CanonicalNameRu,
		Category:        strPtr(c.Category),
		DefaultUnit:     strPtr(c.DefaultUnit),
		CaloriesPer100g: c.CaloriesPer100g,
		ProteinPer100g:  c.ProteinPer100g,
		FatPer100g:      c.FatPer100g,
		CarbsPer100g:    c.CarbsPer100g,
		StorageDays:     intPtr(c.StorageDays),
		Aliases:         aliasesJSON,
	}

	saved, err := h.ingredientRepo.Upsert(ctx, m)
	if err != nil {
		slog.Warn("upsert classified ingredient", "name", c.CanonicalName, "err", err)
		return nil
	}
	return saved
}

// mergeAndDeduplicate combines results from multiple images.
// Items sharing the same name (case-insensitive) have their quantities summed.
func mergeAndDeduplicate(batches [][]gemini.RecognizedItem) []gemini.RecognizedItem {
	seen := make(map[string]*gemini.RecognizedItem)
	var order []string

	for _, batch := range batches {
		for i := range batch {
			item := &batch[i]
			key := normalizeName(item.Name)
			if existing, ok := seen[key]; ok {
				existing.Quantity += item.Quantity
				// Keep the higher confidence estimate.
				if item.Confidence > existing.Confidence {
					existing.Confidence = item.Confidence
				}
			} else {
				seen[key] = item
				order = append(order, key)
			}
		}
	}

	result := make([]gemini.RecognizedItem, 0, len(order))
	for _, key := range order {
		result = append(result, *seen[key])
	}
	return result
}

func normalizeName(s string) string {
	return strings.ToLower(strings.TrimSpace(s))
}

func strPtr(s string) *string {
	if s == "" {
		return nil
	}
	return &s
}

func intPtr(n int) *int {
	return &n
}

// ---------------------------------------------------------------------------
// HTTP helpers
// ---------------------------------------------------------------------------

type errorResponse struct {
	Error string `json:"error"`
}

func writeErrorJSON(w http.ResponseWriter, status int, msg string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(errorResponse{Error: msg})
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}