munch-ease-backend/meals/repository.py

import datetime
from typing import AsyncIterator, List, Optional

from ingredients import (
    Ingredient,
    delete_ingredients_by_meal_id,
    find_ingredients_by_meal_id,
    insert_ingredient,
)
from meals.models import Meal, MealRecipe
from persons.models import Person
from persons.repository import get_by_ids as persons_get_by_ids
from recipes.models import Recipe
from recipes.repository import load_recipe_ingredients, row_to_recipe

from .roles import ROLE_CHEF, ROLE_CLEANUP, ROLE_CONSUMER


async def create(conn):
    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS Meal (
            id INTEGER PRIMARY KEY,
            suggested_date DATETIME,
            consumed_date DATETIME DEFAULT NULL,
            deleted_date DATETIME DEFAULT NULL,
            purchase_date DATETIME DEFAULT NULL
        );"""
    )

    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS MealParticipant (
            meal_id INTEGER,
            person_id INTEGER,
            role TEXT,
            FOREIGN KEY(meal_id) REFERENCES Meal(id),
            FOREIGN KEY(person_id) REFERENCES Person(id)
        );"""
    )
    # Useful indexes
    await conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_meal_participants_meal_role ON MealParticipant(meal_id, role);"
    )

    await conn.execute(
        """
        CREATE TABLE IF NOT EXISTS MealRecipe (
            meal_id INTEGER,
            recipe_id INTEGER,
            servings REAL,
            FOREIGN KEY(meal_id) REFERENCES Meal(id),
            FOREIGN KEY(recipe_id) REFERENCES Recipe(id)
        );"""
    )

    # Index for faster lookup of recipes by meal
    await conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_meal_recipes_meal_id ON MealRecipe(meal_id);"
    )


async def insert_meal_participant(conn, meal_id: int, person_id: int, role: str):
    await conn.execute(
        """
        INSERT INTO MealParticipant (meal_id, person_id, role)
        VALUES (?, ?, ?)
    """,
        (meal_id, person_id, role),
    )


async def sync_meal_participants(conn, meal_id: int, participants: List[Person], role: str):
    await conn.execute(
        """
        DELETE FROM MealParticipant
        WHERE meal_id = ? AND role = ?
    """,
        (meal_id, role),
    )

    for person in participants:
        await insert_meal_participant(conn, meal_id, person.id, role)


async def insert_meal_recipe(conn, r: MealRecipe):
    if r.meal_id < 0:
        raise ValueError("Meal must be inserted before meal recipe")

    if r.recipe_id < 0 and r.recipe:
        r.recipe_id = r.recipe.id

    if r.recipe_id < 0:
        raise ValueError("Recipe must be inserted before meal")

    await conn.execute(
        """
        INSERT INTO MealRecipe (meal_id, recipe_id, servings)
        VALUES (?, ?, ?)
    """,
        (r.meal_id, r.recipe_id, r.servings),
    )


async def insert_meal(conn, meal: Meal):
    async with conn.execute(
        """
        INSERT INTO Meal (suggested_date)
        VALUES (?)
    """,
        (meal.suggested_date.isoformat(),),
    ) as cursor:
        meal.id = cursor.lastrowid

    await sync_meal_participants(conn, meal.id, meal.chefs, ROLE_CHEF)
    await sync_meal_participants(conn, meal.id, meal.cleanup, ROLE_CLEANUP)
    await sync_meal_participants(conn, meal.id, meal.consumers, ROLE_CONSUMER)

    for meal_recipe in meal.recipes:
        meal_recipe.meal_id = meal.id

        await insert_meal_recipe(conn, meal_recipe)

    await sync_extra_ingredients(conn, meal.id, meal.extra_ingredients)


async def find_meal_by_id(conn, meal_id: int) -> Optional[Meal]:
    async with conn.execute(
        f"""
        SELECT {",".join(Meal.KEYS)} FROM Meal
        WHERE id = ?
        LIMIT 1
    """,
        (meal_id,),
    ) as cursor:
        async for row in cursor:
            meal = Meal(**{k: v for k, v in zip(Meal.KEYS, row)})

            await load_participants(conn, meal)
            await load_recipes(conn, meal)
            await load_extra_ingredients(conn, meal)
            return meal
    return None


