"""Unit tests for custom cleaning variant (Clean.custom)."""

import polars as pl
import pytest
import zeolite as z
from zeolite.exceptions import CleanConfigurationError


class TestCleanCustom:
    """Tests for Clean.custom with lambda functions and expressions."""

    def test_clean_custom_with_lambda(self):
        """Test custom clean instantiation and apply with lambda function."""
        ref = z.ref("value")
        clean = z.Clean.custom(lambda col: col.str.to_uppercase())
        expr = clean.apply(ref)

        df = pl.DataFrame({"value": ["hello", "world", "test"]})
        result = df.with_columns(expr).select(ref.clean().col).to_series()

        assert result.to_list() == ["HELLO", "WORLD", "TEST"]

    def test_clean_custom_with_expression(self):
        """Test custom clean with direct Polars expression."""
        ref = z.ref("value")
        clean = z.Clean.custom(pl.col("value").str.to_lowercase().str.strip_chars())
        expr = clean.apply(ref)

        df = pl.DataFrame({"value": ["  HELLO  ", " WORLD ", "TEST"]})
        result = df.with_columns(expr).select(ref.clean().col).to_series()

        assert result.to_list() == ["hello", "world", "test"]

    def test_clean_custom_handles_various_operations(self):
        """Test custom clean with numeric operations and nulls."""
        ref = z.ref("value")
        clean = z.Clean.custom(lambda col: (col * 1.15).round(2), data_type="float")
        expr = clean.apply(ref)

        df = pl.DataFrame({"value": [10.0, 20.0, None, 30.0]})
        result = df.with_columns(expr).select(ref.clean().col).to_series()

        assert result.item(0) == 11.5
        assert result.item(1) == 23.0
        assert result.item(2) is None
        assert result.item(3) == 34.5

    def test_clean_custom_with_complex_expression(self):
        """Test custom clean with complex multi-column expression."""
        df = pl.DataFrame({"price": [100, 200, 300], "tax_rate": [0.1, 0.15, 0.2]})

        ref = z.ref("price")
        clean = z.Clean.custom(
            lambda col: col * (1 + pl.col("tax_rate")), data_type="float"
        )
        expr = clean.apply(ref)

        result = df.with_columns(expr).select(ref.clean().col).to_series()

        assert result.item(0) == pytest.approx(110.0, abs=0.01)
        assert result.item(1) == pytest.approx(230.0, abs=0.01)
        assert result.item(2) == pytest.approx(360.0, abs=0.01)

    def test_clean_custom_requires_callable_or_expr(self):
        """Test that function must be callable or Expr."""
        with pytest.raises(CleanConfigurationError) as exc_info:
            z.Clean.custom("not a function")

        assert "A function or Polars Expression must be provided" in str(exc_info.value)

    def test_clean_custom_function_must_return_expr(self):
        """Test that function must return a Polars Expression."""
        with pytest.raises(CleanConfigurationError) as exc_info:
            z.Clean.custom(lambda col: "not an expression")

        assert "function must return a Polars Expression" in str(exc_info.value)

    def test_clean_custom_whitespace_variations(self, whitespace_df):
        """Test custom clean with whitespace edge cases."""
        ref = z.ref("value")
        clean = z.Clean.custom(lambda col: col.str.strip_chars())
        expr = clean.apply(ref)

        result = whitespace_df.with_columns(expr).select(ref.clean().col).to_series()

        # Should strip whitespace from all values
        assert result.to_list() == ["", "", "", "", "", "test", "test", "test"]

    def test_clean_custom_null_string_variations(self, edge_case_strings):
        """Test custom clean with various null string representations."""
        ref = z.ref("value")
        clean = z.Clean.custom(lambda col: col.str.to_lowercase())
        expr = clean.apply(ref)

        df = pl.DataFrame({"value": edge_case_strings})
        result = df.with_columns(expr).select(ref.clean().col).to_series()

        # Should lowercase all string values
        assert result.to_list() == [
            "",
            "   ",
            "\t",
            "\n",
            "\t\n",
            "null",
            "null",
            "none",
            "none",
        ]

    def test_clean_custom_mixed_input_types(self, mixed_type_df):
        """Test custom clean with mixed input data types."""
        ref = z.ref("mixed_col")
        clean = z.Clean.custom(lambda col: pl.coalesce(col, pl.lit("default")))
        expr = clean.apply(ref)

        result = mixed_type_df.with_columns(expr).select(ref.clean().col).to_series()

        # Should replace nulls with default
        assert result.to_list() == ["123", "456", "default", "abc", "7.89", "", "   "]

    def test_clean_custom_scientific_notation(self, scientific_notation_df):
        """Test custom clean with scientific notation numbers."""
        ref = z.ref("value")
        clean = z.Clean.custom(
            lambda col: col.cast(pl.Float64) * 1000, data_type="float"
        )
        expr = clean.apply(ref)

        result = (
            scientific_notation_df.with_columns(expr)
            .select(ref.clean().col)
            .to_series()
        )

        # Should parse scientific notation and multiply by 1000
        assert result.item(0) == pytest.approx(1e-5 * 1000)  # "1e-5"
        assert result.item(1) == pytest.approx(1.5e10 * 1000, rel=1e-4)  # "1.5E+10"
        assert result.item(2) == pytest.approx(1.0 * 1000)  # "1e+0"
        assert result.item(3) == pytest.approx(-1e-5 * 1000)  # "-1e-5"
        assert result.item(4) == pytest.approx(2.5e-3 * 1000)  # "2.5e-3"
        assert result.item(5) == pytest.approx(1.23e2 * 1000)  # "1.23e+2"

    def test_clean_custom_control_characters(self):
        """Test custom clean with control characters."""
        ref = z.ref("value")
        clean = z.Clean.custom(lambda col: col.str.replace_all(r"[\x00-\x1F]", ""))
        expr = clean.apply(ref)

        df = pl.DataFrame(
            {"value": ["hello\x00world", "test\x01data", "line\nbreak", "tab\there"]}
        )
        result = df.with_columns(expr).select(ref.clean().col).to_series()

        # Should remove control characters
        assert result.to_list() == ["helloworld", "testdata", "linebreak", "tabhere"]

    def test_clean_custom_unicode_normalization(self):
        """Test custom clean with unicode normalization."""
        ref = z.ref("value")
        clean = z.Clean.custom(lambda col: col.str.normalize("NFKD"))
        expr = clean.apply(ref)

        df = pl.DataFrame({"value": ["café", "cafe\u0301", "Müller", "Mu\u0308ller"]})
        result = df.with_columns(expr).select(ref.clean().col).to_series()

        # Should normalize unicode characters
        # Note: The exact result depends on how Polars handles normalization
        assert len(result) == 4
        assert result.item(0) == result.item(
            1
        )  # Should be equivalent after normalization
        assert result.item(2) == result.item(
            3
        )  # Should be equivalent after normalization
