import json
import sys
import tempfile
import unittest
from io import StringIO
from pathlib import Path
from unittest.mock import patch

from comfyui_to_python import ComfyUItoPython, run


class UpscaleModelLoader:
    CATEGORY = "loaders"
    FUNCTION = "load_model"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model_name": ("STRING",),
            }
        }

    def load_model(self, model_name):
        return (model_name,)


class LoadImage:
    CATEGORY = "image"
    FUNCTION = "load_image"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("STRING",),
            }
        }

    def load_image(self, image):
        return (image,)


class ImageUpscaleWithModel:
    CATEGORY = "image/upscaling"
    FUNCTION = "upscale"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "upscale_model": ("UPSCALE_MODEL",),
                "image": ("IMAGE",),
            }
        }

    def upscale(self, upscale_model, image):
        return (image,)


class MergeImages:
    CATEGORY = "image"
    FUNCTION = "merge"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "left": ("IMAGE",),
                "right": ("IMAGE",),
            }
        }

    def merge(self, left, right):
        return (left, right)


class HiddenMetadataFilteredNode:
    CATEGORY = "image"
    FUNCTION = "save"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    def save(self, image):
        return (image,)


class HiddenPromptSeedNode:
    CATEGORY = "sampling"
    FUNCTION = "sample"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "seed": ("INT",),
                "noise_seed": ("INT",),
            },
            "hidden": {
                "prompt": "PROMPT",
            },
        }

    def sample(self, seed, noise_seed, prompt):
        return (seed, noise_seed, prompt)


