from pathlib import Path

import bpy

from ..runtime import JobOperatorBase
from ..properties import BEN2_MODEL_KEY, REALESRGAN_GENERAL_WDN_MODEL_KEY
from .common.ai_job import (
    require_background_input_path,
    da3_parameters,
    require_input_path,
    moge2_parameters,
    require_environment,
    require_model,
)


def _open_job_directory(directory):
    directory = Path(directory)
    if not directory.is_dir():
        raise RuntimeError(f"Job directory does not exist: {directory}")
    result = bpy.ops.wm.path_open(filepath=str(directory))
    if "FINISHED" not in result:
        raise RuntimeError(f"Unable to open job directory: {directory}")


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,
                require_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,
                require_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(require_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}"


class DebugUpscale(JobOperatorBase, bpy.types.Operator):
    bl_idname = "anyimage.debug_upscale"
    bl_label = "Debug Upscale"
    job_type = "debug-upscale"
    starting_message = "Starting Upscale inference"

    model: bpy.props.StringProperty(
        default=REALESRGAN_GENERAL_WDN_MODEL_KEY,
        options={"HIDDEN", "SKIP_SAVE"},
    )

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

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