"""Tests for Polars expression visitor."""

import re

import pytest

pl = pytest.importorskip("polars")

from therismos import F
from therismos.expr._expr import FALSE, TRUE, AllExpr, AnyExpr
from therismos.expr.visitors.polars import PolarsExprVisitor


@pytest.fixture()
def visitor() -> PolarsExprVisitor:
    """Return a PolarsExprVisitor instance."""
    return PolarsExprVisitor()


@pytest.fixture()
def df() -> "pl.DataFrame":
    """Return a sample DataFrame for integration tests."""
    return pl.DataFrame(
        {
            "age": [20, 15, 30, 25],
            "name": ["Alice", "Bob", "Charlie", "alice"],
            "score": [85.0, 72.0, 91.0, 60.0],
            "status": ["active", "inactive", "active", "active"],
            "tag": ["python", "java", "python", "rust"],
            "nullable": [1, None, 3, None],
        }
    )


class TestPolarsExprVisitorComparisons:
    """Tests for comparison expression conversions."""

    def test_eq_produces_polars_expr(self, visitor: PolarsExprVisitor) -> None:
        """Test that Eq produces a valid Polars expression."""
        expr = F("age") == 20
        result = expr.accept(visitor)
        assert isinstance(result, pl.Expr)

    def test_ne_produces_polars_expr(self, visitor: PolarsExprVisitor) -> None:
        """Test that Ne produces a valid Polars expression."""
        expr = F("age") != 20
        result = expr.accept(visitor)
        assert isinstance(result, pl.Expr)

    def test_lt_produces_polars_expr(self, visitor: PolarsExprVisitor) -> None:
        """Test that Lt produces a valid Polars expression."""
        expr = F("age") < 20
        result = expr.accept(visitor)
        assert isinstance(result, pl.Expr)

    def test_le_produces_polars_expr(self, visitor: PolarsExprVisitor) -> None:
        """Test that Le produces a valid Polars expression."""
        expr = F("age") <= 20
        result = expr.accept(visitor)
        assert isinstance(result, pl.Expr)

    def test_gt_produces_polars_expr(self, visitor: PolarsExprVisitor) -> None:
        """Test that Gt produces a valid Polars expression."""
        expr = F("age") > 20
        result = expr.accept(visitor)
        assert isinstance(result, pl.Expr)

    def test_ge_produces_polars_expr(self, visitor: PolarsExprVisitor) -> None:
        """Test that Ge produces a valid Polars expression."""
        expr = F("age") >= 20
        result = expr.accept(visitor)
        assert isinstance(result, pl.Expr)

    @pytest.mark.parametrize(
        "expr,expected_count",
        [
            (F("age") == 20, 1),
            (F("age") != 20, 3),
            (F("age") < 20, 1),
            (F("age") <= 20, 2),
            (F("age") > 20, 2),
            (F("age") >= 20, 3),
        ],
    )
    def test_comparison_filters(
        self, visitor: PolarsExprVisitor, df: "pl.DataFrame", expr: object, expected_count: int
    ) -> None:
        """Test that comparison expressions filter correctly."""
        pl_expr = expr.accept(visitor)  # type: ignore[union-attr]
        result = df.filter(pl_expr)
        assert len(result) == expected_count


