"""
Integration tests for data cleaning through the 6-stage pipeline.

Tests focus on:
- Cleaning interaction with validation (check_on_cleaned=True)
- Multiple cleaners in one schema
- Derived columns depending on cleaned columns
- Refine stage behavior (aliases, column selection)
- Row removal after cleaning
- Table-level validations after cleaning
"""

import polars as pl
import pytest

import zeolite as z
from zeolite.types import ProcessingSuccess, ProcessingFailure


class TestCleaningPipelineIntegration:
    """Integration tests for cleaning through the full 6-stage pipeline."""

    def test_multiple_cleaners_in_single_schema(self):
        """Test that multiple different cleaners execute correctly in one pipeline run."""
        # Arrange - messy data requiring multiple cleaning operations
        df = pl.DataFrame(
            {
                "id": ["  123  ", "456", "789"],
                "name": ["  ALICE  ", "BOB#", "  Charlie$  "],
                "active": ["yes", "no", "true"],
                "amount": ["1,234.56", "7,890.12", "500"],
                "gender": ["Male", "female", "MALE"],
                "date": ["2023-01-15", "2023-12-31", "2024-06-01"],
            }
        )

        schema = z.schema("test").columns(
            z.col.str("id").clean(z.Clean.id(prefix="ORG", sanitise="trim")),
            z.col.str("name").clean(z.Clean.string(sanitise="full")),
            z.col.str("active").clean(
                z.Clean.boolean(
                    true_values={"yes", "true", "1"},
                    false_values={"no", "false", "0"},
                )
            ),
            z.col.str("amount").clean(z.Clean.number()),
            z.col.str("gender").clean(
                z.Clean.enum(
                    enum_map={"male": "M", "female": "F"},
                    sanitise="lowercase",
                )
            ),
            z.col.str("date").clean(z.Clean.date()),
        )

        # Act
        result = schema.apply(df)

        # Assert - all cleaners should have executed successfully
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # Verify each cleaner's output
        assert final_df["id"][0] == "ORG::123"
        assert final_df["name"][0] == "alice"
        assert final_df["active"][0] is True
        assert final_df["amount"][0] == 1234.56
        assert final_df["gender"][0] == "M"
        assert final_df["date"].dtype == pl.Date

    def test_cleaning_preserves_row_order(self):
        """Test that cleaning operations preserve original row order."""
        # Arrange
        df = pl.DataFrame(
            {
                "id": ["999", "001", "500", "250"],
                "value": ["  z  ", "  a  ", "  m  ", "  d  "],
            }
        )

        schema = z.schema("test").columns(
            z.col.str("id").clean(z.Clean.id(prefix="ORG")),
            z.col.str("value").clean(z.Clean.string(sanitise="trim")),
        )

        # Act
        result = schema.apply(df)

        # Assert - order should be preserved (not sorted)
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()
        assert list(final_df["id"]) == ["ORG::999", "ORG::001", "ORG::500", "ORG::250"]
        assert list(final_df["value"]) == ["z", "a", "m", "d"]

    def test_cleaning_with_nulls_across_columns(self):
        """Test that null values are handled consistently across multiple cleaners."""
        # Arrange
        df = pl.DataFrame(
            {
                "id": ["123", None, "789"],
                "name": ["Alice", None, "Charlie"],
                "active": ["yes", None, "no"],
                "amount": ["100", None, "300"],
            }
        )

        schema = z.schema("test").columns(
            z.col.str("id").clean(z.Clean.id(prefix="ORG")),
            z.col.str("name").clean(z.Clean.string(sanitise="lowercase")),
            z.col.str("active").clean(z.Clean.boolean()),
            z.col.str("amount").clean(z.Clean.number()),
        )

        # Act
        result = schema.apply(df)

        # Assert - nulls should be preserved across all cleaned columns
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # Row 2 (index 1) should have nulls in all columns
        assert final_df["id"][1] is None
        assert final_df["name"][1] is None
        assert final_df["active"][1] is None
        assert final_df["amount"][1] is None


