"""Tests for hardened import path resolution and importlib isolation layer.

Covers Approach B from docs/specs/harden-import-path-resolution.html:
- _is_comfyui_directory() structural verification
- _load_module() centralized importlib isolation
- get_comfyui_path() multi-strategy resolution
- sys.path idempotence (no remove/re-insert gap)
- Security: shadowing attack resistance
- Logging hygiene (no print() path leaks)
"""

import os
import sys
import tempfile
import unittest
from unittest.mock import patch

# Real ComfyUI checkout for integration-style unit tests (parent of this repo)
_REAL_COMFYUI = os.path.realpath(
    os.path.join(os.path.dirname(__file__), "..", "..", "ComfyUI")
)


class TestLoadModule(unittest.TestCase):
    """Tests for _load_module() centralized importlib isolation."""

    def tearDown(self):
        # Clean up sys.modules cache between tests to avoid cross-test pollution
        mod_names_to_clean = [
            name for name in sys.modules if name.startswith("_test_mod_")
        ]
        for name in mod_names_to_clean:
            del sys.modules[name]

    def test_loads_module_from_explicit_file_path(self):
        """Given a valid .py file, loads the module and returns it with expected attributes."""
        from comfyui_to_python.node_runtime import _load_module

        # Use a simple self-contained module (nodes.py requires torch)
        with tempfile.TemporaryDirectory() as tmpdir:
            mod_file = os.path.join(tmpdir, "mymodule.py")
            with open(mod_file, "w") as f:
                f.write("VALUE = 42\nNAME = 'test_module'\n")

            mod = _load_module("_test_mod_simple", mod_file)
        self.assertIsNotNone(mod)
        self.assertEqual(mod.VALUE, 42)
        self.assertEqual(mod.NAME, "test_module")

    def test_returns_none_for_missing_file(self):
        """Given a nonexistent file path, returns None without raising."""
        from comfyui_to_python.node_runtime import _load_module

        result = _load_module("_test_mod_missing", "/nonexistent/path/module.py")
        self.assertIsNone(result)


class TestShadowingResistance(unittest.TestCase):
    """Security: _load_module() must ignore shadowed modules at sys.path[0]."""

    def tearDown(self):
        # Clean up sys.modules cache between tests
        mod_names_to_clean = [
            name for name in sys.modules if name.startswith("_sec_test_")
        ]
        for name in mod_names_to_clean:
            del sys.modules[name]

    def test_ignores_shadowed_main_at_syspath(self):
        """Given a malicious main.py at sys.path[0], _load_module() loads from verified path."""
        from comfyui_to_python.node_runtime import _load_module

        with tempfile.TemporaryDirectory() as tmpdir:
            # Legitimate module
            legit_file = os.path.join(tmpdir, "legit_main.py")
            with open(legit_file, "w") as f:
                f.write("ORIGIN = 'legitimate'\n")

            # Malicious shadow at sys.path[0]
            attack_dir = os.path.join(tmpdir, "attack")
            os.makedirs(attack_dir)
            attack_file = os.path.join(attack_dir, "main.py")
            with open(attack_file, "w") as f:
                f.write("ORIGIN = 'malicious'\n")

            # Insert malicious path at sys.path[0]
            old_path = sys.path[:]
            try:
                sys.path.insert(0, attack_dir)
                # _load_module with explicit legit path must return legitimate module
                mod = _load_module("_sec_test_main", legit_file)
            finally:
                sys.path[:] = old_path

        self.assertIsNotNone(mod)
        self.assertEqual(mod.ORIGIN, "legitimate")

    def test_ignores_shadowed_nodes_at_syspath(self):
        """Given a malicious nodes.py at sys.path[0], _load_module() loads from verified path."""
        from comfyui_to_python.node_runtime import _load_module

        with tempfile.TemporaryDirectory() as tmpdir:
            # Legitimate module
            legit_file = os.path.join(tmpdir, "legit_nodes.py")
            with open(legit_file, "w") as f:
                f.write("NODE_CLASS_MAPPINGS = {'RealNode': 'real'}\n")

            # Malicious shadow at sys.path[0]
            attack_dir = os.path.join(tmpdir, "attack")
            os.makedirs(attack_dir)
            attack_file = os.path.join(attack_dir, "nodes.py")
            with open(attack_file, "w") as f:
                f.write("NODE_CLASS_MAPPINGS = {'FakeNode': 'malicious'}\n")

            old_path = sys.path[:]
            try:
                sys.path.insert(0, attack_dir)
                mod = _load_module("_sec_test_nodes", legit_file)
            finally:
                sys.path[:] = old_path

        self.assertIsNotNone(mod)
        self.assertEqual(mod.NODE_CLASS_MAPPINGS, {"RealNode": "real"})


