bambuddy/backend/app/services/print_batch.py

"""Batch order planning: per-plate targets, progress, and staged dispatch (#342).

A batch stores *intent* in :class:`PrintBatchPlate` rows — "this order wants 3
of plate 2" — while its queue items record what was actually dispatched.
Everything here derives one from the other.

The distinction matters for exactly one reason, and it is the reason the
feature exists: a failed or cancelled run does not count towards the target, so
``remaining`` goes back up and the order still says it owes a print. A design
that only counted the items it created could not tell "the user cancelled this
deliberately" apart from "this one burned and needs reprinting".

Batches created before targets existed have no plate rows. They still report
progress — the plate breakdown is derived from their queue items and every
target simply equals the number of items dispatched, so ``remaining`` is zero
and the dispatch endpoint has nothing to do. ``has_targets`` tells callers
which kind of batch they are looking at.
"""

import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from backend.app.models.print_batch import PrintBatch, PrintBatchPlate
from backend.app.models.print_log import PrintLogEntry
from backend.app.models.print_queue import PrintQueueItem, PrintQueueVariant
from backend.app.services.queue_item_printer import drop_printer_resolved
from backend.app.services.queue_position import next_queue_position

logger = logging.getLogger(__name__)

# Statuses that consume a unit of the target. "printing" counts because the
# run is in flight — re-dispatching it would double-print. "failed",
# "cancelled" and "skipped" deliberately do not.
CONSUMING_STATUSES = ("pending", "printing", "completed")

# Queue statuses the roll-up has a counter for. Anything else is ignored rather
# than crashing the page — the queue's status vocabulary is allowed to grow
# without this module having to be updated in lockstep.
COUNTED_STATUSES = ("pending", "printing", "completed", "failed", "cancelled", "skipped")

# Columns copied onto a clone when dispatching more of a plate. This is the
# print *configuration* the user already chose and the API already validated —
# copying the row is what keeps a second dispatch identical to the first
# without re-serialising twenty fields through a template blob that would drift
# from the model the first time someone adds a column.
CLONED_SETTING_COLUMNS = (
    "printer_id",
    "target_model",
    "target_location",
    "required_filament_types",
    "archive_id",
    "library_file_id",
    "project_id",
    "batch_id",
    "ams_mapping",
    # A virtual printer job's slicer pick for its archive (#2700): each copy's
    # print makes or joins an archive too.
    "slicer_ams_mapping",
    "filament_overrides",
    "plate_id",
    "print_time_seconds",
    "gcode_injection",
    "nozzle_mapping",
    "nozzle_rack_choice",
    "require_previous_success",
    "auto_off_after",
    "manual_start",
    "bed_levelling",
    "flow_cali",
    "vibration_cali",
    "layer_inspect",
    "timelapse",
    "use_ams",
    "nozzle_offset_cali",
    "confirm_outcome",
    "preheat_override",
    "preheat_chamber_target_override",
    "skip_filament_check",
    # The scheduler bills and budget-checks from it; dispatch_batch checks the
    # dispatcher may use it.
    "cost_center_id",
)

CLONED_VARIANT_COLUMNS = (
    "position",
    "library_file_id",
    "target_model",
    "plate_id",
    "ams_mapping",
    "nozzle_mapping",
    "nozzle_rack_choice",
    "filament_overrides",
    "required_filament_types",
    "print_time_seconds",
)


class BatchDispatchError(Exception):
    """Raised when more runs are owed but nothing can be cloned to produce them."""


@dataclass
class PlateProgress:
    """Per-plate roll-up for one batch."""

    plate_id: int | None
    plate_name: str | None
    quantity_target: int
    sort_order: int = 0
    pending: int = 0
    printing: int = 0
    completed: int = 0
    failed: int = 0
    cancelled: int = 0
    skipped: int = 0
    # Actual material + energy cost of this plate's finished runs. None when no
    # run has produced a cost yet — reported as "unknown", never as zero.
    actual_cost: float | None = None
    filament_used_grams: float | None = None
    print_time_seconds: int = 0
    # Whether any queue item for this plate still exists, in any status.
    # Dispatch clones one to inherit the print configuration, so a plate with
    # none cannot be queued however many runs it still owes.
    has_source: bool = False

    @property
    def dispatched(self) -> int:
        return self.pending + self.printing + self.completed

    @property
    def remaining(self) -> int:
        return max(0, self.quantity_target - self.dispatched)

    @property
    def can_dispatch(self) -> bool:
        """True when this plate owes runs *and* something can produce them."""
        return self.remaining > 0 and self.has_source

    @property
    def cost_per_run(self) -> float | None:
        """Observed mean cost of this plate's completed runs, or None.

        Deliberately measured rather than estimated from the file: the file's
        estimate ignores what the run actually consumed, and a plate that has
        never completed has no honest number to show.
        """
        if self.completed <= 0 or self.actual_cost is None:
            return None
        return self.actual_cost / self.completed

    @property
    def estimated_remaining_cost(self) -> float | None:
        per_run = self.cost_per_run
        if per_run is None:
            return None
        return per_run * self.remaining


