bambuddy/backend/app/services/filament_requirements.py

"""Parse per-slot filament requirements out of a 3MF file.

The scheduler used to own this logic (`PrintScheduler._get_filament_requirements`)
because it ran during dispatch decisions. Extracted here so the VP queue-mode
write path can use the same parser to populate `filament_overrides` /
`required_filament_types` at upload time (#1188 — Bambuddy was creating queue
items with no filament fields, which made the scheduler fall through to
model-only matching and dispatch onto whatever printer happened to be free
regardless of loaded colour).

The shape returned here matches the `filament_overrides` JSON shape the
scheduler validates against, minus the `force_color_match` flag — callers
add that themselves based on their own setting.
"""

from __future__ import annotations

import logging
import xml.etree.ElementTree as ET
import zipfile
from pathlib import Path

from backend.app.utils.threemf_tools import extract_nozzle_mapping_from_3mf

logger = logging.getLogger(__name__)


def extract_filament_requirements(file_path: Path, plate_id: int | None = None) -> list[dict]:
    """Parse `[{slot_id, type, color, tray_info_idx, used_grams, nozzle_id?}]` from a 3MF.

    Args:
        file_path: Path to the 3MF.
        plate_id: When set, only return filaments used on that plate. When
            None, return every filament with `used_g > 0` across the file.

    Returns:
        Sorted list (by `slot_id`) of filament dicts. Empty list when the
        3MF is unreadable, missing `Metadata/slice_info.config`, or has no
        filaments matching the plate filter — callers treat that as "no
        requirements" rather than an error so a malformed 3MF doesn't break
        the upload path.
    """
    if not file_path.exists():
        return []

    filaments: list[dict] = []
    try:
        with zipfile.ZipFile(file_path, "r") as zf:
            if "Metadata/slice_info.config" not in zf.namelist():
                return []

            content = zf.read("Metadata/slice_info.config").decode()
            root = ET.fromstring(content)  # noqa: S314  # nosec B314

            if plate_id is not None:
                for plate_elem in root.findall("./plate"):
                    plate_index = None
                    for meta in plate_elem.findall("metadata"):
                        if meta.get("key") == "index":
                            try:
                                plate_index = int(meta.get("value", "0"))
                            except ValueError:
                                pass
                            break
                    if plate_index == plate_id:
                        _collect_filaments(plate_elem, filaments)
                        break
            else:
                # Modern BambuStudio format wraps filaments inside <plate> elements.
                # When no plate filter is requested, collect from every plate and
                # deduplicate by slot_id (first occurrence wins after sort).
                plate_elems = root.findall("./plate")
                if plate_elems:
                    for plate_elem in plate_elems:
                        _collect_filaments(plate_elem, filaments)
                    # Deduplicate: same slot_id can appear on multiple plates.
                    # Keep the entry with the highest used_grams; ties go to the
                    # first plate (stable after sort + dict insertion order).
                    seen: dict[int, dict] = {}
                    for f in filaments:
                        sid = f["slot_id"]
                        if sid not in seen or f["used_grams"] > seen[sid]["used_grams"]:
                            seen[sid] = f
                    filaments = list(seen.values())
                else:
                    # Older / non-plate-wrapped format: filaments are direct children of root.
                    _collect_filaments(root, filaments)

            filaments.sort(key=lambda x: x["slot_id"])

            # Dual-nozzle printers (H2D / X2D) — annotate which extruder each
            # slot is fed into. Empty mapping for single-nozzle printers, in
            # which case we just don't add the key.
            # Same plate the filaments above were collected from: a multi-plate
            # file can assign one slot to different extruders per plate, and
            # annotating slot 2 with plate 3's nozzle is worse than not
            # annotating it.
            nozzle_mapping = extract_nozzle_mapping_from_3mf(zf, plate_id=plate_id)
            if nozzle_mapping:
                for filament in filaments:
                    filament["nozzle_id"] = nozzle_mapping.get(filament["slot_id"])
    except Exception as e:
        logger.warning("Failed to parse filament requirements from %s: %s", file_path, e)
        return []

    return filaments


def overrides_for_plate(
    overrides: list[dict],
    file_path: Path | None,
    plate_id: int | None,
) -> list[dict]:
    """Drop the filament overrides whose slots this plate never prints.

    Queueing several plates of one 3MF builds a single override list out of every
    selected plate's filaments and hands that same list to each plate's item. A
    ``force_color_match`` entry blocks dispatch until the printer has that exact
    colour loaded, so a single-colour plate ended up waiting on every colour in
    the batch (#2551). Each item may only demand what its own plate consumes.

    Overrides are kept as-is when the plate's slots cannot be established (whole
    file selected, source gone, unreadable 3MF, malformed entry): an item that
    waits on a colour it does not need is visible and fixable, whereas one that
    silently loses a forced colour can dispatch the print in the wrong filament.
    """
    if not overrides or plate_id is None or file_path is None or not file_path.exists():
        return overrides

    plate_slots = {f["slot_id"] for f in extract_filament_requirements(file_path, plate_id)}
    if not plate_slots:
        logger.warning(
            "Cannot read the filaments of plate %s in %s; keeping all %d filament override(s)",
            plate_id,
            file_path.name,
            len(overrides),
        )
        return overrides

    narrowed = []
    for override in overrides:
        try:
            slot_id = int(override["slot_id"])
        except (KeyError, TypeError, ValueError):
            narrowed.append(override)
            continue
        if slot_id in plate_slots:
            narrowed.append(override)

    if len(narrowed) != len(overrides):
        logger.info(
            "Plate %s: kept %d of %d filament override(s) — the rest belong to other plates",
            plate_id,
            len(narrowed),
            len(overrides),
        )
    return narrowed


def _collect_filaments(parent: ET.Element, into: list[dict]) -> None:
    """Walk every `./filament` child under `parent` and append normalised
    entries to `into`. Skips filaments with `used_g <= 0` (slot present in
    the slicer config but not consumed by this plate)."""
    for filament_elem in parent.findall("./filament"):
        filament_id = filament_elem.get("id")
        if not filament_id:
            continue
        try:
            used_grams = float(filament_elem.get("used_g", "0"))
        except (ValueError, TypeError):
            continue
        if used_grams <= 0:
            continue
        try:
            slot_id = int(filament_id)
        except (ValueError, TypeError):
            continue
        into.append(
            {
                "slot_id": slot_id,
                "type": filament_elem.get("type", ""),
                "color": filament_elem.get("color", ""),
                "tray_info_idx": filament_elem.get("tray_info_idx", ""),
                "used_grams": round(used_grams, 1),
            }
        )