"""Tests for Polars sort specification visitor."""

import pytest

pl = pytest.importorskip("polars")

from therismos.sorting import SortCriterion, SortOrder, SortSpec
from therismos.sorting.visitors.polars import PolarsSortSpec, PolarsSortSpecVisitor


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


@pytest.fixture()
def df() -> "pl.DataFrame":
    """Return a sample DataFrame for integration tests."""
    return pl.DataFrame(
        {
            "age": [25, 30, 20, 25],
            "name": ["Bob", "Alice", "Charlie", "Alice"],
            "score": [85, 91, 72, 60],
        }
    )


class TestPolarsSortSpecVisitorBasic:
    """Tests for basic PolarsSortSpecVisitor functionality."""

    def test_single_ascending_criterion(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test Polars sort spec for single ascending criterion."""
        spec = SortSpec([SortCriterion("age", SortOrder.ASCENDING)])
        result = spec.accept(visitor)
        assert result.by == ("age",)
        assert result.descending == (False,)

    def test_single_descending_criterion(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test Polars sort spec for single descending criterion."""
        spec = SortSpec([SortCriterion("name", SortOrder.DESCENDING)])
        result = spec.accept(visitor)
        assert result.by == ("name",)
        assert result.descending == (True,)

    def test_multiple_criteria(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test Polars sort spec for multiple criteria."""
        spec = SortSpec(
            [
                SortCriterion("age", SortOrder.ASCENDING),
                SortCriterion("name", SortOrder.DESCENDING),
            ]
        )
        result = spec.accept(visitor)
        assert result.by == ("age", "name")
        assert result.descending == (False, True)

    def test_empty_spec(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test Polars sort spec for empty spec."""
        spec = SortSpec()
        result = spec.accept(visitor)
        assert result.by == ()
        assert result.descending == ()

    def test_result_is_polars_sort_spec(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test that result is a PolarsSortSpec."""
        spec = SortSpec([SortCriterion("x", SortOrder.ASCENDING)])
        result = spec.accept(visitor)
        assert isinstance(result, PolarsSortSpec)


class TestPolarsSortSpecVisitorNoneOrder:
    """Tests for handling SortOrder.NONE."""

    def test_single_none_criterion_skipped(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test that NONE order is skipped."""
        spec = SortSpec([SortCriterion("age", SortOrder.NONE)])
        result = spec.accept(visitor)
        assert result.by == ()
        assert result.descending == ()

    def test_none_mixed_with_other_orders(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test that NONE orders are skipped while others are preserved."""
        spec = SortSpec(
            [
                SortCriterion("age", SortOrder.ASCENDING),
                SortCriterion("name", SortOrder.NONE),
                SortCriterion("score", SortOrder.DESCENDING),
            ]
        )
        result = spec.accept(visitor)
        assert result.by == ("age", "score")
        assert result.descending == (False, True)
        assert "name" not in result.by

    def test_all_none_orders(self, visitor: PolarsSortSpecVisitor) -> None:
        """Test spec with all NONE orders produces empty result."""
        spec = SortSpec(
            [
                SortCriterion("a", SortOrder.NONE),
                SortCriterion("b", SortOrder.NONE),
            ]
        )
        result = spec.accept(visitor)
        assert result.by == ()
        assert result.descending == ()


@pytest.mark.parametrize(
    "criteria,expected_by,expected_descending",
    [
        (
            [SortCriterion("x", SortOrder.ASCENDING)],
            ("x",),
            (False,),
        ),
        (
            [SortCriterion("y", SortOrder.DESCENDING)],
            ("y",),
            (True,),
        ),
        (
            [SortCriterion("a", SortOrder.ASCENDING), SortCriterion("b", SortOrder.DESCENDING)],
            ("a", "b"),
            (False, True),
        ),
        (
            [SortCriterion("a", SortOrder.ASCENDING), SortCriterion("b", SortOrder.NONE)],
            ("a",),
            (False,),
        ),
        ([], (), ()),
    ],
)
def test_polars_sort_visitor_parametrized(
    criteria: list[SortCriterion],
    expected_by: tuple[str, ...],
    expected_descending: tuple[bool, ...],
) -> None:
    """Parametrized test for Polars sort visitor with various specs."""
    spec = SortSpec(criteria)
    visitor = PolarsSortSpecVisitor()
    result = spec.accept(visitor)
    assert result.by == expected_by
    assert result.descending == expected_descending


class TestPolarsSortSpecVisitorIntegration:
    """Integration tests for Polars sort visitor with actual DataFrames."""

    def test_sort_ascending(self, visitor: PolarsSortSpecVisitor, df: "pl.DataFrame") -> None:
        """Test ascending sort on DataFrame."""
        spec = SortSpec([SortCriterion("age", SortOrder.ASCENDING)])
        sort = spec.accept(visitor)
        result = df.sort(by=list(sort.by), descending=list(sort.descending))
        ages = result["age"].to_list()
        assert ages == sorted(ages)

    def test_sort_descending(self, visitor: PolarsSortSpecVisitor, df: "pl.DataFrame") -> None:
        """Test descending sort on DataFrame."""
        spec = SortSpec([SortCriterion("score", SortOrder.DESCENDING)])
        sort = spec.accept(visitor)
        result = df.sort(by=list(sort.by), descending=list(sort.descending))
        scores = result["score"].to_list()
        assert scores == sorted(scores, reverse=True)

    def test_multi_column_sort(self, visitor: PolarsSortSpecVisitor, df: "pl.DataFrame") -> None:
        """Test multi-column sort on DataFrame."""
        spec = SortSpec(
            [
                SortCriterion("age", SortOrder.ASCENDING),
                SortCriterion("name", SortOrder.ASCENDING),
            ]
        )
        sort = spec.accept(visitor)
        result = df.sort(by=list(sort.by), descending=list(sort.descending))
        assert result["age"][0] == 20

    def test_sort_lazy_frame(self, visitor: PolarsSortSpecVisitor, df: "pl.DataFrame") -> None:
        """Test sort on LazyFrame."""
        spec = SortSpec([SortCriterion("age", SortOrder.ASCENDING)])
        sort = spec.accept(visitor)
        result = df.lazy().sort(by=list(sort.by), descending=list(sort.descending)).collect()
        ages = result["age"].to_list()
        assert ages == sorted(ages)
