"""Integration tests for struct namespace expressions.

These tests require Ray and test end-to-end struct namespace expression evaluation.
"""

import pandas as pd
import pyarrow as pa
import pytest
from packaging import version

import ray
from ray.data._internal.util import rows_same
from ray.data.expressions import col
from ray.data.tests.conftest import *  # noqa
from ray.tests.conftest import *  # noqa

pytestmark = pytest.mark.skipif(
    version.parse(pa.__version__) < version.parse("19.0.0"),
    reason="Namespace expressions tests require PyArrow >= 19.0",
)


def _create_dataset(items_data, dataset_format, arrow_table=None):
    if dataset_format == "arrow":
        if arrow_table is not None:
            ds = ray.data.from_arrow(arrow_table)
        else:
            table = pa.Table.from_pylist(items_data)
            ds = ray.data.from_arrow(table)
    elif dataset_format == "pandas":
        if arrow_table is not None:
            df = arrow_table.to_pandas()
        else:
            df = pd.DataFrame(items_data)
        ds = ray.data.from_blocks([df])
    return ds


DATASET_FORMATS = ["pandas", "arrow"]


@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
class TestStructNamespace:
    """Tests for struct namespace operations."""

    def test_struct_bracket_bool_index_raises(self, dataset_format):
        """Test struct[bool] raises TypeError instead of being treated as int."""
        del dataset_format  # Unused, required by class-level parametrization.

        with pytest.raises(
            TypeError, match="Struct indices must be strings or integers"
        ):
            col("user").struct[True]

    @pytest.mark.parametrize("bad_index", ["1", True])
    def test_struct_field_by_index_non_integer_raises(self, dataset_format, bad_index):
        """Test struct.field_by_index() rejects non-integer indices."""
        del dataset_format  # Unused, required by class-level parametrization.

        with pytest.raises(TypeError, match="Struct field index must be an integer"):
            col("user").struct.field_by_index(bad_index)

    def test_struct_field_by_index_negative_raises(self, dataset_format):
        """Test struct.field_by_index() rejects negative indices."""
        del dataset_format  # Unused, required by class-level parametrization.

        with pytest.raises(
            ValueError, match="Struct field index must be non-negative, got -1"
        ):
            col("user").struct.field_by_index(-1)

        with pytest.raises(
            ValueError, match="Struct field index must be non-negative, got -1"
        ):
            col("user").struct[-1]

    def test_struct_field(self, ray_start_regular_shared, dataset_format):
        """Test struct.field() extracts field."""
        arrow_table = pa.table(
            {
                "user": pa.array(
                    [
                        {"name": "Alice", "age": 30},
                        {"name": "Bob", "age": 25},
                    ],
                    type=pa.struct(
                        [
                            pa.field("name", pa.string()),
                            pa.field("age", pa.int32()),
                        ]
                    ),
                )
            }
        )
        items_data = [
            {"user": {"name": "Alice", "age": 30}},
            {"user": {"name": "Bob", "age": 25}},
        ]
        ds = _create_dataset(items_data, dataset_format, arrow_table)

        result = ds.with_column("age", col("user").struct.field("age")).to_pandas()
        expected = pd.DataFrame(
            {
                "user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
                "age": [30, 25],
            }
        )
        assert rows_same(result, expected)

    def test_struct_bracket(self, ray_start_regular_shared, dataset_format):
        """Test struct['field'] bracket notation."""
        arrow_table = pa.table(
            {
                "user": pa.array(
                    [
                        {"name": "Alice", "age": 30},
                        {"name": "Bob", "age": 25},
                    ],
                    type=pa.struct(
                        [
                            pa.field("name", pa.string()),
                            pa.field("age", pa.int32()),
                        ]
                    ),
                )
            }
        )
        items_data = [
            {"user": {"name": "Alice", "age": 30}},
            {"user": {"name": "Bob", "age": 25}},
        ]
        ds = _create_dataset(items_data, dataset_format, arrow_table)

        result = ds.with_column("name", col("user").struct["name"]).to_pandas()
        expected = pd.DataFrame(
            {
                "user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
                "name": ["Alice", "Bob"],
            }
        )
        assert rows_same(result, expected)

    def test_struct_field_by_index(self, ray_start_regular_shared, dataset_format):
        """Test struct.field_by_index() extracts field by position."""
        if dataset_format == "pandas":
            pytest.skip(
                "Index-based struct access requires stable Arrow struct field ordering."
            )
        arrow_table = pa.table(
            {
                "user": pa.array(
                    [
                        {"name": "Alice", "age": 30},
                        {"name": "Bob", "age": 25},
                    ],
                    type=pa.struct(
                        [
                            pa.field("name", pa.string()),
                            pa.field("age", pa.int32()),
                        ]
                    ),
                )
            }
        )
        items_data = [
            {"user": {"name": "Alice", "age": 30}},
            {"user": {"name": "Bob", "age": 25}},
        ]
        ds = _create_dataset(items_data, dataset_format, arrow_table)

        result = ds.with_column("age", col("user").struct.field_by_index(1)).to_pandas()
        expected = pd.DataFrame(
            {
                "user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
                "age": [30, 25],
            }
        )
        assert rows_same(result, expected)

    def test_struct_bracket_with_index(self, ray_start_regular_shared, dataset_format):
        """Test struct[index] bracket notation."""
        if dataset_format == "pandas":
            pytest.skip(
                "Index-based struct access requires stable Arrow struct field ordering."
            )
        arrow_table = pa.table(
            {
                "user": pa.array(
                    [
                        {"name": "Alice", "age": 30},
                        {"name": "Bob", "age": 25},
                    ],
                    type=pa.struct(
                        [
                            pa.field("name", pa.string()),
                            pa.field("age", pa.int32()),
                        ]
                    ),
                )
            }
        )
        items_data = [
            {"user": {"name": "Alice", "age": 30}},
            {"user": {"name": "Bob", "age": 25}},
        ]
        ds = _create_dataset(items_data, dataset_format, arrow_table)

        result = ds.with_column("name", col("user").struct[0]).to_pandas()
        expected = pd.DataFrame(
            {
                "user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
                "name": ["Alice", "Bob"],
            }
        )
        assert rows_same(result, expected)

    def test_struct_nested_field(self, ray_start_regular_shared, dataset_format):
        """Test nested struct field access with .field()."""
        arrow_table = pa.table(
            {
                "user": pa.array(
                    [
                        {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
                        {"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
                    ],
                    type=pa.struct(
                        [
                            pa.field("name", pa.string()),
                            pa.field(
                                "address",
                                pa.struct(
                                    [
                                        pa.field("city", pa.string()),
                                        pa.field("zip", pa.string()),
                                    ]
                                ),
                            ),
                        ]
                    ),
                )
            }
        )
        items_data = [
            {"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
            {"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
        ]
        ds = _create_dataset(items_data, dataset_format, arrow_table)

        result = ds.with_column(
            "city", col("user").struct.field("address").struct.field("city")
        ).to_pandas()
        expected = pd.DataFrame(
            {
                "user": [
                    {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
                    {"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
                ],
                "city": ["NYC", "LA"],
            }
        )
        assert rows_same(result, expected)

    def test_struct_nested_bracket(self, ray_start_regular_shared, dataset_format):
        """Test nested struct field access with brackets."""
        arrow_table = pa.table(
            {
                "user": pa.array(
                    [
                        {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
                        {"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
                    ],
                    type=pa.struct(
                        [
                            pa.field("name", pa.string()),
                            pa.field(
                                "address",
                                pa.struct(
                                    [
                                        pa.field("city", pa.string()),
                                        pa.field("zip", pa.string()),
                                    ]
                                ),
                            ),
                        ]
                    ),
                )
            }
        )
        items_data = [
            {"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
            {"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
        ]
        ds = _create_dataset(items_data, dataset_format, arrow_table)

        result = ds.with_column(
            "zip", col("user").struct["address"].struct["zip"]
        ).to_pandas()
        expected = pd.DataFrame(
            {
                "user": [
                    {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
                    {"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
                ],
                "zip": ["10001", "90001"],
            }
        )
        assert rows_same(result, expected)


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))