class TestSysPathIdempotence(unittest.TestCase):
    """Tests for sys.path idempotence — no remove/re-insert gap."""

    def test_add_comfyui_idempotent_no_remove_insert(self):
        """Repeated calls to add_comfyui_directory_to_sys_path() are insert-once only."""
        from comfyui_to_python.node_runtime import (
            add_comfyui_directory_to_sys_path,
        )

        with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
            before_length = len(sys.path)
            add_comfyui_directory_to_sys_path()
            after_first = len(sys.path)

            # Second call — must not change sys.path at all (no remove+re-insert)
            add_comfyui_directory_to_sys_path()
            after_second = len(sys.path)

        self.assertEqual(after_first, after_second, "sys.path changed on second call")
        # ComfyUI should be on sys.path
        self.assertIn(_REAL_COMFYUI, sys.path)

    def test_import_custom_nodes_no_gap_window(self):
        """import_custom_nodes() must never remove ComfyUI from sys.path during execution."""
        import inspect
        from comfyui_to_python.node_runtime import import_custom_nodes

        source = inspect.getsource(import_custom_nodes)
        self.assertNotIn(
            "sys.path.remove",
            source,
            "import_custom_nodes() must not remove ComfyUI from sys.path (gap window)",
        )

    def test_no_print_statements_in_node_runtime(self):
        """node_runtime.py must use logging, not print(), to avoid path leaks to stdout."""
        import inspect
        from comfyui_to_python import node_runtime

        source = inspect.getsource(node_runtime)
        self.assertNotIn(
            "print(",
            source,
            "node_runtime.py must not use print() — use logging.DEBUG instead",
        )


class TestGeneratedScriptIsolation(unittest.TestCase):
    """Tests for generated script bootstrap (Phase 3)."""

    def test_generated_script_embeds_load_module(self):
        """render.py must embed _load_module() in generated scripts."""
        from comfyui_to_python.generator.render import WorkflowRenderer
        from comfyui_to_python.generator.model import GenerationPlan

        plan = GenerationPlan(
            workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
            metadata_workflow_data=None,
            custom_nodes=False,
            import_statements={},
            special_functions_code=[],
            loop_code=["result = some_node()"],
            queue_size=1,
        )
        renderer = WorkflowRenderer()
        generated = renderer.render(plan)

        self.assertIn(
            "_load_module", generated, "Generated script must embed _load_module()"
        )
        self.assertIn(
            "spec_from_file_location",
            generated,
            "Generated script must use importlib isolation",
        )

    def test_generated_script_no_bare_comfyui_imports(self):
        """Generated scripts must have zero bare imports of ComfyUI internals."""
        from comfyui_to_python.generator.render import WorkflowRenderer
        from comfyui_to_python.generator.model import GenerationPlan

        plan = GenerationPlan(
            workflow_data={"1": {"class_type": "CheckpointLoaderSimple"}},
            metadata_workflow_data=None,
            custom_nodes=True,
            import_statements={},
            special_functions_code=[],
            loop_code=["result = some_node()"],
            queue_size=1,
        )
        renderer = WorkflowRenderer()
        generated = renderer.render(plan)

        # These bare imports must NOT appear in generated output
        bare_imports = [
            "from main import",
            "from utils.extra_config import",
            "import execution",
            "import server",
            "from nodes import",
            "import comfy.options",
            "from comfy.cli_args import",
        ]
        for bare in bare_imports:
            self.assertNotIn(
                bare,
                generated,
                f"Generated script must not contain bare import: {bare}",
            )


