"""Tests for CST pruning."""

from __future__ import annotations

import textwrap

import libcst as cst
import pytest

from codeflash_python.context.models import (
    CodeContextType,
    PruneConfig,
)
from codeflash_python.context.pruning import (
    maybe_strip_docstring,
    parse_code_and_prune_cst,
)


class TestParseCodeAndPruneCst:
    """Tests for parse_code_and_prune_cst."""

    def test_read_writable_keeps_target_and_deps(self):
        """
        READ_WRITABLE keeps the target function and used
        assignments, removes unrelated functions and imports.
        """
        code = textwrap.dedent("""\
            import os

            X = 42

            def target():
                return X

            def other():
                return 1
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.READ_WRITABLE,
            {"target"},
        )
        output = result.code
        assert "def target" in output
        assert "X = 42" in output
        assert "def other" not in output
        assert "import os" not in output

    def test_read_writable_keeps_class_init(self):
        """
        READ_WRITABLE keeps __init__ alongside the target
        method inside a class.
        """
        code = textwrap.dedent("""\
            class MyClass:
                def __init__(self):
                    self.x = 1

                def target(self):
                    return self.x

                def unrelated(self):
                    return 0
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.READ_WRITABLE,
            {"MyClass.target"},
        )
        output = result.code
        assert "def target" in output
        assert "def __init__" in output
        assert "def unrelated" not in output

    def test_read_only_excludes_target_keeps_helpers(self):
        """
        READ_ONLY excludes the target function but keeps
        helper functions.
        """
        code = textwrap.dedent("""\
            import os

            def target():
                return 1

            def helper():
                return 2
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.READ_ONLY,
            {"target"},
            {"helper"},
        )
        output = result.code
        assert "def target" not in output
        assert "def helper" in output
        assert "import os" not in output

    def test_read_only_keeps_dunder_methods(self):
        """
        READ_ONLY keeps dunder methods but excludes __init__.
        """
        code = textwrap.dedent("""\
            class MyClass:
                def __init__(self):
                    self.x = 1

                def __repr__(self):
                    return "MyClass"

                def target(self):
                    return self.x
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.READ_ONLY,
            {"MyClass.target"},
            set(),
        )
        output = result.code
        assert "def target" not in output
        assert "__repr__" in output
        assert "__init__" not in output

    def test_hashing_strips_docstrings(self):
        """
        HASHING always strips docstrings from output.
        """
        code = textwrap.dedent("""\
            def target():
                \"\"\"A docstring.\"\"\"
                return 1
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.HASHING,
            {"target"},
        )
        output = result.code
        assert "def target" in output
        assert "docstring" not in output

    def test_hashing_excludes_init_from_targets(self):
        """
        HASHING excludes __init__ even when it is a target.
        """
        code = textwrap.dedent("""\
            class MyClass:
                def __init__(self):
                    self.x = 1

                def target(self):
                    return self.x
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.HASHING,
            {"MyClass.__init__", "MyClass.target"},
        )
        output = result.code
        assert "def target" in output
        assert "__init__" not in output

    def test_testgen_keeps_helpers_and_all_dunders(self):
        """
        TESTGEN keeps target, helpers, and all dunder methods
        including __init__.
        """
        code = textwrap.dedent("""\
            class MyClass:
                def __init__(self):
                    self.x = 1

                def __repr__(self):
                    return "MyClass"

                def target(self):
                    return self.x

                def helper(self):
                    return 2

                def unrelated(self):
                    return 3
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.TESTGEN,
            {"MyClass.target"},
            {"MyClass.helper"},
        )
        output = result.code
        assert "def target" in output
        assert "def helper" in output
        assert "def __init__" in output
        assert "def __repr__" in output
        assert "def unrelated" not in output

    def test_raises_when_no_target_found(self):
        """
        ValueError is raised when the target function is not
        in the code.
        """
        code = "def other():\n    return 1\n"
        with pytest.raises(ValueError, match="No target"):
            parse_code_and_prune_cst(
                code,
                CodeContextType.READ_WRITABLE,
                {"nonexistent"},
            )

    def test_accepts_cst_module_input(self):
        """
        A pre-parsed cst.Module is accepted as input.
        """
        code = "def target():\n    return 1\n"
        module = cst.parse_module(code)
        result = parse_code_and_prune_cst(
            module,
            CodeContextType.READ_WRITABLE,
            {"target"},
        )
        assert "def target" in result.code

    def test_read_writable_keeps_dependency_class(self):
        """
        READ_WRITABLE keeps a class that is used by the
        target function via dependency analysis.
        """
        code = textwrap.dedent("""\
            class Config:
                x = 1

            def target():
                return Config.x
        """)
        result = parse_code_and_prune_cst(
            code,
            CodeContextType.READ_WRITABLE,
            {"target"},
        )
        output = result.code
        assert "class Config" in output
        assert "def target" in output


class TestMaybeStripDocstring:
    """Tests for maybe_strip_docstring."""

    def test_strips_when_configured(self):
        """
        Docstring is removed when remove_docstrings is True.
        """
        code = textwrap.dedent("""\
            def foo():
                \"\"\"A docstring.\"\"\"
                return 1
        """)
        module = cst.parse_module(code)
        func = module.body[0]
        assert isinstance(func, cst.FunctionDef)
        cfg = PruneConfig(remove_docstrings=True)
        result = maybe_strip_docstring(func, cfg)
        assert isinstance(result.body, cst.IndentedBlock)
        assert 1 == len(result.body.body)

    def test_preserves_when_not_configured(self):
        """
        Docstring is kept when remove_docstrings is False.
        """
        code = textwrap.dedent("""\
            def foo():
                \"\"\"A docstring.\"\"\"
                return 1
        """)
        module = cst.parse_module(code)
        func = module.body[0]
        assert isinstance(func, cst.FunctionDef)
        cfg = PruneConfig(remove_docstrings=False)
        result = maybe_strip_docstring(func, cfg)
        assert isinstance(result.body, cst.IndentedBlock)
        assert 2 == len(result.body.body)

    def test_replaces_only_docstring_with_pass(self):
        """
        A function with only a docstring gets a pass
        statement after stripping.
        """
        code = 'def foo():\n    """Only docstring."""\n'
        module = cst.parse_module(code)
        func = module.body[0]
        assert isinstance(func, cst.FunctionDef)
        cfg = PruneConfig(remove_docstrings=True)
        result = maybe_strip_docstring(func, cfg)
        assert isinstance(result.body, cst.IndentedBlock)
        assert 1 == len(result.body.body)
        stmt = result.body.body[0]
        assert isinstance(stmt, cst.SimpleStatementLine)
        assert isinstance(stmt.body[0], cst.Pass)
