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

package recognition

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

	"github.com/go-chi/chi/v5"

	"github.com/food-ai/backend/internal/adapters/ai"
	"github.com/food-ai/backend/internal/domain/dish"
	"github.com/food-ai/backend/internal/domain/product"
	"github.com/food-ai/backend/internal/infra/locale"
	"github.com/food-ai/backend/internal/infra/middleware"
)

// DishRepository is the subset of dish.Repository used by workers and the handler.
type DishRepository interface {
	FindOrCreate(ctx context.Context, name string) (string, bool, error)
	FindOrCreateRecipe(ctx context.Context, dishID string, calories, proteinG, fatG, carbsG float64) (string, bool, error)
	UpsertTranslation(ctx context.Context, dishID, lang, name string) error
	GetTranslation(ctx context.Context, dishID, lang string) (string, bool, error)
	AddRecipe(ctx context.Context, dishID string, req dish.CreateRequest) (string, error)
}

// ProductRepository is the subset of product.Repository used by this handler.
type ProductRepository interface {
	FuzzyMatch(ctx context.Context, name string) (*product.Product, error)
	Upsert(ctx context.Context, catalogProduct *product.Product) (*product.Product, error)
}

// Recognizer is the AI provider interface for image-based food recognition.
type Recognizer interface {
	RecognizeReceipt(ctx context.Context, imageBase64, mimeType, lang string) (*ai.ReceiptResult, error)
	RecognizeProducts(ctx context.Context, imageBase64, mimeType, lang string) ([]ai.RecognizedItem, error)
	RecognizeDish(ctx context.Context, imageBase64, mimeType, lang string) (*ai.DishResult, error)
	ClassifyIngredient(ctx context.Context, name string) (*ai.IngredientClassification, error)
	GenerateRecipeForDish(ctx context.Context, dishName string) (*ai.Recipe, error)
	TranslateDishName(ctx context.Context, name string) (map[string]string, error)
}

// KafkaPublisher publishes job IDs to a Kafka topic.
type KafkaPublisher interface {
	Publish(ctx context.Context, topic, message string) error
}

// Handler handles POST /ai/* recognition endpoints.
type Handler struct {
	recognizer  Recognizer
	productRepo ProductRepository
	jobRepo     JobRepository
	kafkaProducer KafkaPublisher
	sseBroker   *SSEBroker
}

// NewHandler creates a new Handler with async dish recognition support.
func NewHandler(
	recognizer Recognizer,
	productRepo ProductRepository,
	jobRepo JobRepository,
	kafkaProducer KafkaPublisher,
	sseBroker *SSEBroker,
) *Handler {
	return &Handler{
		recognizer:  recognizer,
		productRepo: productRepo,
		jobRepo:     jobRepo,
		kafkaProducer: kafkaProducer,
		sseBroker:   sseBroker,
	}
}

// ---------------------------------------------------------------------------
// 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"`
}

