Fulfillment-Backend/generators/residual_qr_code_generator/generator.py

from io import BytesIO
from typing import Optional

from reportlab.lib.units import mm
from reportlab.pdfgen.canvas import Canvas
from reportlab.platypus import Paragraph, SimpleDocTemplate, PageBreak, Frame
from reportlab_qrcode import QRCodeImage
from sqlalchemy import select
from sqlalchemy.orm import selectinload, joinedload

from barcodes.pdf.pdf_maker import PdfMaker
from generators.base_pdf_card_generator.base_pdf_card_generator import BasePdfCardGenerator
from models import Client, ResidualPallet, ResidualBox


class ResidualQRCodeGenerator(BasePdfCardGenerator):
    async def _get_client_by_id(self, client_id: int) -> Optional[Client]:
        stmt = (
            select(Client)
            .where(Client.id == client_id)
            .options(
                selectinload(Client.boxes),
                selectinload(Client.pallets)
                .selectinload(ResidualPallet.boxes),
            )
        )
        client = (await self._session.execute(stmt)).one_or_none()
        return client[0] if client else None

    @staticmethod
    def _split_string(string: str) -> list[int]:
        if not string:
            return []
        return [int(item) for item in string.split(",")]

    async def generate(self, pallet_ids_str: str, box_ids_str: str):
        pallet_ids = self._split_string(pallet_ids_str)
        box_ids = self._split_string(box_ids_str)

        pallets_buffer = await self.generate_pallets(pallet_ids)
        boxes_buffer = await self.generate_boxes(box_ids)
        return self._merge_pdfs([pallets_buffer, boxes_buffer])

    async def _get_pallets(self, pallet_ids: list[int]) -> list[ResidualPallet]:
        stmt = (
            select(ResidualPallet)
            .options(
                joinedload(ResidualPallet.client),
            )
            .where(ResidualPallet.id.in_(pallet_ids))
            .order_by(ResidualPallet.id.asc())
        )
        pallets = await self._session.execute(stmt)
        return list(pallets.unique().scalars().all())

    def _generate_empty_doc(self) -> BytesIO:
        buffer = BytesIO()
        doc: SimpleDocTemplate = self._create_doc(buffer)
        doc.build([])
        buffer.seek(0)
        return buffer

    async def generate_pallets(self, pallet_ids: list[int]) -> BytesIO:
        if not pallet_ids:
            return self._generate_empty_doc()

        buffer = BytesIO()
        doc: SimpleDocTemplate = self._create_doc(buffer)

        pallet_idx = 0
        pallets = await self._get_pallets(pallet_ids)
        client = pallets[0].client

        def on_page(canvas: Canvas, _):
            nonlocal pallet_idx, pallets
            pallet_id = pallets[pallet_idx].id

            qr = QRCodeImage(f"П{pallet_id}", size=30 * mm)
            qr.drawOn(canvas, 0, 30)

            object_name = Paragraph(f"Паллет", self.small_centered_style)
            pallet_id = Paragraph(f"ID: П{pallet_id}", self.small_centered_style)

            frame = Frame(x1=28 * mm, y1=3 * mm, width=30 * mm, height=30 * mm)
            frame.addFromList([object_name, pallet_id], canvas)

            client_name = Paragraph(f"Клиент: {client.name}", self.small_centered_style)
            frame = Frame(x1=0 * mm, y1=-7 * mm, width=58 * mm, height=20 * mm)
            frame.addFromList([client_name], canvas)

            pallet_idx += 1

        elements = []
        for _ in range(len(pallets)):
            elements.append(Paragraph("", self.medium_style))
            elements.append(PageBreak())

        doc.build(elements, on_page, on_page)

        buffer.seek(0)
        return buffer

    async def _get_boxes(self, box_ids: list[int]) -> list[ResidualBox]:
        stmt = (
            select(ResidualBox)
            .options(
                joinedload(ResidualBox.client),
                selectinload(ResidualBox.pallet)
                .joinedload(ResidualPallet.client),
            )
            .where(ResidualBox.id.in_(box_ids))
            .order_by(ResidualBox.id.asc())
        )
        boxes = await self._session.execute(stmt)
        return list(boxes.unique().scalars().all())

    async def generate_boxes(self, box_ids: list[int]) -> BytesIO:
        if not box_ids:
            return self._generate_empty_doc()

        buffer = BytesIO()
        doc: SimpleDocTemplate = self._create_doc(buffer)

        box_idx = 0
        boxes = await self._get_boxes(box_ids)
        client = boxes[0].client or boxes[0].pallet.client

        def on_page(canvas: Canvas, _):
            nonlocal box_idx
            box_id = boxes[box_idx].id

            qr = QRCodeImage(f"К{box_id}", size=30 * mm)
            qr.drawOn(canvas, 0, 30)

            box_info = [
                Paragraph("Короб", self.small_centered_style),
                Paragraph(f"ID: К{box_id}", self.small_centered_style),
            ]
            if boxes[box_idx].pallet_id:
                box_info.append(Paragraph("На паллете", self.small_centered_style))
                box_info.append(Paragraph(f"ID: П{boxes[box_idx].pallet_id}", self.small_centered_style))

            frame = Frame(x1=28 * mm, y1=8 * mm, width=30 * mm, height=30 * mm)
            frame.addFromList(box_info, canvas)

            client_name = Paragraph(f"Клиент: {client.name}", self.small_centered_style)
            frame = Frame(x1=0 * mm, y1=-7 * mm, width=58 * mm, height=20 * mm)
            frame.addFromList([client_name], canvas)

            box_idx += 1

        elements = []
        for _ in range(len(boxes)):
            elements.append(Paragraph("", self.medium_style))
            elements.append(PageBreak())

        doc.build(elements, on_page, on_page)

        buffer.seek(0)
        return buffer

    def _merge_pdfs(self, buffers: list[BytesIO]) -> BytesIO:
        pdf_maker = PdfMaker((self.page_width, self.page_height))
        for buffer in buffers:
            pdf_maker.add_pdfs(buffer)
        return pdf_maker.get_bytes()