@dataclass
class BatchProgress:
    """Whole-order roll-up, plus the per-plate breakdown it was derived from."""

    plates: list[PlateProgress] = field(default_factory=list)
    has_targets: bool = False

    def _sum(self, attr: str) -> int:
        return sum(getattr(p, attr) for p in self.plates)

    @property
    def pending(self) -> int:
        return self._sum("pending")

    @property
    def printing(self) -> int:
        return self._sum("printing")

    @property
    def completed(self) -> int:
        return self._sum("completed")

    @property
    def failed(self) -> int:
        return self._sum("failed")

    @property
    def cancelled(self) -> int:
        return self._sum("cancelled")

    @property
    def skipped(self) -> int:
        return self._sum("skipped")

    @property
    def target(self) -> int:
        return self._sum("quantity_target")

    @property
    def remaining(self) -> int:
        return self._sum("remaining")

    @property
    def dispatchable_remaining(self) -> int:
        """Of the runs still owed, how many can actually be queued.

        Lower than ``remaining`` when a plate's last queue item was deleted:
        the order still owes the run, but nothing survives to clone its
        printer target, AMS mapping and print options from.
        """
        return sum(p.remaining for p in self.plates if p.has_source)

    @property
    def actual_cost(self) -> float | None:
        costs = [p.actual_cost for p in self.plates if p.actual_cost is not None]
        return sum(costs) if costs else None

    @property
    def estimated_remaining_cost(self) -> float | None:
        estimates = [p.estimated_remaining_cost for p in self.plates if p.estimated_remaining_cost is not None]
        return sum(estimates) if estimates else None

    @property
    def filament_used_grams(self) -> float | None:
        grams = [p.filament_used_grams for p in self.plates if p.filament_used_grams is not None]
        return sum(grams) if grams else None

    @property
    def print_time_seconds(self) -> int:
        return self._sum("print_time_seconds")

    @property
    def is_fulfilled(self) -> bool:
        """True when every target is met and nothing is still in flight.

        A zero total target is never "fulfilled". Without that guard a legacy
        batch whose items were all cancelled one by one would report itself
        completed — its derived target counts only pending/printing/completed
        items, so cancelling the lot leaves a target of zero that trivially
        satisfies ``remaining == 0``.
        """
        return self.target > 0 and self.remaining == 0 and self.pending == 0 and self.printing == 0


