from __future__ import annotations

import textwrap

from codeflash_python.analysis._normalizer import normalize_python_code


class TestBasicNormalization:
    """Tests for local variable renaming to canonical forms."""

    def test_single_variable_renamed(self):
        """
        A single local variable is renamed to var_0.
        """
        code = textwrap.dedent("""\
            def f():
                x = 1
                return x
        """)
        result = normalize_python_code(code)
        assert "var_0 = 1" in result
        assert "return var_0" in result

    def test_multiple_variables_renamed_sequentially(self):
        """
        Multiple local variables are renamed to var_0, var_1, etc. in order of first assignment.
        """
        code = textwrap.dedent("""\
            def f():
                alpha = 1
                beta = 2
                return alpha + beta
        """)
        result = normalize_python_code(code)
        assert "var_0 = 1" in result
        assert "var_1 = 2" in result
        assert "return var_0 + var_1" in result

    def test_same_variable_reused_gets_same_canonical_name(self):
        """
        A variable assigned and then loaded multiple times uses the same canonical name throughout.
        """
        code = textwrap.dedent("""\
            def f():
                total = 0
                total = total + 1
                return total
        """)
        result = normalize_python_code(code)
        assert "var_0 = 0" in result
        assert "var_0 = var_0 + 1" in result
        assert "return var_0" in result


class TestFunctionNamePreservation:
    """Tests that function and class names are not renamed."""

    def test_function_name_preserved(self):
        """
        The function name itself is not renamed.
        """
        code = textwrap.dedent("""\
            def compute_sum():
                x = 1
                return x
        """)
        result = normalize_python_code(code)
        assert "def compute_sum():" in result

    def test_class_name_preserved(self):
        """
        The class name is not renamed.
        """
        code = textwrap.dedent("""\
            class MyClass:
                def method(self):
                    x = 1
                    return x
        """)
        result = normalize_python_code(code)
        assert "class MyClass:" in result


class TestParameterPreservation:
    """Tests that function parameters are not renamed."""

    def test_positional_parameter_preserved(self):
        """
        Positional parameters keep their original names.
        """
        code = textwrap.dedent("""\
            def f(value):
                x = value + 1
                return x
        """)
        result = normalize_python_code(code)
        assert "value" in result
        assert "var_0 = value + 1" in result

    def test_args_and_kwargs_preserved(self):
        """
        *args and **kwargs parameters keep their original names.
        """
        code = textwrap.dedent("""\
            def f(*args, **kwargs):
                x = args
                y = kwargs
                return x, y
        """)
        result = normalize_python_code(code)
        assert "args" in result
        assert "kwargs" in result

    def test_keyword_only_parameter_preserved(self):
        """
        Keyword-only parameters keep their original names.
        """
        code = textwrap.dedent("""\
            def f(*, key):
                x = key
                return x
        """)
        result = normalize_python_code(code)
        assert "key" in result
        assert "var_0 = key" in result

    def test_self_parameter_preserved(self):
        """
        The self parameter on methods is not renamed.
        """
        code = textwrap.dedent("""\
            class C:
                def method(self):
                    x = self
                    return x
        """)
        result = normalize_python_code(code)
        assert "self" in result


class TestImportPreservation:
    """Tests that imported names are not renamed."""

    def test_import_name_preserved(self):
        """
        Names brought in by import statements are not renamed.
        """
        code = textwrap.dedent("""\
            import os
            def f():
                x = os.path.join('a', 'b')
                return x
        """)
        result = normalize_python_code(code)
        assert "os" in result

    def test_from_import_name_preserved(self):
        """
        Names brought in by from-import statements are not renamed.
        """
        code = textwrap.dedent("""\
            from os.path import join
            def f():
                x = join('a', 'b')
                return x
        """)
        result = normalize_python_code(code)
        assert "join" in result

    def test_aliased_import_preserved(self):
        """
        Aliased import names are preserved using the alias.
        """
        code = textwrap.dedent("""\
            import numpy as np
            def f():
                x = np.array([1])
                return x
        """)
        result = normalize_python_code(code)
        assert "np" in result


class TestDocstringRemoval:
    """Tests for docstring handling."""

    def test_docstring_removed_by_default(self):
        """
        Docstrings are removed when remove_docstrings=True (the default).
        """
        code = textwrap.dedent('''\
            def f():
                """This is a docstring."""
                x = 1
                return x
        ''')
        result = normalize_python_code(code)
        assert "docstring" not in result

    def test_module_docstring_removed(self):
        """
        Module-level docstrings are removed.
        """
        code = textwrap.dedent('''\
            """Module docstring."""
            def f():
                x = 1
                return x
        ''')
        result = normalize_python_code(code)
        assert "Module docstring" not in result

    def test_class_docstring_removed(self):
        """
        Class-level docstrings are removed.
        """
        code = textwrap.dedent('''\
            class C:
                """Class docstring."""
                def method(self):
                    x = 1
                    return x
        ''')
        result = normalize_python_code(code)
        assert "Class docstring" not in result


class TestDocstringPreservation:
    """Tests for docstring preservation when remove_docstrings=False."""

    def test_docstring_preserved_when_flag_false(self):
        """
        Docstrings are kept when remove_docstrings=False.
        """
        code = textwrap.dedent('''\
            def f():
                """This is a docstring."""
                x = 1
                return x
        ''')
        result = normalize_python_code(code, remove_docstrings=False)
        assert "This is a docstring." in result

    def test_all_docstrings_preserved_when_flag_false(self):
        """
        Module, class, and function docstrings are all kept when remove_docstrings=False.
        """
        code = textwrap.dedent('''\
            """Module doc."""
            class C:
                """Class doc."""
                def method(self):
                    """Method doc."""
                    x = 1
                    return x
        ''')
        result = normalize_python_code(code, remove_docstrings=False)
        assert "Module doc." in result
        assert "Class doc." in result
        assert "Method doc." in result


class TestSyntaxErrorHandling:
    """Tests that invalid Python input is returned unchanged."""

    def test_syntax_error_returns_original(self):
        """
        Code with a syntax error is returned as-is without modification.
        """
        code = "def f(\n    x = :\n"
        assert code == normalize_python_code(code)

    def test_incomplete_code_returns_original(self):
        """
        Incomplete code that cannot be parsed is returned unchanged.
        """
        code = "def f("
        assert code == normalize_python_code(code)


class TestEmptyInput:
    """Tests for empty and whitespace-only input."""

    def test_empty_string(self):
        """
        An empty string produces an empty string.
        """
        assert "" == normalize_python_code("")

    def test_whitespace_only(self):
        """
        Whitespace-only input normalizes to an empty string.
        """
        result = normalize_python_code("   \n\n  ")
        assert "" == result


class TestMultipleFunctions:
    """Tests for code containing multiple function definitions."""

    def test_variables_in_separate_functions_get_independent_names(self):
        """
        Local variables in different functions are normalized independently per scope.
        """
        code = textwrap.dedent("""\
            def first():
                alpha = 1
                return alpha

            def second():
                beta = 2
                return beta
        """)
        result = normalize_python_code(code)
        assert "def first():" in result
        assert "def second():" in result
        assert "var_0 = 1" in result
        assert "var_0 = 2" in result

    def test_nested_function_variables_normalized(self):
        """
        Variables in nested functions are normalized, inheriting the parent scope counter.
        """
        code = textwrap.dedent("""\
            def outer():
                x = 1
                def inner():
                    y = 2
                    return y
                return x
        """)
        result = normalize_python_code(code)
        assert "def outer():" in result
        assert "def inner():" in result
        assert "var_0 = 1" in result
        assert "var_1 = 2" in result