class TestGetComfyuiPath(unittest.TestCase):
    """Tests for get_comfyui_path() multi-strategy resolution."""

    def test_resolves_via_valid_comfyui_path_env(self):
        """Given COMFYUI_PATH env var set to valid ComfyUI dir, returns that path."""
        from comfyui_to_python.node_runtime import get_comfyui_path

        with patch.dict("os.environ", {"COMFYUI_PATH": _REAL_COMFYUI}, clear=False):
            result = get_comfyui_path()
        self.assertEqual(result, _REAL_COMFYUI)

    def test_falls_through_to_next_strategy_when_env_invalid(self):
        """Given invalid COMFYUI_PATH (no nodes.py), rejects and falls through."""
        from comfyui_to_python.node_runtime import get_comfyui_path

        with patch.dict(
            "os.environ", {"COMFYUI_PATH": "/tmp/not-comfyui"}, clear=False
        ):
            result = get_comfyui_path()
        # The invalid path must be rejected (not returned as-is)
        self.assertNotEqual(result, "/tmp/not-comfyui")
        # In standalone/dev setups, no strategy may succeed -> None is valid
        # In ComfyUI/custom_nodes/ install, relative strategy would find it


class TestIsComfyuiDirectory(unittest.TestCase):
    """Tests for _is_comfyui_directory() structural verification."""

    def test_returns_true_for_real_comfyui_checkout(self):
        """Given a real ComfyUI checkout with nodes.py, returns True."""
        from comfyui_to_python.node_runtime import _is_comfyui_directory

        self.assertTrue(_is_comfyui_directory(_REAL_COMFYUI))

    def test_rejects_directory_lacking_markers(self):
        """Given a directory without nodes.py, returns False."""
        from comfyui_to_python.node_runtime import _is_comfyui_directory

        with tempfile.TemporaryDirectory() as tmpdir:
            # Empty directory has no ComfyUI markers
            self.assertFalse(_is_comfyui_directory(tmpdir))

    def test_rejects_nonexistent_path(self):
        """Given a nonexistent path, returns False."""
        from comfyui_to_python.node_runtime import _is_comfyui_directory

        self.assertFalse(_is_comfyui_directory("/nonexistent/path/to/comfyui"))


class TestLoadModuleFailureCleanup(unittest.TestCase):
    """Tests for _load_module() behavior when exec_module() raises."""

    def tearDown(self):
        mod_names_to_clean = [
            name for name in sys.modules if name.startswith("_test_fail_")
        ]
        for name in mod_names_to_clean:
            del sys.modules[name]

    def test_broken_module_removed_from_sys_modules_on_failure(self):
        """Given a module that raises during exec, it is removed from sys.modules."""
        from comfyui_to_python.node_runtime import _load_module

        with tempfile.TemporaryDirectory() as tmpdir:
            mod_file = os.path.join(tmpdir, "badmodule.py")
            with open(mod_file, "w") as f:
                f.write("raise RuntimeError('mid-execution failure')\n")

            result = _load_module("_test_fail_bad", mod_file)
        self.assertIsNone(result)
        self.assertNotIn(
            "_test_fail_bad",
            sys.modules,
            "Failed module must not remain cached in sys.modules",
        )

    def test_broken_module_retried_cleanly(self):
        """Given a failed load, a subsequent load attempt starts fresh (not from cache)."""
        from comfyui_to_python.node_runtime import _load_module

        with tempfile.TemporaryDirectory() as tmpdir:
            bad_file = os.path.join(tmpdir, "bad.py")
            good_file = os.path.join(tmpdir, "good.py")

            # First load: bad module that raises
            with open(bad_file, "w") as f:
                f.write("raise RuntimeError('fail')\n")
            result_bad = _load_module("_test_fail_retry", bad_file)
            self.assertIsNone(result_bad)

            # Overwrite with good module
            os.remove(bad_file)
            with open(good_file, "w") as f:
                f.write("VALUE = 99\n")

            # Reload same name from a different (working) file
            result_good = _load_module("_test_fail_retry", good_file)
        self.assertIsNotNone(result_good)
        self.assertEqual(result_good.VALUE, 99)