async def load_progress(db: AsyncSession, batch: PrintBatch) -> BatchProgress:
    """Build the per-plate progress roll-up for *batch*.

    Two queries plus one for costs, regardless of how many plates the order
    has — this runs once per batch in the list endpoint.
    """
    plate_rows = (await db.execute(select(PrintBatchPlate).where(PrintBatchPlate.batch_id == batch.id))).scalars().all()

    # (plate_id, status) -> count, plus the time/weight actually recorded.
    item_rows = (
        await db.execute(
            select(
                PrintQueueItem.plate_id,
                PrintQueueItem.status,
                func.count(PrintQueueItem.id),
                func.sum(PrintQueueItem.print_time_seconds),
            )
            .where(PrintQueueItem.batch_id == batch.id)
            .group_by(PrintQueueItem.plate_id, PrintQueueItem.status)
        )
    ).all()

    # Per-run actuals, attributed through the queue item that produced them.
    # PrintLogEntry is the authoritative per-run record (#1378) and is already
    # scoped to the printed plate (#2614), so a multi-plate order gets each
    # plate's own cost rather than the whole file's.
    cost_rows = (
        await db.execute(
            select(
                PrintQueueItem.plate_id,
                func.sum(
                    func.coalesce(PrintLogEntry.cost, 0.0)
                    + func.coalesce(PrintLogEntry.energy_cost, 0.0)
                    + func.coalesce(PrintLogEntry.wear_cost, 0.0)
                ),
                func.sum(PrintLogEntry.filament_used_grams),
            )
            .select_from(PrintLogEntry)
            .join(PrintQueueItem, PrintLogEntry.queue_item_id == PrintQueueItem.id)
            .where(PrintQueueItem.batch_id == batch.id)
            .group_by(PrintQueueItem.plate_id)
        )
    ).all()
    costs = {row[0]: (row[1], row[2]) for row in cost_rows}

    progress = BatchProgress(has_targets=bool(plate_rows))
    by_plate: dict[int | None, PlateProgress] = {}

    for row in plate_rows:
        by_plate[row.plate_id] = PlateProgress(
            plate_id=row.plate_id,
            plate_name=row.plate_name,
            quantity_target=row.quantity_target,
            sort_order=row.sort_order,
        )

    for plate_id, status, count, time_sum in item_rows:
        plate = by_plate.get(plate_id)
        if plate is None:
            # A queue item for a plate the order has no target row for: either
            # a legacy batch, or an item grouped in by hand after the fact.
            # Its own dispatched count becomes its target so it reads as
            # complete rather than as owing work nobody asked for.
            plate = PlateProgress(plate_id=plate_id, plate_name=None, quantity_target=0, sort_order=plate_id or 0)
            by_plate[plate_id] = plate
            if status in CONSUMING_STATUSES:
                plate.quantity_target += count
        elif not progress.has_targets and status in CONSUMING_STATUSES:
            plate.quantity_target += count
        # Any surviving row is a clone source, whatever its status — a
        # cancelled or failed run still carries the configuration a re-queue
        # needs. Set before the status filter below so a status this module
        # has no counter for still marks the plate dispatchable.
        plate.has_source = True
        if status in COUNTED_STATUSES:
            setattr(plate, status, getattr(plate, status) + count)
        else:
            logger.debug("Batch %s: ignoring queue item status %r in progress roll-up", batch.id, status)
        plate.print_time_seconds += int(time_sum or 0)

    for plate_id, (cost_sum, gram_sum) in costs.items():
        plate = by_plate.get(plate_id)
        if plate is None:
            continue
        plate.actual_cost = float(cost_sum) if cost_sum else None
        plate.filament_used_grams = float(gram_sum) if gram_sum else None

    progress.plates = sorted(by_plate.values(), key=lambda p: (p.sort_order, p.plate_id or 0))
    return progress


async def refresh_batch_status(db: AsyncSession, batch: PrintBatch) -> bool:
    """Flip an ``active`` batch to ``completed`` once its targets are met.

    Returns True when the status changed. A ``cancelled`` batch is never
    resurrected, and a ``completed`` batch drops back to ``active`` if its
    targets grow — raising a target on a finished order reopens it rather than
    leaving a "completed" order that still owes prints.
    """
    progress = await load_progress(db, batch)

    if batch.status == "cancelled":
        return False

    if batch.status == "active" and progress.is_fulfilled:
        batch.status = "completed"
        batch.completed_at = datetime.now(timezone.utc)
        logger.info("Batch %s fulfilled — marked completed", batch.id)
        return True

    # A grouping whose every item was cancelled one at a time is finished, but
    # nothing was produced, so "completed" would be a lie and `is_fulfilled`
    # rightly refuses it (its derived target is zero). Left alone it would sit
    # on "active" forever. Cancelled is what it is, and matches what the
    # batch-level Cancel action would have set had it been used.
    #
    # Deliberately not applied to orders: an order states its intent
    # independently of its runs, so cancelling every run still leaves it owing
    # work and offering to re-queue it. A grouping has no such statement — it
    # was only ever the sum of its items.
    if batch.status == "active" and not progress.has_targets and progress.completed == 0:
        settled = progress.pending == 0 and progress.printing == 0
        if settled and progress.cancelled > 0 and progress.failed == 0 and progress.skipped == 0:
            batch.status = "cancelled"
            logger.info("Batch %s had every item cancelled — marked cancelled", batch.id)
            return True
        # Likewise finished, but with real failures among its jobs and nothing
        # produced: neither "completed" nor "cancelled" says that. Not while a
        # job is skipped: resuming the printer puts it back in the queue.
        if settled and progress.failed > 0 and progress.skipped == 0:
            batch.status = "failed"
            logger.info("Batch %s ended with failures and nothing printed — marked failed", batch.id)
            return True

    # A failed grouping that has open work again (a job added to it) is active.
    if batch.status == "failed" and (progress.pending > 0 or progress.printing > 0):
        batch.status = "active"
        logger.info("Batch %s has open work again — reopened", batch.id)
        return True

    if batch.status == "completed" and not progress.is_fulfilled:
        batch.status = "active"
        batch.completed_at = None
        logger.info("Batch %s reopened — targets no longer met", batch.id)
        return True

    return False


