import bpy

from ..properties import BEN2_MODEL_KEY
from ..runtime import JobOperatorBase
from ..common.ai import (
    require_background_input_path,
    invoke_ai_setup_if_needed,
    production_device,
    require_environment,
    require_model,
)

from ..common.image import cleanup_image_input
from ..common.image_target import ImageEditTarget, image_edit_owner


_HIDDEN = {"HIDDEN", "SKIP_SAVE"}


class RemoveImageBackground(bpy.types.Operator):
    bl_idname = "anyimage.remove_image_background"
    bl_label = "Remove Background"
    bl_description = "Remove the image background with BEN2"
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        from ..runtime import runtime

        return image_edit_owner(context) is not None and not runtime.server_busy()

    def invoke(self, _context, _event):
        setup_result = invoke_ai_setup_if_needed()
        if setup_result is not None:
            return setup_result
        return self.execute(_context)

    def execute(self, _context):
        result = bpy.ops.anyimage.run_background_removal()
        if "CANCELLED" not in result:
            return {"FINISHED"}
        return {"CANCELLED"}


class RunBackgroundRemoval(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.run_background_removal"
    bl_label = "Run Background Removal"
    job_type = "remove-background"
    starting_message = "Starting image background removal"

    input_path: bpy.props.StringProperty(options=_HIDDEN)
    delete_input: bpy.props.BoolProperty(default=False, options=_HIDDEN)

    def execute(self, context):
        try:
            self._image_target = ImageEditTarget.capture(context)
            source_path, self.delete_input = self._image_target.prepare(context)
        except RuntimeError as error:
            self.report({"ERROR"}, str(error))
            return {"CANCELLED"}
        self.input_path = str(source_path)
        result = super().execute(context)
        if "CANCELLED" in result:
            self.cleanup()
        return result

    def request(self, context):
        require_environment()
        require_model(BEN2_MODEL_KEY)
        settings = context.scene.anyimage_settings
        return {
            "input": str(require_background_input_path(self.input_path)),
            "model": BEN2_MODEL_KEY,
            "device": production_device().lower(),
            "video_frames": settings.video_frames,
            "start_frame": self._image_target.start_frame,
        }

    def response(self, _context, result):
        self._image_target.apply(result.directory)
        if "FINISHED" not in bpy.ops.ed.undo_push(message="Remove Background"):
            raise RuntimeError("Unable to create the Remove Background Undo step")
        return "Image background is removed"

    def cleanup(self):
        cleanup_image_input(self.input_path, self.delete_input)
        self._image_target = None