bambuddy/backend/app/services/virtual_printer/manager.py

"""Virtual Printer Manager - coordinates SSDP, MQTT, and FTP services."""

import asyncio
import logging
from collections.abc import Callable
from datetime import UTC, datetime
from pathlib import Path

from backend.app.core.config import settings as app_settings
from backend.app.services.virtual_printer.certificate import CertificateService
from backend.app.services.virtual_printer.ftp_server import VirtualPrinterFTPServer
from backend.app.services.virtual_printer.mqtt_server import SimpleMQTTServer
from backend.app.services.virtual_printer.ssdp_server import VirtualPrinterSSDPServer

logger = logging.getLogger(__name__)


# Mapping of SSDP model codes to display names
# These are the codes that slicers expect during discovery
# Sources:
#   - https://gist.github.com/Alex-Schaefer/72a9e2491a42da2ef99fb87601955cc3
#   - https://github.com/psychoticbeef/BambuLabOrcaSlicerDiscovery
VIRTUAL_PRINTER_MODELS = {
    # X1 Series
    "3DPrinter-X1-Carbon": "X1C",  # X1 Carbon
    "3DPrinter-X1": "X1",  # X1
    "C13": "X1E",  # X1E
    # P Series
    "C11": "P1P",  # P1P
    "C12": "P1S",  # P1S
    "N7": "P2S",  # P2S
    # A1 Series
    "N2S": "A1",  # A1
    "N1": "A1 Mini",  # A1 Mini
    # H2 Series
    "O1D": "H2D",  # H2D
    "O1C": "H2C",  # H2C
    "O1S": "H2S",  # H2S
}

# Serial number prefixes for each model (based on Bambu Lab serial number format)
# Format: MMM??RYMDDUUUUU (15 chars total)
#   MMM = Model prefix (3 chars)
#   ?? = Unknown/revision code (2 chars)
#   R = Revision letter (1 char)
#   Y = Year digit (1 char)
#   M = Month (1 char, hex: 1-9, A=Oct, B=Nov, C=Dec)
#   DD = Day (2 chars)
#   UUUUU = Unit number (5 chars)
MODEL_SERIAL_PREFIXES = {
    # X1 Series
    "3DPrinter-X1-Carbon": "00M00A",  # X1C
    "3DPrinter-X1": "00M00A",  # X1
    "C13": "03W00A",  # X1E
    # P Series
    "C11": "01S00A",  # P1P
    "C12": "01P00A",  # P1S
    "N7": "22E00A",  # P2S
    # A1 Series
    "N2S": "03900A",  # A1
    "N1": "03000A",  # A1 Mini
    # H2 Series
    "O1D": "09400A",  # H2D
    "O1C": "09400A",  # H2C
    "O1S": "09400A",  # H2S
}

# Default model
DEFAULT_VIRTUAL_PRINTER_MODEL = "3DPrinter-X1-Carbon"  # X1C