async def backfill_batch_statuses(db: AsyncSession) -> int:
    """Close out ``active`` batches that finished before the status existed.

    ``completed`` only became reachable with #342. Every batch created since
    the feature shipped in April 2026 is therefore still marked ``active``,
    however long ago its last run finished — so without this pass the Batches
    tab opens on months of accumulated history.

    Runs on every startup rather than once behind a marker: it is cheap (only
    batches with nothing in flight are even considered), it is idempotent, and
    repeating it also closes out any order whose last run landed while the
    process was down.

    Returns the number of batches whose status changed.
    """
    candidates = (
        (
            await db.execute(
                select(PrintBatch)
                .where(PrintBatch.status == "active")
                # Anything still queued or printing is by definition unfinished,
                # and re-deriving its progress would change nothing.
                .where(
                    ~select(PrintQueueItem.id)
                    .where(PrintQueueItem.batch_id == PrintBatch.id)
                    .where(PrintQueueItem.status.in_(("pending", "printing")))
                    .exists()
                )
            )
        )
        .scalars()
        .all()
    )

    changed = 0
    for batch in candidates:
        if await refresh_batch_status(db, batch):
            changed += 1

    if changed:
        await db.commit()
        logger.info("Marked %d finished batch(es) as completed at startup (#342)", changed)
    return changed


async def refresh_batches_for_items(db: AsyncSession, queue_item_ids: list[int]) -> None:
    """Re-evaluate every batch owning one of *queue_item_ids*, once each.

    For the paths that end items without a print completing (a cancel, a stop,
    the stranded-item sweep, a library delete): a grouping whose last item was
    cancelled, or an order whose last run was closed, would otherwise stay
    active until the next restart. Changes are left for the caller to commit.
    """
    if not queue_item_ids:
        return
    await db.flush()
    batch_ids = (
        (
            await db.execute(
                select(PrintQueueItem.batch_id)
                .where(PrintQueueItem.id.in_(queue_item_ids))
                .where(PrintQueueItem.batch_id.is_not(None))
                .distinct()
            )
        )
        .scalars()
        .all()
    )
    for batch_id in batch_ids:
        batch = await db.get(PrintBatch, batch_id)
        if batch is not None:
            await refresh_batch_status(db, batch)


async def refresh_batches_after(queue_item_ids: list[int]) -> None:
    """:func:`refresh_batches_for_items` in its own session, after the caller's commit.

    Best-effort, as on the print-completion path: an order's status must not
    be the reason a cancel, a stop or a scheduler step fails or rolls back.
    """
    if not queue_item_ids:
        return
    from backend.app.core import database

    try:
        async with database.async_session() as db:
            await refresh_batches_for_items(db, queue_item_ids)
            await db.commit()
    except Exception as e:
        logger.warning("Failed to refresh the batch status of queue items %s: %s", queue_item_ids, e)


async def refresh_batch_status_for_item(db: AsyncSession, queue_item_id: int) -> None:
    """Re-evaluate the batch owning *queue_item_id*, if it has one.

    Called from the print-completion path so a finished order reports itself
    complete the moment its last run lands, rather than whenever someone next
    opens the page.
    """
    batch_id = (
        await db.execute(select(PrintQueueItem.batch_id).where(PrintQueueItem.id == queue_item_id))
    ).scalar_one_or_none()
    if batch_id is None:
        return
    batch = (await db.execute(select(PrintBatch).where(PrintBatch.id == batch_id))).scalar_one_or_none()
    if batch is None:
        return
    await refresh_batch_status(db, batch)


