import importlib
import sys
import types
import unittest


class DumasH3LongVideosHelperTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls._saved_modules = {
            name: sys.modules.get(name)
            for name in (
                "torch",
                "nodes",
                "comfy",
                "comfy.utils",
                "comfy.samplers",
                "comfy.nested_tensor",
                "comfy.model_management",
                "node_helpers",
                "numpy",
                "PIL",
                "PIL.Image",
                "folder_paths",
                "dumas_image_nodes",
                "dumas_h3_longvideos",
            )
        }

        fake_torch = types.SimpleNamespace(
            cuda=types.SimpleNamespace(OutOfMemoryError=RuntimeError),
            float32="float32",
        )
        fake_numpy = types.SimpleNamespace(
            clip=lambda array, _low, _high: array,
            uint8="uint8",
        )
        fake_pil_image_module = types.SimpleNamespace(fromarray=lambda _array: None)
        fake_pil_module = types.SimpleNamespace(Image=fake_pil_image_module)
        fake_folder_paths = types.SimpleNamespace(
            get_temp_directory=lambda: "/tmp",
            get_output_directory=lambda: "/tmp",
            get_save_image_path=lambda prefix, _out, _width, _height: ("/tmp", prefix, 1, "", prefix),
        )
        fake_nodes = types.SimpleNamespace(common_ksampler=lambda *args, **kwargs: ({},))
        fake_comfy_samplers = types.SimpleNamespace(
            KSampler=types.SimpleNamespace(SAMPLERS=("res_multistep",), SCHEDULERS=("simple",))
        )
        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,
            samplers=fake_comfy_samplers,
            nested_tensor=types.SimpleNamespace(),
            model_management=fake_mm,
        )

        sys.modules["torch"] = fake_torch
        sys.modules["numpy"] = fake_numpy
        sys.modules["PIL"] = fake_pil_module
        sys.modules["PIL.Image"] = fake_pil_image_module
        sys.modules["folder_paths"] = fake_folder_paths
        sys.modules["nodes"] = fake_nodes
        sys.modules["comfy"] = fake_comfy
        sys.modules["comfy.utils"] = fake_comfy_utils
        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["node_helpers"] = types.SimpleNamespace()

        cls.image_module = importlib.import_module("dumas_image_nodes")
        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_extract_wardrobe_is_cached(self):
        fn = self.module.extract_wardrobe
        fn.cache_clear()

        beat = "walks forward\nwardrobe: red jacket, grey shorts\nlooks back"
        self.assertEqual(fn(beat), ("walks forward\nlooks back", "red jacket, grey shorts"))
        self.assertEqual(fn(beat), ("walks forward\nlooks back", "red jacket, grey shorts"))
        self.assertGreater(fn.cache_info().hits, 0)

    def test_dialogue_helpers_keep_existing_outputs_and_cache(self):
        spans_cache = self.module._dialogue_spans_cached
        sec_fn = self.module.dialogue_seconds
        words_fn = self.module.dialogue_words

        spans_cache.cache_clear()
        sec_fn.cache_clear()
        words_fn.cache_clear()

        beat = 'Mara says, "Open it now." Jon replies, "Do it."'
        self.assertEqual(self.module.dialogue_spans(beat), [3, 2])
        self.assertEqual(words_fn(beat), 5)
        self.assertAlmostEqual(sec_fn(beat), 3.5)
        self.assertAlmostEqual(sec_fn(beat, pad=False), 2.5)

        self.module.dialogue_spans(beat)
        sec_fn(beat)
        words_fn(beat)

        self.assertGreater(spans_cache.cache_info().hits, 0)
        self.assertGreater(sec_fn.cache_info().hits, 0)
        self.assertGreater(words_fn.cache_info().hits, 0)

    def test_directive_and_estimate_helpers_are_cached(self):
        directive_fn = self.module.beat_seconds_directive
        estimate_fn = self.module.estimate_beat_seconds
        action_fn = self.module.action_clauses

        directive_fn.cache_clear()
        estimate_fn.cache_clear()
        action_fn.cache_clear()

        beat = 'seconds: 7.5\nShe opens the hatch and climbs inside.'
        self.assertEqual(directive_fn(beat), 7.5)
        self.assertEqual(action_fn(beat), 2)
        self.assertAlmostEqual(estimate_fn(beat), 7.0)

        directive_fn(beat)
        action_fn(beat)
        estimate_fn(beat)

        self.assertGreater(directive_fn.cache_info().hits, 0)
        self.assertGreater(action_fn.cache_info().hits, 0)
        self.assertGreater(estimate_fn.cache_info().hits, 0)

    def test_has_speech_cache_respects_written_text_filter(self):
        fn = self.module.has_speech
        fn.cache_clear()

        written = 'She reads the sign marked "EXIT" and keeps walking.'
        spoken = 'She says, "Exit now." and points to the door.'
        self.assertFalse(fn(written))
        self.assertTrue(fn(spoken))
        fn(written)
        fn(spoken)
        self.assertGreaterEqual(fn.cache_info().hits, 2)

    def test_resolve_tagged_refs_preserves_sparse_socket_numbers(self):
        refs = [None, "img2", None, None, None, None, "img7", None, "img9"]

        text, images, dropped = self.module.resolve_tagged_refs(
            "Mara <Picture 7> turns toward Jon <Picture 2> while <Picture 9> watches.",
            refs,
        )

        self.assertEqual(
            text,
            "Mara <Picture 2> turns toward Jon <Picture 1> while <Picture 3> watches.",
        )
        self.assertEqual(images, ["img2", "img7", "img9"])
        self.assertEqual(dropped, [])

    def test_resolve_tagged_refs_drops_unconnected_sparse_slots(self):
        refs = [None, "img2", None, None, None, None, "img7", None, None]

        text, images, dropped = self.module.resolve_tagged_refs(
            "Use <Picture 7>, skip <Picture 4>, keep <Picture 2>.",
            refs,
        )

        self.assertEqual(text, "Use <Picture 2>, skip, keep <Picture 1>.")
        self.assertEqual(images, ["img2", "img7"])
        self.assertEqual(dropped, [4])

    def test_shot_references_uses_all_connected_sparse_slots(self):
        refs = [None, "img2", None, "img4", None, None, "img7", None, None]

        self.assertEqual(
            self.module.shot_references(refs, "auto ref2v", 0, None),
            ["img2", "img4", "img7"],
        )
        self.assertEqual(
            self.module.shot_references(refs, "first shot", 0, None),
            ["img2", "img4", "img7"],
        )
        self.assertEqual(
            self.module.shot_references(refs, "every shot", 3, None),
            ["img2", "img4", "img7"],
        )

    def test_input_types_expose_nine_ref_slots(self):
        optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]

        for index in range(1, 10):
            self.assertIn(f"ref_image_{index}", optional)
        self.assertIn("plan", optional)
        self.assertIn("plan_scene_index", optional)

    def test_plan_scene_refs_reads_bound_images(self):
        attach = self.image_module.DumasH3PlanAttachSceneImagesNode()
        plan = {"shots": [{"id": "one"}, {"id": "two"}]}
        image2 = object()
        image7 = object()

        plan, _ = attach.attach(plan=plan, scene_index=2, image2=image2, image7=image7)
        refs = self.module._plan_scene_refs(plan, 2)

        self.assertEqual(len(refs), 9)
        self.assertIsNone(refs[0])
        self.assertIs(refs[1], image2)
        self.assertIs(refs[6], image7)

    def test_merge_ref_slots_prefers_direct_refs_over_plan_refs(self):
        merged = self.module._merge_ref_slots(
            (None, "direct2", None, None, "direct5", None, None, None, None),
            ("plan1", "plan2", "plan3", None, "plan5", None, "plan7", None, None),
        )

        self.assertEqual(
            merged,
            ("plan1", "direct2", "plan3", None, "direct5", None, "plan7", None, None),
        )

    def test_ref_mode_defaults_are_ref2v_biased(self):
        optional = self.module.H3LongVideos.INPUT_TYPES()["optional"]

        self.assertEqual(optional["ref_mode"][1]["default"], "auto ref2v")
        self.assertEqual(optional["ref_noise_aug"][1]["default"], 0.95)

    def test_only_canonical_h3_long_videos_node_is_exposed(self):
        self.assertEqual(
            self.module.NODE_CLASS_MAPPINGS,
            {"DumasH3LongVideos": self.module.H3LongVideos},
        )
        self.assertEqual(
            self.module.NODE_DISPLAY_NAME_MAPPINGS,
            {"DumasH3LongVideos": "Dumas H3 Long Videos (FL2VA + REF2VA)"},
        )

    def test_compose_persistent_does_not_expand_ambiguous_plural_to_full_cast(self):
        active = self.module.parse_wardrobe(
            "Maya = she, red jacket\n"
            "Jon = he, navy overalls\n"
            "Becca = she, green coat"
        )

        shot = self.module.compose_persistent(
            "Both of them walk to the door.",
            active,
            "",
            speaking=False,
        )

        self.assertEqual(shot, "Both of them walk to the door.")

    def test_compose_persistent_keeps_two_person_plural_binding(self):
        active = self.module.parse_wardrobe(
            "Maya = she, red jacket\n"
            "Jon = he, navy overalls"
        )

        shot = self.module.compose_persistent(
            "They walk to the door.",
            active,
            "",
            speaking=False,
        )

        self.assertIn("Maya (red jacket)", shot)
        self.assertIn("Jon (navy overalls)", shot)
        self.assertIn("They walk to the door.", shot)

    def test_compose_persistent_all_three_characters_binds_full_cast(self):
        active = self.module.parse_wardrobe(
            "Maya = she, red jacket\n"
            "Jon = he, navy overalls\n"
            "Becca = she, green coat"
        )

        shot = self.module.compose_persistent(
            "The three characters walk to the door.",
            active,
            "",
            speaking=False,
        )

        self.assertIn("Maya (red jacket)", shot)
        self.assertIn("Jon (navy overalls)", shot)
        self.assertIn("Becca (green coat)", shot)


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