class TestCleaningWithValidation:
    """Integration tests for validation with check_on_cleaned=True."""

    def test_validation_check_on_cleaned_detects_duplicates(self):
        """Test that check_on_cleaned=True validates cleaned values, not raw values."""
        # Arrange - raw IDs are different, but cleaned IDs are duplicates
        df = pl.DataFrame({"id": ["  123  ", "123", "456"]})

        schema = z.schema("test").columns(
            z.col.str("id")
            .clean(z.Clean.id(prefix="ORG", sanitise="trim"))
            .validations(z.Check.unique(check_on_cleaned=True, reject="any"))
        )

        # Act
        result = schema.apply(df)

        # Assert - should detect duplicate after cleaning ("ORG::123" appears twice)
        assert isinstance(result, ProcessingFailure)
        assert any("unique" in str(err).lower() for err in result.errors)

    def test_validation_check_on_cleaned_with_enum(self):
        """Test validation on cleaned enum values."""
        # Arrange - various gender inputs that clean to same enum value
        df = pl.DataFrame({"gender": ["male", "MALE", "Male", "female"]})

        schema = z.schema("test").columns(
            z.col.str("gender")
            .clean(
                z.Clean.enum(
                    enum_map={"male": "M", "female": "F"},
                    sanitise="lowercase",
                )
            )
            .validations(z.Check.not_empty(check_on_cleaned=True, warning="any"))
        )

        # Act
        result = schema.apply(df)

        # Assert - should validate cleaned values (all "M" or "F")
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()
        assert all(v in ["M", "F"] for v in final_df["gender"])

    def test_validation_on_raw_vs_cleaned_values(self):
        """Test difference between validation on raw vs cleaned values."""
        # Arrange - IDs with whitespace
        df = pl.DataFrame({"id": ["  123", "  456  ", "  123  "]})

        # Schema 1: Check uniqueness on RAW values (should pass - all different with whitespace)
        schema_raw = z.schema("test").columns(
            z.col.str("id")
            .clean(z.Clean.id(prefix="ORG", sanitise="trim"))
            .validations(z.Check.unique(check_on_cleaned=False, reject="any"))
        )

        # Schema 2: Check uniqueness on CLEANED values (should fail - duplicates after trim)
        schema_cleaned = z.schema("test").columns(
            z.col.str("id")
            .clean(z.Clean.id(prefix="ORG", sanitise="trim"))
            .validations(z.Check.unique(check_on_cleaned=True, reject="any"))
        )

        # Act
        result_raw = schema_raw.apply(df.lazy())
        result_cleaned = schema_cleaned.apply(df.lazy())

        # Assert
        assert isinstance(result_raw, ProcessingSuccess)  # Raw values are unique
        assert isinstance(
            result_cleaned, ProcessingFailure
        )  # Cleaned values have duplicates

    def test_validation_with_row_removal_after_cleaning(self):
        """Test that row removal works correctly based on cleaned column validation."""
        # Arrange
        df = pl.DataFrame(
            {
                "id": ["123", "456", "789", "999"],
                "amount": ["100", "invalid", "300", "not-a-number"],
            }
        )

        schema = z.schema("test").columns(
            z.col.str("id").clean(z.Clean.id(prefix="ORG")),
            z.col.str("amount")
            .clean(z.Clean.number())
            .validations(
                z.Check.not_empty(
                    check_on_cleaned=True,
                    reject=0.6,
                    remove_row_on_fail=True,
                )
            ),
        )

        # Act
        result = schema.apply(df)

        # Assert - rows with invalid amounts should be removed
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # Should only have 2 rows left (123, 789) - the ones with valid amounts
        assert final_df.height == 2
        assert list(final_df["id"]) == ["ORG::123", "ORG::789"]


