"""Tests for import gathering and addition."""

from __future__ import annotations

import textwrap

import libcst as cst

from codeflash_python.context.imports import (
    DottedImportCollector,
    FutureAliasedImportTransformer,
    add_needed_imports_from_module,
    gather_source_imports,
)


class TestDottedImportCollector:
    """Tests for DottedImportCollector."""

    def test_bare_import(self):
        """
        Bare import is collected as the module name.
        """
        module = cst.parse_module("import os\n")
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os" in collector.imports

    def test_dotted_import(self):
        """
        Dotted import is collected as the full dotted name.
        """
        module = cst.parse_module("import os.path\n")
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os.path" in collector.imports

    def test_aliased_bare_import(self):
        """
        Aliased bare import uses the alias in the name.
        """
        module = cst.parse_module("import numpy as np\n")
        collector = DottedImportCollector()
        module.visit(collector)
        assert "numpy.np" in collector.imports

    def test_from_import(self):
        """
        From-import is collected as module.name.
        """
        module = cst.parse_module(
            "from pathlib import Path\n",
        )
        collector = DottedImportCollector()
        module.visit(collector)
        assert "pathlib.Path" in collector.imports

    def test_aliased_from_import(self):
        """
        Aliased from-import uses the alias name.
        """
        module = cst.parse_module(
            "from os.path import join as pjoin\n",
        )
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os.path.pjoin" in collector.imports

    def test_star_import_skipped(self):
        """
        Star imports are not collected.
        """
        module = cst.parse_module(
            "from os.path import *\n",
        )
        collector = DottedImportCollector()
        module.visit(collector)
        assert collector.imports == set()

    def test_function_level_import_skipped(self):
        """
        Imports inside functions are not collected.
        """
        code = textwrap.dedent("""\
            def f():
                import os
        """)
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert collector.imports == set()

    def test_conditional_imports_collected(self):
        """
        Imports inside if blocks are collected.
        """
        code = textwrap.dedent("""\
            import sys
            if sys.version_info >= (3, 11):
                from typing import Self
        """)
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "sys" in collector.imports
        assert "typing.Self" in collector.imports

    def test_try_imports_collected(self):
        """
        Imports inside try blocks are collected.
        """
        code = textwrap.dedent("""\
            try:
                import ujson
            except ImportError:
                pass
        """)
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "ujson" in collector.imports

    def test_multiple_from_imports(self):
        """
        Multiple names from one from-import are all
        collected.
        """
        module = cst.parse_module(
            "from os.path import join, exists\n",
        )
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os.path.join" in collector.imports
        assert "os.path.exists" in collector.imports


class TestFutureAliasedImportTransformer:
    """Tests for FutureAliasedImportTransformer."""

    def test_removes_aliased_future(self):
        """
        Aliased __future__ imports are removed entirely.
        """
        code = "from __future__ import annotations as ann\n"
        module = cst.parse_module(code)
        result = module.visit(
            FutureAliasedImportTransformer(),
        )
        assert "annotations" not in result.code

    def test_keeps_non_aliased_future(self):
        """
        Non-aliased __future__ imports are preserved.
        """
        code = "from __future__ import annotations\n"
        module = cst.parse_module(code)
        result = module.visit(
            FutureAliasedImportTransformer(),
        )
        assert "annotations" in result.code

    def test_mixed_future_imports(self):
        """
        Mixed: keeps non-aliased, removes aliased.
        """
        code = "from __future__ import annotations, division as div\n"
        module = cst.parse_module(code)
        result = module.visit(
            FutureAliasedImportTransformer(),
        )
        assert "annotations" in result.code
        assert "division" not in result.code


class TestGatherSourceImports:
    """Tests for gather_source_imports."""

    def test_gathers_module_imports(self, tmp_path):
        """
        Module-level imports are gathered from source.
        """
        src = tmp_path / "src.py"
        src.write_text(
            "import os\nfrom pathlib import Path\n",
        )
        result = gather_source_imports(
            src.read_text(),
            src,
            tmp_path,
        )
        assert result is not None

    def test_returns_none_when_no_imports(self, tmp_path):
        """
        Returns None when source has no imports.
        """
        src = tmp_path / "src.py"
        src.write_text("x = 1\n")
        result = gather_source_imports(
            src.read_text(),
            src,
            tmp_path,
        )
        assert result is None

    def test_ignores_function_level_imports(
        self,
        tmp_path,
    ):
        """
        Imports inside functions are not gathered.
        """
        src = tmp_path / "src.py"
        src.write_text(
            textwrap.dedent("""\
            def f():
                import os
                return os.getcwd()
        """)
        )
        result = gather_source_imports(
            src.read_text(),
            src,
            tmp_path,
        )
        assert result is None


class TestAddNeededImportsFromModule:
    """Tests for add_needed_imports_from_module."""

    def test_adds_missing_import(self, tmp_path):
        """
        Missing import from source is added to destination.
        """
        src = tmp_path / "src.py"
        src.write_text(
            textwrap.dedent("""\
            import math

            def f():
                return math.sqrt(2)
        """)
        )
        dst = tmp_path / "dst.py"
        dst.write_text(
            textwrap.dedent("""\
            def f():
                return math.sqrt(2)
        """)
        )

        result = add_needed_imports_from_module(
            src_module_code=src.read_text(),
            dst_module_code=dst.read_text(),
            src_path=src,
            dst_path=dst,
            project_root=tmp_path,
        )
        assert "import math" in result
        assert "def f" in result

    def test_skips_helper_fqns(self, tmp_path):
        """
        Imports for helper functions are skipped.
        """
        src = tmp_path / "src.py"
        src.write_text(
            textwrap.dedent("""\
            from utils import helper

            def f():
                return helper()
        """)
        )
        dst = tmp_path / "dst.py"
        dst.write_text(
            textwrap.dedent("""\
            def f():
                return helper()
        """)
        )

        result = add_needed_imports_from_module(
            src_module_code=src.read_text(),
            dst_module_code=dst.read_text(),
            src_path=src,
            dst_path=dst,
            project_root=tmp_path,
            helper_fqns={"utils.helper"},
        )
        assert "from utils import helper" not in result

    def test_skips_existing_import(self, tmp_path):
        """
        Already-present imports are not duplicated.
        """
        src = tmp_path / "src.py"
        src.write_text(
            textwrap.dedent("""\
            import math

            def f():
                return math.sqrt(2)
        """)
        )
        dst = tmp_path / "dst.py"
        dst.write_text(
            textwrap.dedent("""\
            import math

            def f():
                return math.sqrt(2)
        """)
        )

        result = add_needed_imports_from_module(
            src_module_code=src.read_text(),
            dst_module_code=dst.read_text(),
            src_path=src,
            dst_path=dst,
            project_root=tmp_path,
        )
        assert result.count("import math") == 1

    def test_uses_precomputed_gatherer(self, tmp_path):
        """
        Pre-computed gatherer is reused without re-parsing.
        """
        src = tmp_path / "src.py"
        src.write_text(
            textwrap.dedent("""\
            import math

            def f():
                return math.sqrt(2)
        """)
        )
        dst = tmp_path / "dst.py"
        dst.write_text(
            textwrap.dedent("""\
            def f():
                return math.sqrt(2)
        """)
        )

        gatherer = gather_source_imports(
            src.read_text(),
            src,
            tmp_path,
        )
        result = add_needed_imports_from_module(
            src_module_code=src.read_text(),
            dst_module_code=dst.read_text(),
            src_path=src,
            dst_path=dst,
            project_root=tmp_path,
            gathered_imports=gatherer,
        )
        assert "import math" in result