async def find_meal_by_id_scoped(conn, meal_id: int, household_id: int) -> Optional[Meal]:
    async with conn.execute(
        f"""
        SELECT {",".join(Meal.KEYS)} FROM Meal
        WHERE id = ? AND household_id = ?
        LIMIT 1
    """,
        (meal_id, household_id),
    ) as cursor:
        async for row in cursor:
            meal = Meal(**{k: v for k, v in zip(Meal.KEYS, row)})

            await load_participants(conn, meal)
            await load_recipes(conn, meal)
            await load_extra_ingredients(conn, meal)
            return meal
    return None


async def find_upcoming_meals_by_date_range(
    conn, start: datetime.datetime, end: datetime.datetime
) -> AsyncIterator[Meal]:
    async with conn.execute(
        f"""
        SELECT {",".join(Meal.KEYS)} FROM Meal
        WHERE suggested_date >= ? AND suggested_date <= ? AND consumed_date IS NULL AND deleted_date IS NULL
    """,
        (start, end),
    ) as cursor:
        async for row in cursor:
            yield Meal(**{k: v for k, v in zip(Meal.KEYS, row)})


async def find_upcoming_meals_by_date_range_scoped(
    conn, start: datetime.datetime, end: datetime.datetime, household_id: int
) -> AsyncIterator[Meal]:
    async with conn.execute(
        f"""
        SELECT {",".join(Meal.KEYS)} FROM Meal
        WHERE suggested_date >= ? AND suggested_date <= ? AND consumed_date IS NULL AND deleted_date IS NULL AND household_id = ?
    """,
        (start, end, household_id),
    ) as cursor:
        async for row in cursor:
            yield Meal(**{k: v for k, v in zip(Meal.KEYS, row)})


async def load_participants(conn, meal: Meal) -> None:
    # Fetch all participant links
    links: list[tuple[int, str]] = []
    async with conn.execute(
        """
        SELECT person_id, role FROM MealParticipant
        WHERE meal_id = ?
    """,
        (meal.id,),
    ) as cursor:
        async for row in cursor:
            links.append((int(row[0]), str(row[1])))

    if not links:
        return

    # Bulk load persons by id
    unique_ids = sorted({pid for pid, _ in links})
    people = await persons_get_by_ids(conn, unique_ids)

    for pid, role in links:
        person = people.get(pid)
        if role == ROLE_CHEF:
            if person:
                meal.chefs.append(person)
        elif role == ROLE_CLEANUP:
            if person:
                meal.cleanup.append(person)
        elif role == ROLE_CONSUMER:
            if person:
                meal.consumers.append(person)
        else:
            raise Exception(f"Unknown role: {role}")


async def bulk_load_participants(conn, meals: List[Meal]) -> None:
    """Populate participants for many meals in one query to avoid N+1.

    For each meal, fills meal.chefs, meal.cleanup, meal.consumers using a bulk
    lookup of MealParticipant rows and a single persons.get_by_ids fetch.
    """
    if not meals:
        return

    meal_ids = [m.id for m in meals]
    placeholders = ",".join(["?"] * len(meal_ids))

    # Collect (meal_id -> [(person_id, role), ...]) and dedupe person IDs
    links_by_meal: dict[int, list[tuple[int, str]]] = {mid: [] for mid in meal_ids}
    person_ids: set[int] = set()
    async with conn.execute(
        f"""
        SELECT meal_id, person_id, role
        FROM MealParticipant
        WHERE meal_id IN ({placeholders})
        """,
        meal_ids,
    ) as cursor:
        async for row in cursor:
            mid, pid, role = int(row[0]), int(row[1]), str(row[2])
            links_by_meal.setdefault(mid, []).append((pid, role))
            person_ids.add(pid)

    if not person_ids:
        return

    # Bulk load persons once
    people = await persons_get_by_ids(conn, sorted(person_ids))

    # Assign per meal
    by_id = {m.id: m for m in meals}
    for mid, links in links_by_meal.items():
        meal = by_id.get(mid)
        if not meal:
            continue
        # Reset roles to avoid duplicates
        meal.chefs = []
        meal.cleanup = []
        meal.consumers = []
        for pid, role in links:
            person = people.get(pid)
            if not person:
                continue
            if role == ROLE_CHEF:
                meal.chefs.append(person)
            elif role == ROLE_CLEANUP:
                meal.cleanup.append(person)
            elif role == ROLE_CONSUMER:
                meal.consumers.append(person)
            else:
                raise Exception(f"Unknown role: {role}")


