printer_utils.py

"""Printer handling and detection utilities for the Sticker Factory."""

import logging
import subprocess
import tempfile
import time
import os
from pathlib import Path
from brother_ql.models import ModelsManager
from brother_ql.backends import backend_factory
from brother_ql import labels
from brother_ql.raster import BrotherQLRaster
from brother_ql.conversion import convert
from brother_ql.backends.helpers import send
import usb.core
from dataclasses import dataclass

import streamlit as st
from job_queue import print_queue
from config_manager import PRIVACY_MODE, DEBUG_MODE, FALLBACK_LABEL_TYPE, FALLBACK_MODELS

logger = logging.getLogger("sticker_factory.printer_utils")

def safe_filename(text):
    epoch_time = int(time.time())
    return f"{epoch_time}_{text}.png"

@dataclass
class PrinterInfo:
    identifier: str
    backend: str
    protocol: str
    vendor_id: str
    product_id: str
    serial_number: str
    name: str = "Brother QL Printer"
    model: str = "QL-570"
    status: str = "unknown"
    label_type: str = "unknown"
    label_size : str = "unknown"
    label_width: int = 0
    label_height: int = 0
    
    def __getitem__(self, item):
        return getattr(self, item)
    
    def __setitem__(self, key, value):
        setattr(self, key, value)


def create_virtual_printer():
    """Create a virtual printer for debug mode."""
    virtual_printer = PrinterInfo(
        identifier="virtual/debug/0000",
        backend="virtual",
        model="QL-570",
        protocol="virtual",
        vendor_id="0000",
        product_id="0000",
        serial_number="DEBUG-0000",
        name="Virtual Debug Printer",
        status="Waiting to receive",
        label_type=FALLBACK_LABEL_TYPE,
        label_size=f"{FALLBACK_LABEL_TYPE}mm",
        label_width=get_label_width(FALLBACK_LABEL_TYPE),
        label_height=None,
    )
    logger.info("Created virtual debug printer")
    return virtual_printer


def find_and_parse_printer():
    logger.info("Searching for Brother QL printers...")
    model_manager = ModelsManager()
    
    found_printers = []
    
    # Add virtual printer if debug mode is enabled
    if DEBUG_MODE:
        virtual_printer = create_virtual_printer()
        found_printers.append(virtual_printer)
        logger.info("DEBUG MODE: Added virtual printer to available printers")

    for backend_name in ["pyusb", "linux_kernel"]:
        try:
            logger.debug(f"Trying backend: {backend_name}")
            backend = backend_factory(backend_name)
            available_devices = backend["list_available_devices"]()
            logger.debug(f"Found {len(available_devices)} devices with {backend_name} backend")
            
            for printer in available_devices:
                logger.debug(f"Found device: {printer}")
                identifier = printer["identifier"]
                parts = identifier.split("/")

                if len(parts) < 4:
                    logger.warning(f"Skipping device with invalid identifier format: {identifier}")
                    continue

                protocol = parts[0]
                device_info = parts[2]
                serial_number = parts[3]
                
                try:
                    vendor_id, product_id = device_info.split(":")
                except ValueError:
                    logger.warning(f"Invalid device info format: {device_info}")
                    continue
                
                try:
                    product_id_int = int(product_id, 16)
                    for m in model_manager.iter_elements():
                        if m.product_id == product_id_int:
                            model = m.identifier
                            break
                    logger.debug(f"Matched printer model: {model}")
                except ValueError:
                    logger.warning(f"Invalid product ID format: {product_id}")
                    continue

                printer_info = PrinterInfo(
                    identifier=identifier,
                    backend=backend_name,
                    model=model,
                    protocol=protocol,
                    vendor_id=vendor_id,
                    product_id=product_id,
                    serial_number=serial_number,
                )

                found_printers.append(printer_info)   
                printer_info['name'] = f"{printer_info['model']} - {printer_info['serial_number'][-4:]}"
                get_printer_status(printer_info)
                logger.debug(f"Added printer: {printer_info}")

        except Exception as e:
            logger.error(f"Error with backend {backend_name}: {str(e)}")
            continue    
    return found_printers


