import unittest

import libcst as cst

from core.languages.python.cst_utils import any_ellipsis_in_cst, ellipsis_in_cst_not_types
from core.languages.python.testgen.generate import did_generate_ellipsis


class TestEllipsisInCst(unittest.TestCase):
    def test_ellipsis_in_cst_with_ellipsis(self) -> None:
        """Test that ellipsis_in_cst detects an ellipsis in the function body."""
        code = """
def example_function():
    return ...
    """
        module = cst.parse_module(code)
        assert ellipsis_in_cst_not_types(module) is True

    def test_ellipsis_in_cst_without_ellipsis(self) -> None:
        """Test that ellipsis_in_cst does not detect an ellipsis when there is none."""
        code = """
def example_function():
    return 42
    """
        module = cst.parse_module(code)
        assert ellipsis_in_cst_not_types(module) is False

    def test_ellipsis_in_cst_in_type_annotation(self) -> None:
        """Test that ellipsis_in_cst ignores ellipses in type annotations."""
        code = """
def example_function(param: tuple[int, ...]) -> None:
    pass
    """
        module = cst.parse_module(code)
        assert ellipsis_in_cst_not_types(module) is False

    def test_ellipsis_in_cst_in_type_annotation_union(self) -> None:
        """Test that ellipsis_in_cst ignores ellipses in type annotations with unions."""
        code = """
def example_function(param: tuple[int, ...] | list[int]) -> None:
    pass
    """
        module = cst.parse_module(code)
        assert ellipsis_in_cst_not_types(module) is False

    def test_ellipsis_in_cst(self) -> None:
        code = """def example_fn():
    return ..."""
        module = cst.parse_module(code)
        assert any_ellipsis_in_cst(module) is True

    def test_ellipsis_in_protocol_method_body_is_allowed(self) -> None:
        """Test that ellipsis as function body (Protocol/abstract methods) is allowed."""
        code = """
from typing import Protocol

class MyProtocol(Protocol):
    def method(self) -> None:
        ...

    async def async_method(self) -> str:
        ...
"""
        module = cst.parse_module(code)
        # Ellipsis as function body should NOT be flagged
        assert ellipsis_in_cst_not_types(module) is False

    def test_ellipsis_in_return_statement_is_flagged(self) -> None:
        """Test that ellipsis in return statement is still flagged."""
        code = """
def example_function():
    return ...
"""
        module = cst.parse_module(code)
        # Ellipsis as return value should be flagged
        assert ellipsis_in_cst_not_types(module) is True

    def test_ellipsis_mixed_valid_and_invalid(self) -> None:
        """Test that invalid ellipsis is flagged even when valid ones exist."""
        code = """
from typing import Protocol

class MyProtocol(Protocol):
    def method(self) -> None:
        ...

def other_function():
    return ...
"""
        module = cst.parse_module(code)
        # Should be flagged because of the return ...
        assert ellipsis_in_cst_not_types(module) is True


class TestContextDidGenerateEllipsis(unittest.TestCase):
    def test_single_context_flags_generated_ellipsis(self) -> None:
        """Single context should flag ellipsis when original had none."""
        source_code = "def foo():\n    return 42"
        generated_code = "def test_foo():\n    return ..."
        assert did_generate_ellipsis(source_code, generated_code) is True

    def test_single_context_allows_ellipsis_when_original_had_ellipsis(self) -> None:
        """Single context should not flag if original already had ellipsis."""
        source_code = "def foo():\n    ..."
        generated_code = "def test_foo():\n    return ..."
        assert did_generate_ellipsis(source_code, generated_code) is False

    def test_multi_context_flags_when_any_section_had_no_ellipsis(self) -> None:
        """Multi context should flag if ANY original section had no ellipsis.

        This tests the fix for using any() instead of all() - if one section
        had ellipsis but another didn't, we should still flag generated ellipsis.
        """
        # Section 1: has ellipsis (protocol method body - valid)
        # Section 2: no ellipsis
        source_code = """```python:module/protocol.py
from typing import Protocol

class MyProtocol(Protocol):
    def method(self) -> None:
        ...
```

```python:module/impl.py
def implementation():
    return 42
```"""
        generated_code = "def test_implementation():\n    return ..."
        # Should be flagged because module/impl.py had no ellipsis
        assert did_generate_ellipsis(source_code, generated_code) is True

    def test_multi_context_allows_when_all_sections_had_ellipsis(self) -> None:
        """Multi context should not flag if all original sections had ellipsis."""
        source_code = """```python:module/a.py
def foo():
    ...
```

```python:module/b.py
def bar():
    ...
```"""
        generated_code = "def test_foo():\n    return ..."
        # Should not be flagged because all original sections had ellipsis
        assert did_generate_ellipsis(source_code, generated_code) is False

    def test_multi_context_no_flag_when_generated_has_no_invalid_ellipsis(self) -> None:
        """Multi context should not flag if generated code has no invalid ellipsis."""
        source_code = """```python:module/impl.py
def implementation():
    return 42
```"""
        generated_code = "def test_implementation():\n    assert implementation() == 42"
        assert did_generate_ellipsis(source_code, generated_code) is False