def _clone_queue_item(source: PrintQueueItem, *, position: int, created_by_id: int | None) -> PrintQueueItem:
    """Copy *source*'s print configuration into a fresh pending item.

    Lifecycle state (status, timestamps, retry counters, scheduler flags) is
    deliberately not copied — the clone is a new run, not a resurrection.

    ``scheduled_time`` is dropped too: dispatching more of a plate is a
    "queue this now" action, and replaying the original's scheduled time would
    either fire immediately (it is in the past) or silently park the new run
    until a moment the user chose for a different print.

    ``cleanup_library_after_dispatch`` is forced off. It comes from the
    Printers-page direct-print flow and from a queue virtual printer's "Remove
    from File Manager once it prints", and deletes the library row at dispatch
    — replaying that on a clone would delete the source file out from under the
    rest of the order.
    """
    clone = PrintQueueItem(
        status="pending",
        position=position,
        created_by_id=created_by_id if created_by_id is not None else source.created_by_id,
        cleanup_library_after_dispatch=False,
    )
    for column in CLONED_SETTING_COLUMNS:
        setattr(clone, column, getattr(source, column))
    # A printer the scheduler picked for a model-based job (or one with
    # printer alternatives) was that run's pick, not the user's: each copy
    # waits for any printer of the model again.
    if source.printer_id is not None and (source.target_model or source.variants):
        clone.printer_id = None
        drop_printer_resolved(clone)
    return clone


def _plate_label(plate: PlateProgress) -> str:
    """How a plate is named in an error the user reads.

    Prefers the plate name the order stored, because that is what the Batches
    tab shows; falls back to the plate number for orders created before names
    were recorded.
    """
    if plate.plate_name:
        return plate.plate_name
    return f"Plate {plate.plate_id if plate.plate_id is not None else 1}"


async def dispatch_remaining(
    db: AsyncSession,
    batch: PrintBatch,
    *,
    plate_id: int | None = None,
    only_plate: bool = False,
    limit: int | None = None,
    created_by_id: int | None = None,
) -> list[PrintQueueItem]:
    """Create queue items for the runs *batch* still owes.

    ``only_plate`` restricts the dispatch to the single plate named by
    ``plate_id`` (which may legitimately be ``None`` for a single-plate file);
    otherwise every plate with work outstanding is dispatched in plate order.
    ``limit`` caps the total number of items created across all plates.

    Raises :class:`BatchDispatchError` when nothing at all can be produced
    because no plate that owes runs has an item to clone — the order can
    describe work whose every run has since been deleted, and there is no
    configuration to copy in that case. A plate in that state is skipped
    rather than aborting the whole order: one unrecoverable plate must not
    block the plates that are still perfectly dispatchable.
    """
    progress = await load_progress(db, batch)
    if not progress.has_targets:
        return []

    targets = [p for p in progress.plates if p.remaining > 0]
    if only_plate:
        targets = [p for p in targets if p.plate_id == plate_id]

    created: list[PrintQueueItem] = []
    stranded: list[PlateProgress] = []

    for plate in targets:
        if limit is not None and len(created) >= limit:
            break

        source = (
            await db.execute(
                select(PrintQueueItem)
                .options(selectinload(PrintQueueItem.variants))
                .where(PrintQueueItem.batch_id == batch.id)
                .where(PrintQueueItem.plate_id == plate.plate_id)
                .order_by(PrintQueueItem.id.desc())
                .limit(1)
            )
        ).scalar_one_or_none()

        if source is None:
            stranded.append(plate)
            continue

        wanted = plate.remaining
        if limit is not None:
            wanted = min(wanted, limit - len(created))

        # Clones for this plate are appended to the end of the queue (#3200),
        # after whatever is already pending, in plate order.
        position = await next_queue_position(db)

        for _ in range(wanted):
            clone = _clone_queue_item(source, position=position, created_by_id=created_by_id)
            position += 1
            db.add(clone)
            await db.flush()
            for variant in source.variants:
                cloned_variant = PrintQueueVariant(queue_item_id=clone.id)
                for column in CLONED_VARIANT_COLUMNS:
                    setattr(cloned_variant, column, getattr(variant, column))
                db.add(cloned_variant)
            created.append(clone)

    if not created and stranded:
        names = ", ".join(_plate_label(p) for p in stranded)
        raise BatchDispatchError(
            f"{names} {'have' if len(stranded) > 1 else 'has'} no queued or finished run to copy settings from. "
            "Queue the plate once from the file, then dispatch the rest from here."
        )

    if created:
        # Dispatching more work can only ever un-fulfil an order, but run the
        # check anyway so a reopened batch flips back from completed.
        await db.flush()
        await refresh_batch_status(db, batch)

    if stranded:
        logger.info("Batch %s: skipped %d plate(s) with no item to clone", batch.id, len(stranded))
    logger.info("Dispatched %d item(s) for batch %s", len(created), batch.id)
    return created