def get_printer_status(printer):
    printer['status'] = "unknown"
    printer['label_type'] = "unknown"
    printer['label_size'] = "unknown"
    printer['label_width'] = 0
    printer['label_height'] = 0
    logger.debug(f"Checking if '{printer['model']}' is in FALLBACK_MODELS: {FALLBACK_MODELS}")
    if str(printer['model']) in FALLBACK_MODELS:
        printer['label_type'] = FALLBACK_LABEL_TYPE
        printer['label_width'] = get_label_width(FALLBACK_LABEL_TYPE)
        printer['label_height'] = 0
        printer['status'] = "Waiting to receive"
        logger.debug(f"Using fallback label type {printer['label_type']} for model {printer['model']}")
    else:
        try:
            cmd = f"brother_ql -b pyusb --model {printer['model']} -p {printer['identifier']} status"
            logger.debug(f"Running status command: {cmd}")
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=5)
            
            # Log the raw output for debugging
            if result.stdout:
                logger.debug(f"Status command stdout:\n{result.stdout}")
            if result.stderr:
                logger.warning(f"Status command stderr:\n{result.stderr}")
            if result.returncode != 0:
                logger.warning(f"Status command returned non-zero exit code: {result.returncode}")
                
            for line in result.stdout.splitlines():
                if "Phase:" in line:
                    printer['status'] = line.split("Phase:")[1].strip()
                    logger.debug(f"Detected status: {printer['status']}")
                if "Media size:" in line:
                    printer['label_size'] = line.split("Media size:")[1].strip()
                    size_str = line.split("Media size:")[1].strip().split('x')[0].strip()
                    try:
                        media_width_mm = int(size_str)
                        label_sizes = {
                            12: "12", 29: "29", 38: "38", 50: "50", 54: "54",
                            62: "62", 102: "102", 103: "103", 104: "104"
                        }
                        if media_width_mm in label_sizes:
                            label_type = label_sizes[media_width_mm]
                            printer['label_type'] = label_type
                            printer['label_width'] = get_label_width(label_type)
                            printer['label_height'] = None
                            logger.debug(f"Detected label type: {label_type} from width: {media_width_mm}mm")
                    except Exception as e:
                        logger.warning(f"Exception parsing media width: {str(e)}")
            logger.info(f"Printer {printer['name']}: label type: {printer['label_type']}, status: {printer['status']}")

        except subprocess.TimeoutExpired:
            logger.error(f"Timeout getting status for printer {printer['name']} - USB might be busy")
            printer['status'] = "timeout"
        except Exception as e:
            logger.warning(f"Error getting status for printer {printer['name']}: {str(e)}")
            printer['status'] = str(e)

def get_label_width(label_type):
    """Get the pixel width of a label type."""
    label_definitions = labels.ALL_LABELS
    for label in label_definitions:
        if label.identifier == label_type:
            width = label.dots_printable[0]
            logger.debug(f"Label type {label_type} width: {width} dots")
            return width
    raise ValueError(f"Label type {label_type} not found in label definitions")

