grotti/grotti/send_queue.py

"""Hintergrund-Queue für Senden und Starten an Moonraker."""

from __future__ import annotations

import threading
import time
import uuid
from collections import deque
from queue import Empty, Queue

from grotti import jobs
from grotti.db import connect
from grotti.errors import GrottiError
from grotti.events_bus import publish as publish_live


class SendQueue:
    def __init__(self, database_url, cfg, workers=4):
        self.database_url = database_url
        self.cfg = cfg
        self._workers_n = max(1, int(workers))
        self._lock = threading.Lock()
        self._pending = Queue()
        self._jobs = {}
        self._order = deque(maxlen=60)
        self._threads = []
        self._stop = threading.Event()

    def start(self):
        if self._threads:
            return
        self._stop.clear()
        for index in range(self._workers_n):
            thread = threading.Thread(
                target=self._worker,
                name="grotti-send-queue-%s" % (index + 1),
                daemon=True,
            )
            thread.start()
            self._threads.append(thread)

    def stop(self):
        self._stop.set()
        for _ in self._threads:
            try:
                self._pending.put_nowait(None)
            except Exception:
                pass

    def enqueue(
        self,
        conn,
        file_id,
        host,
        confirm,
        override,
        spool_id,
        ace_map=None,
        bed_leveling=False,
    ):
        jobs_out = self.enqueue_many(
            conn,
            file_id,
            [host],
            confirm,
            override,
            spool_id,
            {host: ace_map or []} if ace_map else {},
            bed_leveling=bed_leveling,
        )
        return jobs_out[0]

    def enqueue_many(
        self,
        conn,
        file_id,
        hosts,
        confirm,
        override,
        spool_id,
        ace_maps=None,
        bed_leveling=False,
    ):
        if not confirm:
            raise GrottiError("Start braucht eine Bestätigung.", "bestaetigung")
        hosts = [str(host or "").strip() for host in (hosts or []) if str(host or "").strip()]
        if not hosts:
            raise GrottiError("Mindestens einen Drucker wählen.", "drucker")
        seen = set()
        unique = []
        for host in hosts:
            if host in seen:
                continue
            seen.add(host)
            unique.append(host)
        ace_maps = ace_maps if isinstance(ace_maps, dict) else {}
        # Eine Batch-Option für alle Drucker im Multi-Send (nicht pro Host).
        want_level = bool(bed_leveling)
        created = []
        for host in unique:
            # Mapping zuerst, damit Materialprüfung die ACE-Spulen sieht.
            mapping = ace_maps.get(host) or []
            printer, row, chosen = jobs.prepare_send(
                conn, file_id, host, True, override, spool_id, ace_map=mapping
            )
            if not mapping:
                mapping = ace_maps.get(printer["name"]) or []
            job_id = uuid.uuid4().hex[:12]
            job = {
                "id": job_id,
                "status": "pending",
                "progress": 0,
                "step": "warteschlange",
                "message": "In der Warteschlange…",
                "file_id": row["id"],
                "filename": row["name"],
                "host": host,
                "printer": printer["name"],
                "spool_id": chosen,
                "override_machine": bool(override),
                "ace_map": mapping,
                "bed_leveling": want_level,
                "result": None,
                "error": None,
                "created_at": time.time(),
                "updated_at": time.time(),
            }
            with self._lock:
                self._jobs[job_id] = job
                self._order.appendleft(job_id)
            self._pending.put(
                {
                    "id": job_id,
                    "file_id": row["id"],
                    "host": host,
                    "override": bool(override),
                    "spool_id": chosen,
                    "ace_map": mapping,
                    "bed_leveling": want_level,
                }
            )
            created.append(self.snapshot(job_id))
        return created

    def get(self, job_id):
        return self.snapshot(job_id)

    def list_jobs(self, limit=20):
        with self._lock:
            ids = list(self._order)[: max(1, int(limit))]
            return [dict(self._jobs[job_id]) for job_id in ids if job_id in self._jobs]

    def snapshot(self, job_id):
        with self._lock:
            job = self._jobs.get(job_id)
            return dict(job) if job else None

    def _update(self, job_id, **fields):
        with self._lock:
            job = self._jobs.get(job_id)
            if not job:
                return
            job.update(fields)
            job["updated_at"] = time.time()
        publish_live("senden", {})

    def _worker(self):
        while not self._stop.is_set():
            try:
                item = self._pending.get(timeout=1)
            except Empty:
                continue
            if item is None:
                break
            self._run(item)

    def _run(self, item):
        job_id = item["id"]

        def on_progress(step, progress, message):
            self._update(
                job_id,
                status="running",
                step=step,
                progress=int(progress),
                message=message,
            )

        self._update(
            job_id,
            status="running",
            step="vorbereiten",
            progress=5,
            message="Vorbereitung…",
        )
        try:
            with connect(self.database_url) as conn:
                result = jobs.send_file(
                    conn,
                    self.cfg,
                    item["file_id"],
                    item["host"],
                    True,
                    item["override"],
                    item["spool_id"],
                    on_progress=on_progress,
                    prepared=True,
                    ace_map=item.get("ace_map"),
                    bed_leveling=bool(item.get("bed_leveling")),
                )
                conn.commit()
            done_msg = "Gestartet auf %s: %s" % (result["printer"], result["filename"])
            warnings = result.get("warnings") or []
            if warnings:
                done_msg += " · " + " · ".join(warnings)
            self._update(
                job_id,
                status="done",
                step="fertig",
                progress=100,
                message=done_msg,
                result=result,
                error=None,
            )
        except GrottiError as exc:
            self._update(
                job_id,
                status="error",
                step="fehler",
                progress=100,
                message=exc.message,
                error={"message": exc.message, "code": exc.code},
            )
        except Exception as exc:
            self._update(
                job_id,
                status="error",
                step="fehler",
                progress=100,
                message="Senden ist fehlgeschlagen.",
                error={"message": str(exc) or "Senden ist fehlgeschlagen.", "code": "intern"},
            )