class VirtualPrinterManager:
    """Manages the virtual printer lifecycle and coordinates all services."""

    # Fixed configuration
    PRINTER_NAME = "Bambuddy"
    SERIAL_SUFFIX = "391800001"  # Fixed suffix for virtual printer

    def __init__(self):
        """Initialize the virtual printer manager."""
        self._session_factory: Callable | None = None
        self._enabled = False
        self._access_code = ""
        self._mode = "immediate"
        self._model = DEFAULT_VIRTUAL_PRINTER_MODEL

        # Service instances
        self._ssdp: VirtualPrinterSSDPServer | None = None
        self._ftp: VirtualPrinterFTPServer | None = None
        self._mqtt: SimpleMQTTServer | None = None

        # Background tasks
        self._tasks: list[asyncio.Task] = []

        # Directories
        self._base_dir = app_settings.base_dir / "virtual_printer"
        self._upload_dir = self._base_dir / "uploads"
        self._cert_dir = self._base_dir / "certs"

        # Certificate service
        self._cert_service = CertificateService(self._cert_dir)

        # Track pending uploads for MQTT correlation
        self._pending_files: dict[str, Path] = {}

    def _get_serial_for_model(self, model: str) -> str:
        """Get appropriate serial number for the given model.

        Args:
            model: SSDP model code (e.g., 'BL-P001', 'C11')

        Returns:
            Serial number with correct prefix for the model
        """
        prefix = MODEL_SERIAL_PREFIXES.get(model, "00M09A")
        return f"{prefix}{self.SERIAL_SUFFIX}"

    @property
    def printer_serial(self) -> str:
        """Get the current printer serial number based on model."""
        return self._get_serial_for_model(self._model)

    def set_session_factory(self, session_factory: Callable) -> None:
        """Set the database session factory.

        Args:
            session_factory: Async context manager for database sessions
        """
        self._session_factory = session_factory

    @property
    def is_enabled(self) -> bool:
        """Check if virtual printer is enabled."""
        return self._enabled

    @property
    def is_running(self) -> bool:
        """Check if virtual printer services are running."""
        return len(self._tasks) > 0 and all(not t.done() for t in self._tasks)

    async def configure(
        self,
        enabled: bool,
        access_code: str = "",
        mode: str = "immediate",
        model: str = "",
    ) -> None:
        """Configure and start/stop virtual printer.

        Args:
            enabled: Whether to enable the virtual printer
            access_code: Authentication password for slicer connections
            mode: Archive mode - 'immediate' or 'queue'
            model: SSDP model code (e.g., 'BL-P001' for X1C)
        """
        if enabled and not access_code:
            raise ValueError("Access code is required when enabling virtual printer")

        # Validate model if provided
        new_model = model if model and model in VIRTUAL_PRINTER_MODELS else self._model
        model_changed = new_model != self._model
        old_model = self._model

        logger.debug(
            f"configure() called: enabled={enabled}, self._enabled={self._enabled}, "
            f"model={model}, new_model={new_model}, old_model={old_model}, model_changed={model_changed}"
        )

        self._access_code = access_code
        self._mode = mode
        self._model = new_model

        if enabled and not self._enabled:
            logger.info("Starting virtual printer (was disabled)")
            await self._start()
        elif not enabled and self._enabled:
            logger.info("Stopping virtual printer (was enabled)")
            await self._stop()
        elif enabled and self._enabled and model_changed:
            # Model changed while running - restart services
            logger.info(f"Model changed from {old_model} to {new_model}, restarting...")
            await self._stop()
            # Give time for ports to be released
            await asyncio.sleep(0.5)
            await self._start()
            logger.info("Virtual printer restarted with new model")
        else:
            logger.debug(
                f"No state change needed (enabled={enabled}, self._enabled={self._enabled}, model_changed={model_changed})"
            )

        self._enabled = enabled

    async def _start(self) -> None:
        """Start all virtual printer services."""
        logger.info("Starting virtual printer services...")

        # Update certificate service with current serial (based on model)
        current_serial = self.printer_serial
        self._cert_service.serial = current_serial

        # Regenerate printer cert if serial changed (CA is preserved)
        self._cert_service.delete_printer_certificate()
        cert_path, key_path = self._cert_service.generate_certificates()
        logger.info(f"Generated certificate for serial: {current_serial}")

        # Create directories
        self._upload_dir.mkdir(parents=True, exist_ok=True)
        (self._upload_dir / "cache").mkdir(exist_ok=True)

        # Initialize services
        self._ssdp = VirtualPrinterSSDPServer(
            name=self.PRINTER_NAME,
            serial=self.printer_serial,
            model=self._model,
        )

        self._ftp = VirtualPrinterFTPServer(
            upload_dir=self._upload_dir,
            access_code=self._access_code,
            cert_path=cert_path,
            key_path=key_path,
            on_file_received=self._on_file_received,
        )

        self._mqtt = SimpleMQTTServer(
            serial=self.printer_serial,
            access_code=self._access_code,
            cert_path=cert_path,
            key_path=key_path,
            on_print_command=self._on_print_command,
        )

        # Start services as background tasks
        # Wrap each in error handler so one failure doesn't stop others
        async def run_with_logging(coro, name):
            try:
                await coro
            except Exception as e:
                logger.error(f"Virtual printer {name} failed: {e}")

        self._tasks = [
            asyncio.create_task(run_with_logging(self._ssdp.start(), "SSDP"), name="virtual_printer_ssdp"),
            asyncio.create_task(run_with_logging(self._ftp.start(), "FTP"), name="virtual_printer_ftp"),
            asyncio.create_task(run_with_logging(self._mqtt.start(), "MQTT"), name="virtual_printer_mqtt"),
        ]

        logger.info(f"Virtual printer '{self.PRINTER_NAME}' started (serial: {self.printer_serial})")

    async def _stop(self) -> None:
        """Stop all virtual printer services."""
        logger.info("Stopping virtual printer services...")

        # Stop services first - this closes servers and cancels active sessions
        if self._ftp:
            await self._ftp.stop()
            self._ftp = None

        if self._mqtt:
            await self._mqtt.stop()
            self._mqtt = None

        if self._ssdp:
            await self._ssdp.stop()
            self._ssdp = None

        # Cancel remaining tasks with short timeout
        for task in self._tasks:
            task.cancel()

        if self._tasks:
            try:
                await asyncio.wait_for(asyncio.gather(*self._tasks, return_exceptions=True), timeout=1.0)
            except TimeoutError:
                logger.debug("Some tasks didn't stop in time")

        self._tasks = []

        logger.info("Virtual printer stopped")

    async def _on_file_received(self, file_path: Path, source_ip: str) -> None:
        """Handle file upload completion from FTP.

        Args:
            file_path: Path to uploaded file
            source_ip: IP address of the uploading slicer
        """
        logger.info(f"Virtual printer received file: {file_path.name} from {source_ip}")

        # Store file reference for MQTT correlation
        self._pending_files[file_path.name] = file_path

        # In immediate mode, archive right away
        # In queue mode, create pending upload record
        if self._mode == "immediate":
            await self._archive_file(file_path, source_ip)
        else:
            await self._queue_file(file_path, source_ip)

    async def _on_print_command(self, filename: str, data: dict) -> None:
        """Handle print command from MQTT.

        In a real printer, this would start the print. For virtual printer,
        we just log it since archiving is handled by file upload.

        Args:
            filename: Name of the file to print
            data: Print command data (contains settings like timelapse, bed_leveling, etc.)
        """
        logger.info(f"Virtual printer received print command for: {filename}")
        logger.debug(f"Print command data: {data}")

        # The file should already be archived from FTP upload
        # This command just confirms the slicer's intent to "print"

    async def _archive_file(self, file_path: Path, source_ip: str) -> None:
        """Archive file immediately.

        Args:
            file_path: Path to the 3MF file
            source_ip: IP address of uploader
        """
        if not self._session_factory:
            logger.error("Cannot archive: no database session factory configured")
            return

        # Only archive 3MF files
        if file_path.suffix.lower() != ".3mf":
            logger.debug(f"Skipping non-3MF file: {file_path.name}")
            # Remove from pending and clean up
            self._pending_files.pop(file_path.name, None)
            try:
                file_path.unlink()
            except Exception:
                pass
            return

        try:
            from backend.app.services.archive import ArchiveService

            async with self._session_factory() as db:
                service = ArchiveService(db)

                # Archive the print
                archive = await service.archive_print(
                    printer_id=None,  # No physical printer
                    source_file=file_path,
                    print_data={
                        "status": "archived",
                        "source": "virtual_printer",
                        "source_ip": source_ip,
                    },
                )

                if archive:
                    logger.info(f"Archived virtual printer upload: {archive.id} - {archive.print_name}")

                    # Clean up uploaded file (it's now copied to archive)
                    try:
                        file_path.unlink()
                    except Exception:
                        pass

                    # Remove from pending
                    self._pending_files.pop(file_path.name, None)
                else:
                    logger.error(f"Failed to archive file: {file_path.name}")

        except Exception as e:
            logger.error(f"Error archiving file: {e}")

    async def _queue_file(self, file_path: Path, source_ip: str) -> None:
        """Queue file for user review.

        Args:
            file_path: Path to the 3MF file
            source_ip: IP address of uploader
        """
        if not self._session_factory:
            logger.error("Cannot queue: no database session factory configured")
            return

        # Only queue 3MF files
        if file_path.suffix.lower() != ".3mf":
            logger.warning(f"Skipping non-3MF file: {file_path.name}")
            return

        try:
            from backend.app.models.pending_upload import PendingUpload

            async with self._session_factory() as db:
                pending = PendingUpload(
                    filename=file_path.name,
                    file_path=str(file_path),
                    file_size=file_path.stat().st_size,
                    source_ip=source_ip,
                    status="pending",
                    uploaded_at=datetime.now(UTC),
                )
                db.add(pending)
                await db.commit()

                logger.info(f"Queued virtual printer upload: {pending.id} - {file_path.name}")

                # Remove from pending files dict
                self._pending_files.pop(file_path.name, None)

        except Exception as e:
            logger.error(f"Error queueing file: {e}")

    def get_status(self) -> dict:
        """Get virtual printer status.

        Returns:
            Status dictionary with enabled, running, mode, etc.
        """
        return {
            "enabled": self._enabled,
            "running": self.is_running,
            "mode": self._mode,
            "name": self.PRINTER_NAME,
            "serial": self.printer_serial,
            "model": self._model,
            "model_name": VIRTUAL_PRINTER_MODELS.get(self._model, self._model),
            "pending_files": len(self._pending_files),
        }


# Global instance
virtual_printer_manager = VirtualPrinterManager()