// recognizeDishRequest is the body for POST /ai/recognize-dish.
type recognizeDishRequest struct {
	ImageBase64    string  `json:"image_base64"`
	MimeType       string  `json:"mime_type"`
	TargetDate     *string `json:"target_date"`
	TargetMealType *string `json:"target_meal_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 []ai.UnrecognizedItem `json:"unrecognized"`
}

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

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

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

	lang := locale.FromContext(request.Context())
	result, recognizeError := handler.recognizer.RecognizeReceipt(request.Context(), req.ImageBase64, req.MimeType, lang)
	if recognizeError != nil {
		slog.Error("recognize receipt", "err", recognizeError)
		writeErrorJSON(responseWriter, http.StatusServiceUnavailable, "recognition failed, please try again")
		return
	}

	enriched := handler.enrichItems(request.Context(), result.Items)
	writeJSON(responseWriter, http.StatusOK, ReceiptResponse{
		Items:        enriched,
		Unrecognized: result.Unrecognized,
	})
}

// RecognizeProducts handles POST /ai/recognize-products.
// Body: {"images": [{"image_base64": "...", "mime_type": "image/jpeg"}, ...]}
func (handler *Handler) RecognizeProducts(responseWriter http.ResponseWriter, request *http.Request) {
	var req imagesRequest
	if decodeError := json.NewDecoder(request.Body).Decode(&req); decodeError != nil || len(req.Images) == 0 {
		writeErrorJSON(responseWriter, 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
	}

	lang := locale.FromContext(request.Context())
	allItems := make([][]ai.RecognizedItem, len(req.Images))
	var wg sync.WaitGroup
	for i, img := range req.Images {
		wg.Add(1)
		go func(index int, imageReq imageRequest) {
			defer wg.Done()
			items, recognizeError := handler.recognizer.RecognizeProducts(request.Context(), imageReq.ImageBase64, imageReq.MimeType, lang)
			if recognizeError != nil {
				slog.Warn("recognize products from image", "index", index, "err", recognizeError)
				return
			}
			allItems[index] = items
		}(i, img)
	}
	wg.Wait()

	merged := MergeAndDeduplicate(allItems)
	enriched := handler.enrichItems(request.Context(), merged)
	writeJSON(responseWriter, http.StatusOK, map[string]any{"items": enriched})
}

// RecognizeDish handles POST /ai/recognize-dish (async).
// Enqueues the image for AI processing and returns 202 Accepted with a job_id.
// Body: {"image_base64": "...", "mime_type": "image/jpeg", "target_date": "2006-01-02", "target_meal_type": "lunch"}
func (handler *Handler) RecognizeDish(responseWriter http.ResponseWriter, request *http.Request) {
	var req recognizeDishRequest
	if decodeError := json.NewDecoder(request.Body).Decode(&req); decodeError != nil || req.ImageBase64 == "" {
		writeErrorJSON(responseWriter, http.StatusBadRequest, "image_base64 is required")
		return
	}

	userID := middleware.UserIDFromCtx(request.Context())
	userPlan := middleware.UserPlanFromCtx(request.Context())
	lang := locale.FromContext(request.Context())

	job := &Job{
		UserID:         userID,
		UserPlan:       userPlan,
		ImageBase64:    req.ImageBase64,
		MimeType:       req.MimeType,
		Lang:           lang,
		TargetDate:     req.TargetDate,
		TargetMealType: req.TargetMealType,
	}
	if insertError := handler.jobRepo.InsertJob(request.Context(), job); insertError != nil {
		slog.Error("insert recognition job", "err", insertError)
		writeErrorJSON(responseWriter, http.StatusInternalServerError, "failed to create job")
		return
	}

	position, positionError := handler.jobRepo.QueuePosition(request.Context(), userPlan, job.CreatedAt)
	if positionError != nil {
		position = 0
	}

	topic := TopicFree
	if userPlan == "paid" {
		topic = TopicPaid
	}
	if publishError := handler.kafkaProducer.Publish(request.Context(), topic, job.ID); publishError != nil {
		slog.Error("publish recognition job", "job_id", job.ID, "err", publishError)
		writeErrorJSON(responseWriter, http.StatusInternalServerError, "failed to enqueue job")
		return
	}

	estimatedSeconds := (position + 1) * 6
	writeJSON(responseWriter, http.StatusAccepted, map[string]any{
		"job_id":            job.ID,
		"queue_position":    position,
		"estimated_seconds": estimatedSeconds,
	})
}

// ListTodayJobs handles GET /ai/jobs — returns today's unlinked jobs for the current user.
func (handler *Handler) ListTodayJobs(responseWriter http.ResponseWriter, request *http.Request) {
	userID := middleware.UserIDFromCtx(request.Context())

	summaries, listError := handler.jobRepo.ListTodayUnlinked(request.Context(), userID)
	if listError != nil {
		slog.Error("list today unlinked jobs", "err", listError)
		writeErrorJSON(responseWriter, http.StatusInternalServerError, "failed to list jobs")
		return
	}

	// Return an empty array instead of null when there are no results.
	if summaries == nil {
		summaries = []*JobSummary{}
	}
	writeJSON(responseWriter, http.StatusOK, summaries)
}

// ListAllJobs handles GET /ai/jobs/history — returns all recognition jobs for the current user.
func (handler *Handler) ListAllJobs(responseWriter http.ResponseWriter, request *http.Request) {
	userID := middleware.UserIDFromCtx(request.Context())

	summaries, listError := handler.jobRepo.ListAll(request.Context(), userID)
	if listError != nil {
		slog.Error("list all jobs", "err", listError)
		writeErrorJSON(responseWriter, http.StatusInternalServerError, "failed to list jobs")
		return
	}

	if summaries == nil {
		summaries = []*JobSummary{}
	}
	writeJSON(responseWriter, http.StatusOK, summaries)
}

// GetJobStream handles GET /ai/jobs/{id}/stream — SSE endpoint for job updates.
func (handler *Handler) GetJobStream(responseWriter http.ResponseWriter, request *http.Request) {
	handler.sseBroker.ServeSSE(responseWriter, request)
}

// GetJob handles GET /ai/jobs/{id} — fetches a job result (for app re-open after backgrounding).
func (handler *Handler) GetJob(responseWriter http.ResponseWriter, request *http.Request) {
	jobID := chi.URLParam(request, "id")
	userID := middleware.UserIDFromCtx(request.Context())

	job, fetchError := handler.jobRepo.GetJobByID(request.Context(), jobID)
	if fetchError != nil {
		writeErrorJSON(responseWriter, http.StatusNotFound, "job not found")
		return
	}
	if job.UserID != userID {
		writeErrorJSON(responseWriter, http.StatusForbidden, "forbidden")
		return
	}
	writeJSON(responseWriter, http.StatusOK, job)
}

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

// enrichItems matches each recognized item against the product catalog.
// Items without a match trigger a classification call and upsert into the DB.
func (handler *Handler) enrichItems(ctx context.Context, items []ai.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
		}

		catalogProduct, matchError := handler.productRepo.FuzzyMatch(ctx, item.Name)
		if matchError != nil {
			slog.Warn("fuzzy match product", "name", item.Name, "err", matchError)
		}

		if catalogProduct != nil {
			id := catalogProduct.ID
			enriched.MappingID = &id
			if catalogProduct.DefaultUnit != nil {
				enriched.Unit = *catalogProduct.DefaultUnit
			}
			if catalogProduct.StorageDays != nil {
				enriched.StorageDays = *catalogProduct.StorageDays
			}
			if catalogProduct.Category != nil {
				enriched.Category = *catalogProduct.Category
			}
		} else {
			classification, classifyError := handler.recognizer.ClassifyIngredient(ctx, item.Name)
			if classifyError != nil {
				slog.Warn("classify unknown product", "name", item.Name, "err", classifyError)
			} else {
				saved := handler.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 classification into the product catalog.
func (handler *Handler) saveClassification(ctx context.Context, classification *ai.IngredientClassification) *product.Product {
	if classification == nil || classification.CanonicalName == "" {
		return nil
	}

	catalogProduct := &product.Product{
		CanonicalName:   classification.CanonicalName,
		Category:        strPtr(classification.Category),
		DefaultUnit:     strPtr(classification.DefaultUnit),
		CaloriesPer100g: classification.CaloriesPer100g,
		ProteinPer100g:  classification.ProteinPer100g,
		FatPer100g:      classification.FatPer100g,
		CarbsPer100g:    classification.CarbsPer100g,
		StorageDays:     intPtr(classification.StorageDays),
	}

	saved, upsertError := handler.productRepo.Upsert(ctx, catalogProduct)
	if upsertError != nil {
		slog.Warn("upsert classified product", "name", classification.CanonicalName, "err", upsertError)
		return nil
	}

	return saved
}

// MergeAndDeduplicate combines results from multiple images.
// Items sharing the same name (case-insensitive) have their quantities summed.
func MergeAndDeduplicate(batches [][]ai.RecognizedItem) []ai.RecognizedItem {
	seen := make(map[string]*ai.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
				if item.Confidence > existing.Confidence {
					existing.Confidence = item.Confidence
				}
			} else {
				seen[key] = item
				order = append(order, key)
			}
		}
	}

	result := make([]ai.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(responseWriter http.ResponseWriter, status int, msg string) {
	responseWriter.Header().Set("Content-Type", "application/json")
	responseWriter.WriteHeader(status)
	_ = json.NewEncoder(responseWriter).Encode(errorResponse{Error: msg})
}

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