class TestLoadModuleTemp(unittest.TestCase):
    """Tests for _load_module_temp() removal from sys.modules."""

    def tearDown(self):
        mod_names_to_clean = [
            name for name in sys.modules if name.startswith("_test_temp_")
        ]
        for name in mod_names_to_clean:
            del sys.modules[name]

    def test_module_removed_from_sys_modules_after_temp_load(self):
        """After _load_module_temp(), the module key is absent from sys.modules."""
        from comfyui_to_python.node_runtime import _load_module_temp

        with tempfile.TemporaryDirectory() as tmpdir:
            mod_file = os.path.join(tmpdir, "tempmod.py")
            with open(mod_file, "w") as f:
                f.write("VALUE = 123\n")

            mod = _load_module_temp("_test_temp_mod", mod_file)
        self.assertIsNotNone(mod)
        self.assertEqual(mod.VALUE, 123)
        self.assertNotIn(
            "_test_temp_mod",
            sys.modules,
            "Temp module must be removed from sys.modules after load",
        )

    def test_module_removed_on_temp_load_failure(self):
        """After _load_module_temp() fails, the module key is absent from sys.modules."""
        from comfyui_to_python.node_runtime import _load_module_temp

        with tempfile.TemporaryDirectory() as tmpdir:
            mod_file = os.path.join(tmpdir, "badtemp.py")
            with open(mod_file, "w") as f:
                f.write("raise RuntimeError('fail')\n")

            result = _load_module_temp("_test_temp_bad", mod_file)
        self.assertIsNone(result)
        self.assertNotIn(
            "_test_temp_bad",
            sys.modules,
            "Failed temp module must be removed from sys.modules",
        )


class TestRepeatedNodeLoad(unittest.TestCase):
    """Tests that repeated _load_module('nodes', ...) calls respect the cache."""

    def tearDown(self):
        if "_test_nodes" in sys.modules:
            del sys.modules["_test_nodes"]

    def test_repeated_load_preserves_cached_attributes(self):
        """Second _load_module() for same name returns cached module, not a fresh copy."""
        from comfyui_to_python.node_runtime import _load_module

        with tempfile.TemporaryDirectory() as tmpdir:
            mod_file = os.path.join(tmpdir, "mynodes.py")
            with open(mod_file, "w") as f:
                f.write("NODE_CLASS_MAPPINGS = {'A': 1}\n")

            first = _load_module("_test_nodes", mod_file)
            self.assertIsNotNone(first)

            # Mutate the cached module to simulate init_extra_nodes()
            first.NODE_CLASS_MAPPINGS["B"] = 2

            # Second load with same name must return cached copy, preserving mutation
            second = _load_module("_test_nodes", mod_file)
        self.assertIs(first, second)
        self.assertEqual(second.NODE_CLASS_MAPPINGS.get("B"), 2)


class TestGetComfyuiPathThirdStrategy(unittest.TestCase):
    """Tests for get_comfyui_path() third-strategy verification."""

    def test_fallback_rejects_non_comfyui_directory(self):
        """Given a directory named ComfyUI without nodes.py, third strategy rejects it."""
        from comfyui_to_python.node_runtime import get_comfyui_path

        with tempfile.TemporaryDirectory() as tmpdir:
            # Create a fake "ComfyUI" directory (no nodes.py)
            fake_dir = os.path.join(tmpdir, "ComfyUI")
            os.makedirs(fake_dir)

            original_cwd = os.getcwd()
            try:
                os.chdir(tmpdir)
                with patch.dict("os.environ", {}, clear=False):
                    # Strip COMFYUI_PATH so fallback strategies activate
                    env_copy = dict(os.environ)
                    env_copy.pop("COMFYUI_PATH", None)
                    with patch.dict("os.environ", env_copy, clear=True):
                        result = get_comfyui_path()
            finally:
                os.chdir(original_cwd)

        # Must NOT return the fake directory
        self.assertNotEqual(result, fake_dir)


class TestImportCustomNodesSysPathRestoration(unittest.TestCase):
    """Tests for sys.path restoration in import_custom_nodes()."""

    def test_sys_path_restored_on_crash(self):
        """Given a crash during import_custom_nodes(), sys.path is still restored."""
        from comfyui_to_python.node_runtime import import_custom_nodes

        original_path = sys.path[:]

        with patch(
            "comfyui_to_python.node_runtime._load_module",
            side_effect=RuntimeError("simulated crash during server load"),
        ):
            try:
                import_custom_nodes()
            except RuntimeError:
                pass  # Expected crash

        # sys.path should still contain all original entries (no dangling state)
        for entry in original_path:
            self.assertIn(entry, sys.path)


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