class TestPolarsExprVisitorRegex:
    """Tests for regex expression conversion."""

    def test_regex_without_flags(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test regex match without flags."""
        expr = F("name").matches(r"^Alice$")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 1
        assert result["name"][0] == "Alice"

    def test_regex_case_insensitive(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test regex match with IGNORECASE flag."""
        expr = F("name").matches(r"^alice$", flags=re.IGNORECASE)
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 2

    def test_regex_partial_match(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test regex partial string match."""
        expr = F("tag").matches(r"py")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 2

    def test_regex_multiline_flag(self, visitor: PolarsExprVisitor) -> None:
        """Test that MULTILINE flag is converted to inline modifier."""
        expr = F("text").matches(r"^line", flags=re.MULTILINE)
        pl_expr = expr.accept(visitor)
        assert isinstance(pl_expr, pl.Expr)

    def test_regex_dotall_flag(self, visitor: PolarsExprVisitor) -> None:
        """Test that DOTALL flag is converted to inline modifier."""
        expr = F("text").matches(r"a.b", flags=re.DOTALL)
        pl_expr = expr.accept(visitor)
        assert isinstance(pl_expr, pl.Expr)

    def test_regex_combined_flags(self, visitor: PolarsExprVisitor) -> None:
        """Test regex with multiple flags combined."""
        expr = F("text").matches(r"hello", flags=re.IGNORECASE | re.MULTILINE)
        pl_expr = expr.accept(visitor)
        assert isinstance(pl_expr, pl.Expr)


class TestPolarsExprVisitorIn:
    """Tests for IN expression conversion."""

    def test_in_filter(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test IN expression filters correctly."""
        expr = F("tag").is_in("python", "java")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 3

    def test_in_single_value(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test IN expression with a single value."""
        expr = F("status").is_in("inactive")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 1

    def test_in_no_match(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test IN expression with no matching values."""
        expr = F("tag").is_in("cobol", "fortran")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 0


class TestPolarsExprVisitorIsNull:
    """Tests for IsNull expression conversion."""

    def test_is_null_filter(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test IsNull expression filters null values correctly."""
        expr = F("nullable").is_null()
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 2

    def test_is_not_null_filter(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test IsNull(is_null=False) filters non-null values correctly."""
        expr = F("nullable").is_not_null()
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 2


class TestPolarsExprVisitorConstants:
    """Tests for TrueExpr and FalseExpr conversion."""

    def test_true_expr_matches_all(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test TrueExpr matches all rows."""
        pl_expr = TRUE.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == len(df)

    def test_false_expr_matches_none(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test FalseExpr matches no rows."""
        pl_expr = FALSE.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 0


class TestPolarsExprVisitorLogical:
    """Tests for logical expression conversion (AND, OR, NOT)."""

    def test_all_expr(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test AllExpr (AND) filters correctly."""
        expr = (F("age") > 18) & (F("status") == "active")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert all(result["age"] > 18)
        assert all(result["status"] == "active")

    def test_any_expr(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test AnyExpr (OR) filters correctly."""
        expr = (F("age") > 25) | (F("status") == "inactive")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 2

    def test_not_expr(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test NotExpr negates correctly."""
        expr = ~(F("status") == "active")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 1
        assert result["status"][0] == "inactive"

    def test_empty_all_expr_matches_all(
        self, visitor: PolarsExprVisitor, df: "pl.DataFrame"
    ) -> None:
        """Test empty AllExpr (pl.lit(True)) matches all rows."""
        expr = AllExpr()
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == len(df)

    def test_empty_any_expr_matches_none(
        self, visitor: PolarsExprVisitor, df: "pl.DataFrame"
    ) -> None:
        """Test empty AnyExpr (pl.lit(False)) matches no rows."""
        expr = AnyExpr()
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 0

    def test_nested_logical(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test nested logical expressions."""
        expr = ((F("age") >= 20) & (F("age") <= 25)) | (F("name") == "Bob")
        pl_expr = expr.accept(visitor)
        result = df.filter(pl_expr)
        assert len(result) == 3


class TestPolarsExprVisitorLazyFrame:
    """Integration tests for Polars LazyFrame compatibility."""

    def test_filter_lazy_frame(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test that the visitor result works with LazyFrame.filter."""
        expr = F("age") > 20
        pl_expr = expr.accept(visitor)
        result = df.lazy().filter(pl_expr).collect()
        assert len(result) == 2

    def test_complex_lazy_filter(self, visitor: PolarsExprVisitor, df: "pl.DataFrame") -> None:
        """Test complex expression on LazyFrame."""
        expr = (F("age") >= 20) & (F("status") == "active")
        pl_expr = expr.accept(visitor)
        result = df.lazy().filter(pl_expr).collect()
        assert all(result["age"] >= 20)
        assert all(result["status"] == "active")
