import importlib
import sys
import types
import unittest


class DumasH3LongVideosUpstreamWrapperTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls._saved_modules = {
            name: sys.modules.get(name)
            for name in (
                "torch",
                "nodes",
                "comfy",
                "comfy.utils",
                "comfy.sample",
                "comfy.samplers",
                "comfy.nested_tensor",
                "comfy.model_management",
                "latent_preview",
                "node_helpers",
                "folder_paths",
                "dumas_h3_longvideos",
                "dumas_h3_longvideos_upstream",
            )
        }

        fake_torch = types.SimpleNamespace(
            cuda=types.SimpleNamespace(OutOfMemoryError=RuntimeError),
            float32="float32",
            float16="float16",
            bfloat16="bfloat16",
            zeros=lambda *args, **kwargs: None,
            empty=lambda *args, **kwargs: None,
            cat=lambda *args, **kwargs: None,
            stack=lambda *args, **kwargs: None,
            tensor=lambda *args, **kwargs: None,
            no_grad=lambda: _NullContext(),
            inference_mode=lambda: _NullContext(),
        )
        fake_nodes = types.SimpleNamespace(
            NODE_CLASS_MAPPINGS={},
            common_ksampler=lambda *args, **kwargs: ({},),
        )
        fake_comfy_samplers = types.SimpleNamespace(
            KSampler=types.SimpleNamespace(
                SAMPLERS=("res_multistep", "euler"),
                SCHEDULERS=("simple", "normal"),
            )
        )
        fake_comfy_utils = types.SimpleNamespace(ProgressBar=lambda total: None)
        fake_mm = types.SimpleNamespace(
            current_loaded_models=[],
            free_memory=lambda *args, **kwargs: None,
            get_torch_device=lambda: "cpu",
            soft_empty_cache=lambda *args, **kwargs: None,
            unload_all_models=lambda *args, **kwargs: None,
            get_free_memory=lambda *args, **kwargs: 0,
            get_total_memory=lambda *args, **kwargs: 0,
        )
        fake_comfy = types.SimpleNamespace(
            utils=fake_comfy_utils,
            sample=types.SimpleNamespace(),
            samplers=fake_comfy_samplers,
            nested_tensor=types.SimpleNamespace(),
            model_management=fake_mm,
        )

        sys.modules["torch"] = fake_torch
        sys.modules["nodes"] = fake_nodes
        sys.modules["comfy"] = fake_comfy
        sys.modules["comfy.utils"] = fake_comfy_utils
        sys.modules["comfy.sample"] = fake_comfy.sample
        sys.modules["comfy.samplers"] = fake_comfy_samplers
        sys.modules["comfy.nested_tensor"] = fake_comfy.nested_tensor
        sys.modules["comfy.model_management"] = fake_mm
        sys.modules["latent_preview"] = types.SimpleNamespace()
        sys.modules["node_helpers"] = types.SimpleNamespace()
        sys.modules["folder_paths"] = types.SimpleNamespace(
            get_folder_paths=lambda name: [],
            get_filename_list=lambda name: [],
            get_full_path=lambda name, filename: None,
            get_temp_directory=lambda: "/tmp",
            get_output_directory=lambda: "/tmp",
            models_dir="/tmp",
        )

        cls.module = importlib.import_module("dumas_h3_longvideos")

    @classmethod
    def tearDownClass(cls):
        for name, module in cls._saved_modules.items():
            if module is None:
                sys.modules.pop(name, None)
            else:
                sys.modules[name] = module

    def test_dumas_key_wraps_upstream_node(self):
        mappings = self.module.NODE_CLASS_MAPPINGS

        self.assertIs(mappings["DumasH3LongVideos"], self.module.H3LongVideos)
        self.assertIs(mappings["H3LongVideos"], self.module.H3LongVideos)
        self.assertIs(mappings["H3LongVideosREF2VA"], self.module.H3LongVideos)

    def test_upstream_schema_is_exposed_under_dumas_key(self):
        node_cls = self.module.NODE_CLASS_MAPPINGS["DumasH3LongVideos"]
        schema = node_cls.INPUT_TYPES()

        self.assertIn("prompt", schema["required"])
        self.assertTrue(schema["required"]["prompt"][1]["forceInput"])
        self.assertIn("first_frame", schema["optional"])
        self.assertIn("ref_image_1", schema["optional"])
        self.assertIn("latent_upscale", schema["optional"])
        self.assertEqual(node_cls.RETURN_NAMES[0:4], ("images", "audio", "info", "script"))


class _NullContext:
    def __enter__(self):
        return None

    def __exit__(self, *_exc):
        return False


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