class UpscaleModelLoaderExportTest(unittest.TestCase):
    def test_top_level_module_preserves_exporter_entrypoints(self):
        from comfyui_to_python import main

        self.assertTrue(callable(ComfyUItoPython))
        self.assertTrue(callable(run))
        self.assertTrue(callable(main))

    def test_export_defers_comfyui_bootstrap_until_main(self):
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "example.png",
                },
            }
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "LoadImage": LoadImage,
            },
        )

        generated = output.getvalue()

        self.assertIn("def bootstrap_comfyui_runtime()", generated)
        self.assertIn("def cleanup_comfyui_runtime(", generated)
        self.assertIn("import comfy.options", generated)
        self.assertIn("comfy.options.enable_args_parsing()", generated)
        self.assertIn("import cuda_malloc", generated)
        self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated)
        self.assertIn(
            "def main(unload_models: bool | None = None):\n"
            "    bootstrap_comfyui_runtime()\n"
            "    add_extra_model_paths()",
            generated,
        )
        self.assertLess(
            generated.index("def bootstrap_comfyui_runtime()"),
            generated.index("def main(unload_models: bool | None = None):"),
        )
        main_section = generated[
            generated.index("def main(unload_models: bool | None = None):") :
        ]
        self.assertIn(
            "def main(unload_models: bool | None = None):\n"
            "    bootstrap_comfyui_runtime()\n"
            "    add_extra_model_paths()",
            main_section,
        )
        self.assertLess(
            main_section.index("bootstrap_comfyui_runtime()"),
            main_section.index("add_extra_model_paths()"),
        )
        self.assertLess(
            main_section.index("add_extra_model_paths()"),
            main_section.index("import torch"),
        )
        self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
        self.assertIn(
            "    finally:\n        cleanup_comfyui_runtime(unload_models=unload_models)",
            main_section,
        )

    def test_generated_module_import_does_not_parse_cli_args(self):
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "example.png",
                },
            }
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "LoadImage": LoadImage,
            },
        )

        generated = output.getvalue()

        self.assertNotIn("\nimport torch\n", generated)

        globals_dict = {"__name__": "generated_workflow_module"}
        with patch.object(sys, "argv", ["generated_workflow.py", "--wrapper-flag"]):
            exec(generated, globals_dict)

        self.assertTrue(callable(globals_dict["main"]))

    def test_upscale_workflow_uses_direct_upscale_model_loader_init(self):
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "example.png",
                },
            },
            "2": {
                "class_type": "UpscaleModelLoader",
                "inputs": {
                    "model_name": "RealESRGAN_x4plus.safetensors",
                },
            },
            "3": {
                "class_type": "ImageUpscaleWithModel",
                "inputs": {
                    "upscale_model": ["2", 0],
                    "image": ["1", 0],
                },
            },
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "LoadImage": LoadImage,
                "UpscaleModelLoader": UpscaleModelLoader,
                "ImageUpscaleWithModel": ImageUpscaleWithModel,
            },
        )

        generated = output.getvalue()

        self.assertIn("    from nodes import NODE_CLASS_MAPPINGS", generated)
        self.assertIn(f"    from {LoadImage.__module__} import (", generated)
        self.assertIn("LoadImage,", generated)
        self.assertIn("UpscaleModelLoader,", generated)
        self.assertIn("ImageUpscaleWithModel,", generated)
        self.assertIn("loadimage = LoadImage()", generated)
        self.assertIn("upscalemodelloader = UpscaleModelLoader()", generated)
        self.assertIn(
            "imageupscalewithmodel_3 = imageupscalewithmodel.upscale(",
            generated,
        )
        self.assertNotIn('NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()', generated)

    def test_frontend_workflow_metadata_is_embedded_for_reimport(self):
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "example.png",
                },
            }
        }
        frontend_workflow = {
            "version": 0.4,
            "last_node_id": 1,
            "last_link_id": 0,
            "nodes": [],
            "links": [],
            "groups": [],
            "config": {},
            "extra": {},
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            frontend_workflow=json.dumps(frontend_workflow),
            output_file=output,
            node_class_mappings={
                "LoadImage": LoadImage,
            },
        )

        generated = output.getvalue()

        self.assertIn("def build_extra_pnginfo()", generated)
        self.assertIn('"workflow": {', generated)
        self.assertIn('"version": 0.4', generated)
        self.assertIn('"nodes": []', generated)
        self.assertNotIn('"workflow": json.loads(', generated)
        self.assertNotIn('"source": "workflow_api"', generated)

    def test_export_without_frontend_workflow_leaves_png_workflow_metadata_absent(self):
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "example.png",
                },
            }
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "LoadImage": LoadImage,
            },
        )

        generated = output.getvalue()

        self.assertIn("def build_extra_pnginfo()", generated)
        self.assertIn("return None", generated)
        self.assertNotIn('"workflow": json.loads(', generated)

    def test_export_preserves_unique_variable_names_for_subgraph_node_ids(self):
        workflow = {
            "1:23": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "left.png",
                },
            },
            "12:3": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "right.png",
                },
            },
            "20": {
                "class_type": "MergeImages",
                "inputs": {
                    "left": ["1:23", 0],
                    "right": ["12:3", 0],
                },
            },
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "LoadImage": LoadImage,
                "MergeImages": MergeImages,
            },
        )

        generated = output.getvalue()

        self.assertIn("loadimage_1_23 = loadimage.load_image(", generated)
        self.assertIn("loadimage_12_3 = loadimage.load_image(", generated)
        self.assertIn("mergeimages_20 = mergeimages.merge(", generated)
        self.assertIn("left=get_value_at_index(loadimage_1_23, 0)", generated)
        self.assertIn("right=get_value_at_index(loadimage_12_3, 0)", generated)

    def test_run_cli_export_leaves_png_workflow_metadata_absent(self):
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "example.png",
                },
            }
        }

        with tempfile.TemporaryDirectory() as tmpdir:
            tmpdir_path = Path(tmpdir)
            input_file = tmpdir_path / "workflow_api.json"
            output_file = tmpdir_path / "workflow_api.py"

            input_file.write_text(json.dumps(workflow), encoding="utf-8")

            with patch(
                "comfyui_to_python.get_node_class_mappings",
                return_value={"LoadImage": LoadImage},
            ), patch("comfyui_to_python.import_custom_nodes"):
                run(
                    input_file=str(input_file),
                    output_file=str(output_file),
                    queue_size=1,
                )

            generated = output_file.read_text(encoding="utf-8")

        self.assertIn("def build_extra_pnginfo()", generated)
        self.assertIn("return None", generated)
        self.assertNotIn('"workflow": json.loads(', generated)

    def test_export_structures_generated_script_into_readable_sections(self):
        workflow = {
            "1": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": "example.png",
                },
            },
            "2": {
                "class_type": "UpscaleModelLoader",
                "inputs": {
                    "model_name": "RealESRGAN_x4plus.safetensors",
                },
            },
            "3": {
                "class_type": "ImageUpscaleWithModel",
                "inputs": {
                    "upscale_model": ["2", 0],
                    "image": ["1", 0],
                },
            },
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "LoadImage": LoadImage,
                "UpscaleModelLoader": UpscaleModelLoader,
                "ImageUpscaleWithModel": ImageUpscaleWithModel,
            },
        )

        generated = output.getvalue()

        self.assertIn("# Imports", generated)
        self.assertIn("# Workflow data", generated)
        self.assertIn("# Workflow execution", generated)
        self.assertIn("# Entrypoint", generated)
        self.assertIn("def build_workflow()", generated)
        self.assertIn("def build_extra_pnginfo()", generated)
        self.assertIn("def main(unload_models: bool | None = None)", generated)
        self.assertIn("bootstrap_comfyui_runtime()", generated)
        self.assertIn("cleanup_comfyui_runtime(unload_models=unload_models)", generated)
        self.assertNotIn("def initialize_workflow()", generated)
        self.assertNotIn("def run_once(", generated)
        self.assertIn("with torch.inference_mode():", generated)
        self.assertIn("finally:", generated)
        self.assertIn("for q in range(1):", generated)
        self.assertIn("workflow = build_workflow()", generated)
        self.assertIn("extra_pnginfo = build_extra_pnginfo()", generated)
        self.assertNotIn('workflow = json.loads("', generated)
        self.assertLess(
            generated.index("def build_workflow()"),
            generated.index("def main(unload_models: bool | None = None)"),
        )

    def test_hidden_metadata_kwargs_follow_function_signature(self):
        workflow = {
            "1": {
                "class_type": "HiddenMetadataFilteredNode",
                "inputs": {
                    "image": "example.png",
                },
            }
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "HiddenMetadataFilteredNode": HiddenMetadataFilteredNode,
            },
        )

        generated = output.getvalue()

        self.assertIn("hiddenmetadatafilterednode_1 =", generated)
        self.assertNotIn("prompt=prompt", generated)
        self.assertNotIn("extra_pnginfo=extra_pnginfo", generated)

    def test_randomized_seed_inputs_update_prompt_metadata_before_execution(self):
        workflow = {
            "1": {
                "class_type": "HiddenPromptSeedNode",
                "inputs": {
                    "seed": 1,
                    "noise_seed": 2,
                },
            }
        }

        output = StringIO()
        ComfyUItoPython(
            workflow=json.dumps(workflow),
            output_file=output,
            node_class_mappings={
                "HiddenPromptSeedNode": HiddenPromptSeedNode,
            },
        )

        generated = output.getvalue()

        self.assertIn(
            'node_1_seed = prompt["1"]["inputs"]["seed"] = random.randint(1, 2**64)',
            generated,
        )
        self.assertIn(
            'node_1_noise_seed = prompt["1"]["inputs"]["noise_seed"] = random.randint(',
            generated,
        )
        self.assertIn("1, 2**64", generated)
        self.assertIn("seed=node_1_seed", generated)
        self.assertIn("noise_seed=node_1_noise_seed", generated)
        self.assertIn("prompt=prompt", generated)


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