class TestCleaningWithDerivedColumns:
    """Integration tests for derived columns depending on cleaned columns."""

    def test_derived_from_cleaned_column(self):
        """Test derived column using cleaned column reference."""
        # Arrange
        df = pl.DataFrame({"price": [100, 200, 300]})

        schema = z.schema("test").columns(
            z.col.float("price").clean(
                z.Clean.custom(
                    lambda col: col * 1.15,
                    data_type="float",
                )
            ),
            z.col.derived(
                "high_price",
                function=z.ref("price").clean().col > 200,
                data_type="boolean",
            ),
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # After 15% markup: 115, 230, 345
        assert list(final_df["high_price"]) == [False, True, True]

    def test_derived_from_multiple_cleaned_columns(self):
        """Test derived column referencing multiple cleaned columns."""
        # Arrange
        df = pl.DataFrame(
            {
                "first_name": ["  alice  ", "  bob  "],
                "last_name": ["  SMITH  ", "  JONES  "],
            }
        )

        schema = z.schema("test").columns(
            z.col.str("first_name").clean(z.Clean.string(sanitise="full")),
            z.col.str("last_name").clean(z.Clean.string(sanitise="full")),
            z.col.derived(
                "full_name",
                function=z.ref("first_name").clean().col
                + " "
                + z.ref("last_name").clean().col,
                data_type="string",
            ),
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()
        assert list(final_df["full_name"]) == ["alice smith", "bob jones"]

    def test_derived_with_validation_on_cleaned(self):
        """Test complex chain: clean → derive from cleaned → validate derived."""
        # Arrange
        df = pl.DataFrame({"price_str": ["100", "200", "300"]})

        schema = z.schema("test").columns(
            z.col.str("price_str").clean(z.Clean.number()),
            z.col.derived(
                "price_with_tax",
                function=z.ref("price_str").clean().col * 1.15,
                data_type="float",
            ).validations(z.Check.in_range(min_value=0, max_value=400)),
        )

        # Act
        result = schema.apply(df)

        # Assert - should pass because all prices with tax are under 400
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()
        assert all(v <= 400 for v in final_df["price_with_tax"])


class TestCleaningRefineStage:
    """Integration tests for cleaned column behavior in refine stage (aliases, selection)."""

    def test_clean_with_alias(self):
        """Test that cleaned column uses alias name in final output."""
        # Arrange
        df = pl.DataFrame({"raw_name": ["  Alice  ", "  Bob  "]})
        schema = z.schema("test").columns(
            z.col.str("raw_name").clean(
                z.Clean.string(sanitise="full").alias("clean_name")
            )
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # Should have clean_name column instead of raw_name__clean
        assert "clean_name" in final_df.columns
        assert "raw_name" in final_df.columns
        assert list(final_df["clean_name"]) == ["alice", "bob"]

    def test_clean_without_alias(self):
        """Test that without alias, original column is replaced with cleaned value."""
        # Arrange
        df = pl.DataFrame({"name": ["  Alice  ", "  Bob  "]})
        schema = z.schema("test").columns(
            z.col.str("name").clean(z.Clean.string(sanitise="full"))
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # Cleaned value should replace original
        assert "name" in final_df.columns
        assert list(final_df["name"]) == ["alice", "bob"]

    def test_custom_clean_with_alias(self):
        """Test custom clean uses custom alias in final output."""
        # Arrange
        df = pl.DataFrame({"raw_name": ["alice", "bob"]})

        schema = z.schema("test").columns(
            z.col.str("raw_name").clean(
                z.Clean.custom(
                    lambda col: col.str.to_uppercase(),
                    alias="clean_name",
                )
            ),
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()
        assert "clean_name" in final_df.columns
        assert "raw_name" in final_df.columns
        assert list(final_df["clean_name"]) == ["ALICE", "BOB"]

    def test_multiple_aliases_in_schema(self):
        """Test multiple cleaned columns with different alias behaviors."""
        # Arrange
        df = pl.DataFrame(
            {
                "col1": ["  alice  ", "  bob  "],
                "col2": ["  CHARLIE  ", "  DIANA  "],
                "col3": ["100", "200"],
            }
        )

        schema = z.schema("test").columns(
            # col1: with alias (keeps both raw and cleaned)
            z.col.str("col1").clean(
                z.Clean.string(sanitise="full").alias("name_lower")
            ),
            # col2: without alias (replaces original)
            z.col.str("col2").clean(z.Clean.string(sanitise="full")),
            # col3: with alias (keeps both raw and cleaned)
            z.col.str("col3").clean(z.Clean.number().alias("amount")),
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # Check final column names
        assert "name_lower" in final_df.columns  # alias - cleaned column
        assert "col1" in final_df.columns  # alias - raw column kept
        assert "col2" in final_df.columns  # no alias - original replaced with cleaned
        assert "amount" in final_df.columns  # alias - cleaned column
        assert "col3" in final_df.columns  # alias - raw column kept

        # Check cleaned values
        assert list(final_df["name_lower"]) == ["alice", "bob"]
        assert list(final_df["col2"]) == [
            "charlie",
            "diana",
        ]  # Replaced raw with cleaned
        assert list(final_df["amount"]) == [100, 200]

        # Check raw values are preserved when alias is used
        assert list(final_df["col1"]) == ["  alice  ", "  bob  "]  # Raw preserved
        assert list(final_df["col3"]) == ["100", "200"]  # Raw preserved


class TestCleaningWithTableValidations:
    """Integration tests for table-level validations after cleaning."""

    def test_table_validation_removed_rows_after_cleaning(self):
        """Test that table-level validation runs after row removal from cleaning validation."""
        # Arrange
        df = pl.DataFrame(
            {
                "id": ["123", "456", "789", "999", "111"],
                "amount": ["100", "invalid", "300", "not-a-number", "500"],
            }
        )

        schema = (
            z.schema("test")
            .columns(
                z.col.str("id").clean(z.Clean.id(prefix="ORG")),
                z.col.str("amount")
                .clean(z.Clean.number())
                .validations(
                    z.Check.not_empty(
                        check_on_cleaned=True,
                        remove_row_on_fail=True,
                    )
                ),
            )
            .table_validation(
                # Reject if more than 50% of rows removed
                z.TableCheck.removed_rows(reject=0.5)
            )
        )

        # Act
        result = schema.apply(df)

        # Assert - should fail because 2/5 (40%) rows removed (close to 50% threshold)
        # Actually should pass because 40% < 50%, but let's verify the behavior
        if isinstance(result, ProcessingSuccess):
            final_df = result.data.collect()
            # Should have 3 rows left (removed 2 invalid amounts)
            assert final_df.height == 3
        else:
            # If it fails, it's because removed_rows threshold was exceeded
            assert isinstance(result, ProcessingFailure)

    def test_table_validation_min_rows_after_cleaning(self):
        """Test minimum row count validation after cleaning removes rows."""
        # Arrange - 5 rows, but 3 have invalid amounts
        df = pl.DataFrame(
            {
                "id": ["123", "456", "789", "999", "111"],
                "amount": ["100", "invalid", "invalid", "invalid", "500"],
            }
        )

        schema = (
            z.schema("test")
            .columns(
                z.col.str("id").clean(z.Clean.id(prefix="ORG")),
                z.col.str("amount")
                .clean(z.Clean.number())
                .validations(
                    z.Check.not_empty(
                        check_on_cleaned=True,
                        remove_row_on_fail=True,
                    )
                ),
            )
            .table_validation(
                # Require at least 3 rows after cleaning
                z.TableCheck.min_rows(reject=3)
            )
        )

        # Act
        result = schema.apply(df)

        # Assert - should fail because only 2 rows remain (< 3 required)
        assert isinstance(result, ProcessingFailure)
        assert any("min_rows" in str(err).lower() for err in result.errors)


class TestCleaningEdgeCases:
    """Integration tests for edge cases in the cleaning pipeline."""

    def test_cleaning_all_null_column(self):
        """Test cleaning a column where all values are null."""
        # Arrange
        df = pl.DataFrame(
            {"id": [None, None, None], "name": ["Alice", "Bob", "Charlie"]}
        )

        schema = z.schema("test").columns(
            z.col.str("id").clean(z.Clean.id(prefix="ORG")),
            z.col.str("name").clean(z.Clean.string(sanitise="lowercase")),
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()
        assert all(v is None for v in final_df["id"])
        assert list(final_df["name"]) == ["alice", "bob", "charlie"]

    def test_cleaning_with_temporary_column(self):
        """Test that temporary columns are excluded from final output after cleaning."""
        # Arrange
        df = pl.DataFrame(
            {
                "raw_value": ["100", "200", "300"],
                "multiplier": [1.1, 1.2, 1.3],
            }
        )

        schema = z.schema("test").columns(
            z.col.str("raw_value").clean(z.Clean.number()),
            z.col.float("multiplier").temporary(
                True
            ),  # Should be excluded from final output
            z.col.derived(
                "adjusted_value",
                function=z.ref("raw_value").clean().col * z.ref("multiplier").col,
                data_type="float",
            ),
        )

        # Act
        result = schema.apply(df)

        # Assert
        assert isinstance(result, ProcessingSuccess)
        final_df = result.data.collect()

        # multiplier should not be in final output
        assert "multiplier" not in final_df.columns
        assert "raw_value" in final_df.columns
        assert "adjusted_value" in final_df.columns

        # Check calculated values: 100*1.1=110, 200*1.2=240, 300*1.3=390
        assert final_df["adjusted_value"][0] == pytest.approx(110.0)
        assert final_df["adjusted_value"][1] == pytest.approx(240.0)
        assert final_df["adjusted_value"][2] == pytest.approx(390.0)
