Epicure/apps/web/lib/ai/features/estimate-nutrition.ts

import { generateObject } from "ai";
import { z } from "zod";
import { resolveModel, type AiConfig } from "../factory";
import { estimateGrams } from "@/lib/ingredient-grams";
import { lookupUsdaNutrients, type UsdaNutrientsPer100g } from "@/lib/usda";

// Only used for its inferred type below (estimateNutrition builds this
// shape by hand from USDA + AI sub-totals, not via generateObject anymore).
// eslint-disable-next-line @typescript-eslint/no-unused-vars
const NutritionSchema = z.object({
  perServing: z.object({
    calories: z.number().describe("Total calories per serving (kcal)"),
    proteinG: z.number().describe("Protein in grams per serving"),
    carbsG: z.number().describe("Total carbohydrates in grams per serving"),
    fatG: z.number().describe("Total fat in grams per serving"),
    fiberG: z.number().describe("Dietary fiber in grams per serving"),
    sodiumMg: z.number().describe("Sodium in milligrams per serving"),
  }),
});

export type NutritionEstimate = z.infer<typeof NutritionSchema>;

type Totals = { calories: number; proteinG: number; carbsG: number; fatG: number; fiberG: number; sodiumMg: number };

const ZERO_TOTALS: Totals = { calories: 0, proteinG: 0, carbsG: 0, fatG: 0, fiberG: 0, sodiumMg: 0 };

function addTotals(a: Totals, b: Totals): Totals {
  return {
    calories: a.calories + b.calories,
    proteinG: a.proteinG + b.proteinG,
    carbsG: a.carbsG + b.carbsG,
    fatG: a.fatG + b.fatG,
    fiberG: a.fiberG + b.fiberG,
    sodiumMg: a.sodiumMg + b.sodiumMg,
  };
}

function scalePer100g(per100g: UsdaNutrientsPer100g, grams: number): Totals {
  const factor = grams / 100;
  return {
    calories: per100g.calories * factor,
    proteinG: per100g.proteinG * factor,
    carbsG: per100g.carbsG * factor,
    fatG: per100g.fatG * factor,
    fiberG: per100g.fiberG * factor,
    sodiumMg: per100g.sodiumMg * factor,
  };
}

const IngredientTotalSchema = z.object({
  calories: z.number().describe("Total calories (kcal) contributed by these ingredients, as used"),
  proteinG: z.number().describe("Total protein in grams"),
  carbsG: z.number().describe("Total carbohydrates in grams"),
  fatG: z.number().describe("Total fat in grams"),
  fiberG: z.number().describe("Total dietary fiber in grams"),
  sodiumMg: z.number().describe("Total sodium in milligrams"),
});

/** AI fallback for whichever ingredients didn't get a USDA match — asked
 * for the total contribution of just this subset, at the quantities as
 * used (not per serving), so it composes directly with the USDA-derived
 * totals via addTotals rather than needing to be reconciled against a
 * separate whole-recipe estimate. */
async function estimateIngredientsTotal(
  ingredients: Array<{ rawName: string; quantity?: string | number | null; unit?: string | null }>,
  config?: AiConfig
): Promise<Totals> {
  const model = resolveModel(config);
  const ingredientList = ingredients
    .map((i) => {
      const parts: string[] = [];
      if (i.quantity) parts.push(String(i.quantity));
      if (i.unit) parts.push(i.unit);
      parts.push(i.rawName);
      return `- ${parts.join(" ")}`;
    })
    .join("\n");

  const { object } = await generateObject({
    model,
    schema: IngredientTotalSchema,
    system:
      "You are an expert nutritionist with deep knowledge of food composition. Estimate the total nutritional contribution of the given ingredient list at the quantities specified, not per serving. When quantities are not specified, use typical amounts. Provide realistic estimates based on standard nutritional databases.",
    prompt: `Estimate the total nutrition contributed by these ingredients, as used:\n\n${ingredientList}`,
  });

  return object;
}

/**
 * Per-ingredient USDA FoodData Central lookup (converting quantity+unit to
 * grams first) summed across matches, with AI estimation as the fallback
 * for whichever ingredients couldn't be converted to grams or had no
 * usable USDA match — including every ingredient, transparently, when
 * USDA_API_KEY isn't configured at all. Divides the combined total by
 * baseServings for the final per-serving figures callers expect.
 */
export async function estimateNutrition(
  recipe: {
    title: string;
    baseServings: number;
    ingredients: Array<{
      rawName: string;
      quantity?: string | number | null;
      unit?: string | null;
    }>;
  },
  config?: AiConfig
): Promise<NutritionEstimate> {
  const usdaAttempts = await Promise.all(
    recipe.ingredients.map(async (ingredient) => {
      const grams = estimateGrams(ingredient.quantity, ingredient.unit);
      if (grams == null) return { ingredient, totals: null };
      const per100g = await lookupUsdaNutrients(ingredient.rawName);
      return { ingredient, totals: per100g ? scalePer100g(per100g, grams) : null };
    })
  );

  const usdaTotal = usdaAttempts.reduce(
    (sum, attempt) => (attempt.totals ? addTotals(sum, attempt.totals) : sum),
    ZERO_TOTALS
  );
  const remaining = usdaAttempts.filter((a) => !a.totals).map((a) => a.ingredient);

  const remainingTotal = remaining.length > 0
    ? await estimateIngredientsTotal(remaining, config)
    : ZERO_TOTALS;

  const total = addTotals(usdaTotal, remainingTotal);
  const servings = recipe.baseServings > 0 ? recipe.baseServings : 1;

  return {
    perServing: {
      calories: total.calories / servings,
      proteinG: total.proteinG / servings,
      carbsG: total.carbsG / servings,
      fatG: total.fatG / servings,
      fiberG: total.fiberG / servings,
      sodiumMg: total.sodiumMg / servings,
    },
  };
}