Fulfillment-Backend/marketplaces/wildberries/core.py

import time
from collections import defaultdict

from sqlalchemy import select
from sqlalchemy.orm import selectinload

from external.marketplace.wildberries.core import WildberriesMarketplaceApi
from marketplaces.base.core import BaseMarketplaceController
from models import Product, ProductBarcode, ProductImage, WildberriesProduct


class WildberriesProductSynchronizer:
    products: list[Product]
    barcodes: list[ProductBarcode]
    images: list[ProductImage]
    wildberries_products: list[WildberriesProduct]

    def __init__(self, session, marketplace, api):
        self.session = session
        self.marketplace = marketplace
        self.api = api
        self.products = []
        self.barcodes = []
        self.images = []
        self.wildberries_products = []

    def _clear(self):
        self.products = []
        self.barcodes = []
        self.images = []
        self.wildberries_products = []

    async def _write(self):
        instances = self.products + self.wildberries_products + self.barcodes + self.images
        start = time.time()
        self.session.add_all(instances)
        await self.session.commit()
        print(f'Write time: {time.time() - start}')

    async def _get_synchronized_nm_uuids(self):
        marketplace_id = self.marketplace.id
        synchronized_nm_uuids = set(
            (
                await self.session.scalars(
                    select(
                        WildberriesProduct.nm_uuid
                    )
                    .where(
                        WildberriesProduct.marketplace_id == marketplace_id
                    )
                )
            ).all()
        )
        return synchronized_nm_uuids

    async def _create_product(self, card, size_value):
        return Product(
            client_id=self.marketplace.client_id,
            name=card['title'],
            article=card['vendorCode'],
            size=size_value
        )

    async def _create_barcodes(self, product, skus):
        barcodes = []
        for sku in skus:
            barcode = ProductBarcode(
                product=product,
                barcode=sku
            )
            barcodes.append(barcode)
        return barcodes

    async def _create_images(self, product, photos):
        images = []
        for photo in photos[:1]:
            image = ProductImage(
                product=product,
                image_url=photo['big']
            )
            images.append(image)
        return images

    async def _create_wildberries_product(self, product, nm_uuid):
        return WildberriesProduct(
            marketplace_id=self.marketplace.id,
            product=product,
            nm_uuid=nm_uuid
        )

    async def _update_product_info(self, product, card):
        product.name = card['title']
        product.article = card['vendorCode']

    async def _update_barcodes(self, product, skus):
        existing_barcodes = {barcode.barcode for barcode in product.barcodes}
        new_barcodes = []
        for sku in skus:
            if sku not in existing_barcodes:
                barcode = ProductBarcode(
                    product=product,
                    barcode=sku
                )
                new_barcodes.append(barcode)
        return new_barcodes

    async def _update_images(self, product, photos):
        existing_images = {image.image_url for image in product.images}
        new_images = []
        for photo in photos[:1]:
            if photo['big'] not in existing_images:
                image = ProductImage(
                    product=product,
                    image_url=photo['big']
                )
                new_images.append(image)
        return new_images

    async def _process_product(self, card, size_value, nm_uuid):
        product = await self._create_product(card, size_value)
        barcodes = await self._create_barcodes(product, card.get('sizes')[0].get('skus') or [])
        images = await self._create_images(product, card.get('photos') or [])
        wildberries_product = await self._create_wildberries_product(product, nm_uuid)

        self.products.append(product)
        self.barcodes.extend(barcodes)
        self.images.extend(images)
        self.wildberries_products.append(wildberries_product)

    async def create_products(self):
        self._clear()
        synchronized_nm_uuids = await self._get_synchronized_nm_uuids()
        async for card in self.api.get_all_products():
            nm_uuid = card['nmUUID']
            if nm_uuid in synchronized_nm_uuids:
                continue
            sizes = card.get('sizes') or []
            for size in sizes:
                tech_size = size.get('techSize')
                wb_size = size.get('wbSize')
                size_value = tech_size or wb_size
                await self._process_product(
                    card,
                    size_value,
                    nm_uuid
                )
        await self._write()

    async def synchronize_products(self):
        self._clear()
        synchronized_products_stmt = (
            select(Product)
            .join(
                WildberriesProduct
            )
            .options(
                selectinload(Product.barcodes),
                selectinload(Product.wildberries_products)
            )
            .where(WildberriesProduct.marketplace_id == self.marketplace.id)
        )
        synchronized_products = await self.session.execute(synchronized_products_stmt)
        synchronized_products = synchronized_products.scalars().all()
        synchronized_products_nm_id_dict = defaultdict(list)
        for product in synchronized_products:
            for wb_product in product.wildberries_products:
                synchronized_products_nm_id_dict[wb_product.nm_uuid].append(product)
        synchronized_nm_uuids = list(synchronized_products_nm_id_dict.keys())
        async for card in self.api.get_all_products():
            nm_uuid = card['nmUUID']
            if nm_uuid not in synchronized_nm_uuids:
                continue
            products = synchronized_products_nm_id_dict[nm_uuid]
            existing_sizes = {product.size for product in products}
            size_product_dict = {product.size: product for product in products}

            sizes = card.get('sizes') or []
            for size in sizes:
                tech_size = size.get('techSize')
                wb_size = size.get('wbSize')
                size_value = tech_size or wb_size
                if size_value in existing_sizes:
                    product = size_product_dict[size_value]
                    await self._update_product_info(product, card)
                    self.barcodes.extend(
                        await self._update_barcodes(product, size.get('skus') or [])
                    )
                    self.images.extend(
                        await self._update_images(product, card.get('photos') or [])
                    )
                    continue
                await self._process_product(
                    card,
                    size_value,
                    nm_uuid
                )
        await self._write()


class WildberriesController(
    BaseMarketplaceController
):
    api: WildberriesMarketplaceApi

    def __init__(self, session, marketplace):
        super().__init__(session, marketplace)
        self.product_synchronizer = WildberriesProductSynchronizer(session, marketplace, self.api)

    async def create_products(self):
        await self.product_synchronizer.create_products()

    async def synchronize_products(self):
        await self.product_synchronizer.synchronize_products()