import bpy

from ..runtime import JobOperatorBase
from ..properties import BEN2_MODEL_KEY
from .common import (
    background_input_path,
    da3_parameters,
    input_path,
    moge2_parameters,
    open_job_directory,
    require_environment,
    require_model,
)


class DebugDa3(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.debug_da3"
    bl_label = "Debug DA3"
    job_type = "debug-da3"
    starting_message = "Starting DA3 API inference"

    def request(self, context):
        require_environment()
        settings = context.scene.anyimage_settings
        model = settings.model
        require_model(model)
        return {
            **da3_parameters(
                settings,
                input_path(settings.input_path),
                model,
                settings.process_res,
                1,
            ),
            "device": settings.device.lower(),
            "inference_mode": "joint",
            "remove_background_model": "",
        }

    def response(self, _context, result):
        open_job_directory(result.directory)
        return f"Saved DA3 API prediction to {result.directory}"


class DebugMoge2(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.debug_moge2"
    bl_label = "Debug MoGe-2"
    job_type = "debug-moge2"
    starting_message = "Starting MoGe-2 inference"

    def request(self, context):
        require_environment()
        settings = context.scene.anyimage_settings
        require_model(settings.moge_model)
        return {
            **moge2_parameters(
                settings,
                input_path(settings.input_path),
                settings.moge_model,
                settings.moge_resolution_level,
                1,
            ),
            "device": settings.device.lower(),
        }

    def response(self, _context, result):
        open_job_directory(result.directory)
        return f"Exported MoGe-2 data to {result.directory}"


class DebugBen2(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.debug_ben2"
    bl_label = "Debug BEN2"
    job_type = "debug-ben2"
    starting_message = "Starting BEN2 background removal"

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

    def response(self, _context, result):
        open_job_directory(result.directory)
        return f"Removed background to {result.directory}"