import unittest

try:
    from .test_pro_edit_helpers import import_plugin_submodule, load_plugin_package
except ImportError:
    from test_pro_edit_helpers import import_plugin_submodule, load_plugin_package


class TestProEditRegistration(unittest.TestCase):
    def test_plugin_registers_pro_edit_nodes(self):
        plugin = load_plugin_package()

        self.assertIn("LLSProImageEditPrepare", plugin.NODE_CLASS_MAPPINGS)
        self.assertIn("LLSProKSamplerBridge", plugin.NODE_CLASS_MAPPINGS)
        self.assertIn("LLSProImageEditFinish", plugin.NODE_CLASS_MAPPINGS)
        self.assertEqual(
            plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSProImageEditPrepare"],
            "LLS Pro Image Edit Prepare",
        )
        self.assertEqual(
            plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSProKSamplerBridge"],
            "LLS Pro KSampler Bridge",
        )
        self.assertEqual(
            plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSProImageEditFinish"],
            "LLS Pro Image Edit Finish",
        )

    def test_prepare_node_schema_matches_contract(self):
        plugin = load_plugin_package()
        node_cls = plugin.NODE_CLASS_MAPPINGS["LLSProImageEditPrepare"]
        schema = node_cls.INPUT_TYPES()
        required = schema["required"]
        optional = schema["optional"]

        self.assertEqual(node_cls.CATEGORY, "LLS/Image Edit")
        self.assertEqual(node_cls.FUNCTION, "prepare")
        self.assertEqual(
            node_cls.RETURN_TYPES,
            ("LATENT", "IMAGE", "MASK", "LLS_EDIT_INFO", "FLOAT", "CONDITIONING", "CONDITIONING"),
        )
        self.assertEqual(
            node_cls.RETURN_NAMES,
            ("latent", "work_image", "work_mask", "edit_info", "recommended_denoise", "positive", "negative"),
        )
        self.assertEqual(
            set(required),
            {
                "image",
                "mask",
                "vae",
                "positive",
                "negative",
                "backend_mode",
                "edit_scope",
                "mask_grow",
                "mask_blur",
                "mask_threshold",
                "invert_mask",
                "crop_context",
                "crop_context_factor",
                "min_size",
                "max_size",
                "resize_mode",
                "expand_left",
                "expand_right",
                "expand_top",
                "expand_bottom",
                "canvas_fill",
                "auto_recommend",
            },
        )
        self.assertEqual(required["image"], ("IMAGE",))
        self.assertEqual(required["mask"], ("MASK",))
        self.assertEqual(required["vae"], ("VAE",))
        self.assertEqual(required["positive"], ("CONDITIONING",))
        self.assertEqual(required["negative"], ("CONDITIONING",))
        self.assertEqual(required["backend_mode"][0], ["auto", "sdxl", "flux"])
        self.assertEqual(required["edit_scope"][0], ["auto", "region", "crop", "canvas"])
        self.assertEqual(optional["model"], ("MODEL",))
        self.assertEqual(optional["model_info"], ("STRING",))

    def test_bridge_and_finish_schemas_match_contract(self):
        plugin = load_plugin_package()

        bridge_cls = plugin.NODE_CLASS_MAPPINGS["LLSProKSamplerBridge"]
        bridge_schema = bridge_cls.INPUT_TYPES()
        bridge_required = bridge_schema["required"]
        bridge_optional = bridge_schema.get("optional", {})
        self.assertEqual(bridge_cls.CATEGORY, "LLS/Image Edit")
        self.assertEqual(bridge_cls.FUNCTION, "sample")
        self.assertEqual(bridge_cls.RETURN_TYPES, ("LATENT",))
        self.assertEqual(bridge_cls.RETURN_NAMES, ("latent",))
        self.assertEqual(
            set(bridge_required),
            {
                "model",
                "add_noise",
                "noise_seed",
                "steps",
                "cfg",
                "sampler_name",
                "scheduler",
                "positive",
                "negative",
                "latent_image",
                "start_at_step",
                "end_at_step",
                "return_with_leftover_noise",
            },
        )
        self.assertEqual(bridge_required["add_noise"][0], ["enable", "disable"])
        self.assertEqual(bridge_required["return_with_leftover_noise"][0], ["disable", "enable"])
        self.assertEqual(bridge_optional, {})

        finish_cls = plugin.NODE_CLASS_MAPPINGS["LLSProImageEditFinish"]
        finish_schema = finish_cls.INPUT_TYPES()
        finish_required = finish_schema["required"]
        finish_optional = finish_schema["optional"]
        self.assertEqual(finish_cls.CATEGORY, "LLS/Image Edit")
        self.assertEqual(finish_cls.FUNCTION, "finish")
        self.assertEqual(finish_cls.RETURN_TYPES, ("IMAGE", "IMAGE"))
        self.assertEqual(finish_cls.RETURN_NAMES, ("final_image", "preview_image"))
        self.assertEqual(finish_required["original_image"], ("IMAGE",))
        self.assertEqual(finish_required["generated_image"], ("IMAGE",))
        self.assertEqual(finish_required["edit_info"], ("LLS_EDIT_INFO",))
        self.assertEqual(finish_optional["work_mask"], ("MASK",))
        self.assertEqual(finish_optional["sample_info"], ("STRING",))

    def test_backend_registry_exposes_built_in_scaffolds(self):
        plugin = load_plugin_package()
        registry = import_plugin_submodule(plugin, "pro_edit.backends.registry")

        self.assertEqual(registry.get_backend("sdxl").backend_name, "sdxl")
        self.assertEqual(registry.get_backend("flux").backend_name, "flux")

    def test_normalize_edit_info_accepts_jsonish_string_payloads(self):
        plugin = load_plugin_package()
        utils = import_plugin_submodule(plugin, "pro_edit.pro_edit_utils")

        info = utils.normalize_edit_info('{"backend_name":"sdxl","edit_scope":"crop"}')
        self.assertEqual(info["backend_name"], "sdxl")
        self.assertEqual(info["edit_scope"], "crop")
        self.assertEqual(info["edit_payload_version"], "1.0")

        legacy_info = utils.normalize_edit_info("backend_name=flux|edit_scope=canvas")
        self.assertEqual(legacy_info["backend_name"], "flux")
        self.assertEqual(legacy_info["edit_scope"], "canvas")
        self.assertEqual(legacy_info["edit_payload_version"], "1.0")


if __name__ == "__main__":
    unittest.main()