from PIL import Image

from sd_bmab.base import cache
from sd_bmab.base import ProcessorBase
from sd_bmab.base.context import Context
from sd_bmab.external import load_external_module


class ICLight(ProcessorBase):

    def __init__(self) -> None:
        super().__init__()
        self.iclight_opt = {}
        self.enabled = False
        self.enable_before_upscale = False
        self.style = 'normal'
        self.prompt = ''
        self.preference = 'None'
        self.use_background_image = False
        self.blending = 0.5

    def preprocess(self, context: Context, image: Image):
        self.iclight_opt = context.args.get('module_config', {}).get('iclight', {})
        self.enabled = self.iclight_opt.get('enabled', self.enabled)
        self.enable_before_upscale = self.iclight_opt.get('enable_before_upscale', self.enable_before_upscale)
        self.style = self.iclight_opt.get('style', self.style)
        self.prompt = self.iclight_opt.get('prompt', self.prompt)
        self.preference = self.iclight_opt.get('preference', self.preference)
        self.use_background_image = self.iclight_opt.get('use_background_image', self.use_background_image)
        self.blending = self.iclight_opt.get('blending', self.blending)
        return self.enabled

    def process(self, context: Context, image: Image):
        mod = load_external_module('iclight', 'bmabiclight')
        bg_image = self.get_background_image() if self.use_background_image else None
        return mod.bmab_relight(context, self.style, image, bg_image, self.prompt, self.blending, self.preference)

    def postprocess(self, context: Context, image: Image):
        pass

    @staticmethod
    def get_background_image():
        img = cache.get_image_from_cache('iclight_background.png')
        if img is not None:
            return img
        return Image.new('RGB', (512, 768), (0, 0, 0))

    @staticmethod
    def put_backgound_image(img):
        cache.put_image_to_cache('iclight_background.png', img)

    @staticmethod
    def get_styles():
        return ['intensive', 'less intensive', 'normal', 'soft']


class ICLightBeforeUpsacle(ICLight):

    def preprocess(self, context: Context, image: Image):
        super().preprocess(context, image)
        return self.enabled and self.enable_before_upscale


class ICLightAfterUpsacle(ICLight):

    def preprocess(self, context: Context, image: Image):
        super().preprocess(context, image)
        return self.enabled and not self.enable_before_upscale