# =====================================================
# Apckeyl_RealESRGAN
# Version 2.0
# batch_processor.py
# =====================================================

from image_upscaler import upscale_image
from priority_queue import PriorityQueue


class BatchProcessor:

    """
    Batch Processor.

    Version 2.0

    Использует Priority Queue.
    """

    def __init__(self):

        self.results = []

        self.queue = PriorityQueue()

    # -------------------------------------------------

    def clear(self):

        self.results.clear()

        self.queue.clear()

    # -------------------------------------------------

    def add_task(

        self,

        image,

        mode,

        quality,

        denoise,

        owner=False,

    ):

        self.queue.add(

            {

                "image": image,

                "mode": mode,

                "quality": quality,

                "denoise": denoise,

            },

            owner=owner,

        )

    # -------------------------------------------------

    def process_queue(self):

        self.results.clear()

        while self.queue.size() > 0:

            task = self.queue.get_next()

            result = upscale_image(

                image=task["image"],

                mode=task["mode"],

                quality=task["quality"],

                denoise=task["denoise"],

            )

            self.results.append(

                result

            )

        return self.results

    # -------------------------------------------------

    def process(

        self,

        images,

        mode,

        quality,

        denoise,

        owner=False,

    ):

        self.clear()

        for image in images:

            self.add_task(

                image=image,

                mode=mode,

                quality=quality,

                denoise=denoise,

                owner=owner,

            )

        return self.process_queue()

    # -------------------------------------------------

    def process_with_limit(

        self,

        images,

        mode,

        quality,

        denoise,

        owner=False,

        max_images=10,

    ):

        if images is None:

            return []

        images = images[:max_images]

        return self.process(

            images=images,

            mode=mode,

            quality=quality,

            denoise=denoise,

            owner=owner,

        )

    # -------------------------------------------------

    def count(self):

        return len(

            self.results

        )

    # -------------------------------------------------

    def queue_size(self):

        return self.queue.size()