async def load_recipes(conn, meal: Meal) -> None:
    async with conn.execute(
        f"""
        SELECT {",".join(Recipe.KEYS)}, MealRecipe.servings as requested_servings
        FROM Recipe
        JOIN MealRecipe ON MealRecipe.recipe_id = Recipe.id
        WHERE MealRecipe.meal_id = ?
    """,
        (meal.id,),
    ) as cursor:
        async for row in cursor:
            recipe = row_to_recipe(list(zip(Recipe.KEYS, row[:-1])))
            await load_recipe_ingredients(conn, recipe)

            meal.recipes.append(
                MealRecipe(meal_id=meal.id, recipe_id=recipe.id, servings=row[-1], recipe=recipe)
            )


async def load_extra_ingredients(conn, meal: Meal) -> None:
    async for ingredient in find_ingredients_by_meal_id(conn, meal.id):
        meal.extra_ingredients.append(ingredient)


async def delete_meal(conn, meal_id: int) -> None:
    await conn.execute(
        """
        UPDATE Meal
        SET deleted_date = ?
        WHERE id = ?
    """,
        (datetime.datetime.now().astimezone().isoformat(), meal_id),
    )


async def sync_extra_ingredients(conn, meal_id: int, ingredients: List[Ingredient]) -> None:
    await delete_ingredients_by_meal_id(conn, meal_id)

    for ingredient in ingredients:
        ingredient.meal_id = meal_id
        ingredient.recipe_id = None

        await insert_ingredient(conn, ingredient)


async def sync_meal_recipes(conn, meal_id: int, recipes: List[MealRecipe]) -> None:
    await conn.execute(
        """
        DELETE FROM MealRecipe
        WHERE meal_id = ?
    """,
        (meal_id,),
    )

    for meal_recipe in recipes:
        if meal_recipe.meal_id >= 0 and meal_recipe.meal_id != meal_id:
            raise ValueError("Already associated with another meal")

        meal_recipe.meal_id = meal_id
        await insert_meal_recipe(conn, meal_recipe)


async def update_meal(conn, meal: Meal) -> None:
    await conn.execute(
        """
        UPDATE Meal
        SET suggested_date = ?
        WHERE id = ?
    """,
        (meal.suggested_date.isoformat(), meal.id),
    )

    await sync_meal_participants(conn, meal.id, meal.chefs, ROLE_CHEF)
    await sync_meal_participants(conn, meal.id, meal.cleanup, ROLE_CLEANUP)
    await sync_meal_participants(conn, meal.id, meal.consumers, ROLE_CONSUMER)

    await sync_extra_ingredients(conn, meal.id, meal.extra_ingredients)
    await sync_meal_recipes(conn, meal.id, meal.recipes)


async def mark_consumed(conn, meal: Meal, date: datetime.datetime) -> None:
    meal.consumed_date = date

    await conn.execute(
        """
        UPDATE Meal
        SET consumed_date = ?
        WHERE id = ?
    """,
        (date.isoformat(), meal.id),
    )


async def mark_purchased(conn, meal: Meal) -> Meal:
    meal.purchase_date = datetime.datetime.now().astimezone()

    await conn.execute(
        """
        UPDATE Meal
        SET purchase_date = ?
        WHERE id = ?
    """,
        (meal.purchase_date.isoformat(), meal.id),
    )

    return meal