def print_image(image, printer_info, rotate=0, dither=False):
    """Queue a print job."""
    temp_dir = tempfile.gettempdir()
    os.makedirs(temp_dir, exist_ok=True)

    with tempfile.NamedTemporaryFile(suffix=".png", delete=False, dir=temp_dir) as temp_file:
        temp_file_path = temp_file.name
        image.save(temp_file_path, "PNG")
        logger.info(f"{temp_file_path} added to print queue for printer {printer_info['name']}")

    logger.debug(f"Using label type: {printer_info['label_type']}")

    job_id = print_queue.add_job(
        image,
        rotate=rotate,
        dither=dither,
        printer_info=printer_info,
        temp_file_path=temp_file_path,
        label_type=printer_info["label_type"]
    )

    status = print_queue.get_job_status(job_id)
    status_container = st.empty()
    
    while status.status in ["pending", "processing"]:
        status_container.info(f"Print job status: {status.status}")
        time.sleep(0.5)
        status = print_queue.get_job_status(job_id)

    if status.status == "completed":
        status_container.success("Print job completed successfully!")
        if PRIVACY_MODE:
            status_container.info("Privacy mode is enabled; sticker not saved locally.")
        else:
            filename = safe_filename("Stikka-")
            file_path = os.path.join("labels", filename)
            image.save(file_path, "PNG")
            status_container.success(f"Sticker saved as {filename}")
        
        # Record statistics - DISABLED on Raspberry Pi due to SIGILL compatibility issues
        # Uncomment below if stats module works on your system
        # try:
        #     import importlib
        #     stats_module = importlib.import_module('stats_utils')
        #     record_print = getattr(stats_module, 'record_print', None)
        #     if record_print:
        #         printer_name = printer_info['name']
        #         printer_model = getattr(printer_info, 'model', None)
        #         record_print(printer_name, printer_model)
        # except Exception:
        #     pass
        
        return True
    else:
        status_container.error(f"Print job failed: {status.error}")
        return False


def process_print_job(image, printer_info, temp_file_path, rotate=0, dither=False, label_type="102"):
    """
    Process a single print job.
    Returns (success, error_message)
    """

    try:
        # If debug mode is enabled, use virtual printer (save to debug directory)
        if DEBUG_MODE:
            debug_dir = Path("debug")
            debug_dir.mkdir(exist_ok=True)
            
            # Generate a filename with timestamp
            timestamp = int(time.time())
            filename = f"{timestamp}_debug_print_{printer_info['name'].replace(' ', '_')}.png"
            output_path = debug_dir / filename
            
            # Copy the image to debug directory
            image.save(output_path, "PNG")
            logger.info(f"DEBUG MODE: Virtual printer saved file to {output_path}")
            logger.debug(f"""
            Debug print parameters:
            - Label type: {label_type}
            - Rotate: {rotate}
            - Dither: {dither}
            - Model: {printer_info['model']}
            - Output: {output_path}
            """)
            return True, None
        
        # Prepare the image for printing
        qlr = BrotherQLRaster(printer_info["model"])
        
        logger.debug(f"Printing {temp_file_path} on label type {label_type} on printer {printer_info['name']}")
        
        instructions = convert(
            qlr=qlr,
            images=[temp_file_path],
            label=label_type,
            rotate=rotate,
            threshold=70,
            dither=dither,
            compress=True,
            red=False,
            dpi_600=False,
            hq=False,
            cut=True,
        )


        logger.debug(f"""
        Print parameters:
        - Label type: {label_type}
        - Rotate: {rotate}
        - Dither: {dither}
        - Model: {printer_info['model']}
        - Backend: {printer_info['backend']}
        - Identifier: {printer_info['identifier']}
        """)

        # Try to print using Python API
        success = send(
            instructions=instructions,
            printer_identifier=printer_info["identifier"],
            backend_identifier="pyusb"
        )
        
        if not success:
            return False, "Failed to print using Python API"

        return True, None

    except usb.core.USBError as e:
        # Treat timeout errors as successful since they often occur after print completion
        if e.errno == 110:  # Operation timed out
            logger.error("USB timeout occurred - this is normal and the print likely completed")
            return True, "Print completed (timeout is normal)"
        error_msg = f"USBError encountered: {e}"
        logger.error(error_msg)
        return False, error_msg

    except Exception as e:
        error_msg = f"Unexpected error during printing: {str(e)}"
        logger.error(error_msg)
        return False, error_msg

    finally:
        # Clean up temporary file
        try:
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)
                logger.debug(f"Temporary file {temp_file_path} deleted.")
        except Exception as e:
            logger.warning(f"Failed to delete temporary file {temp_file_path}: {str(e)}")