"""
Generate PySpark parity fixtures for robin-sparkless tests.

This script is not used by the Rust build; it is a helper to produce
`tests/fixtures/*.json` files that encode:

- Input schema + rows
- A sequence of PySpark-style operations
- Expected schema + rows

See `docs/TEST_CREATION_GUIDE.md` for the full workflow.
"""

from __future__ import annotations

import json
import tempfile
from pathlib import Path
from typing import Any, Dict, List

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType


def schema_to_json(schema: StructType) -> list[dict[str, Any]]:
    """Convert a Spark StructType to a list of {name, type} dicts."""
    return [{"name": f.name, "type": f.dataType.simpleString()} for f in schema.fields]


def df_to_rows(df: Any) -> list[list[Any]]:
    """Collect DataFrame rows as list of lists (one list per row)."""
    return [list(r) for r in df.collect()]


def case_filter_age_gt_30(spark: SparkSession) -> dict[str, Any]:
    data = [(1, 25, "Alice"), (2, 30, "Bob"), (3, 35, "Charlie")]
    df = spark.createDataFrame(data, ["id", "age", "name"])

    out_df = df.filter("age > 30").select("name", "age").orderBy("name")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "filter_age_gt_30",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "filter", "expr": "col('age') > 30"},
            {"op": "select", "columns": ["name", "age"]},
            {"op": "orderBy", "columns": ["name"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_count(spark: SparkSession) -> dict[str, Any]:
    data = [
        (1, "Alice", "Sales"),
        (2, "Bob", "Sales"),
        (3, "Charlie", "Engineering"),
        (4, "David", "Engineering"),
        (5, "Eve", "Sales"),
    ]
    df = spark.createDataFrame(data, ["id", "name", "department"])

    out_df = df.groupBy("department").count().orderBy("department")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_count",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["department"]},
            {"op": "agg", "aggregations": [{"func": "count", "alias": "count"}]},
            {"op": "orderBy", "columns": ["department"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_read_csv(spark: SparkSession) -> dict[str, Any]:
    """Test reading CSV file and applying operations."""
    # Create a temporary CSV file
    csv_content = "id,age,name\n1,25,Alice\n2,30,Bob\n3,35,Charlie\n"

    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
        f.write(csv_content)
        csv_path = f.name

    try:
        # Read CSV using PySpark
        df = (
            spark.read.option("header", "true")
            .option("inferSchema", "true")
            .csv(csv_path)
        )

        # Apply operations
        out_df = df.filter("age > 30").select("name", "age").orderBy("name")

        input_schema = schema_to_json(df.schema)
        input_rows = df_to_rows(df)

        expected_schema = schema_to_json(out_df.schema)
        expected_rows = df_to_rows(out_df)

        return {
            "name": "read_csv",
            "pyspark_version": spark.version,
            "input": {
                "schema": input_schema,
                "rows": input_rows,
                "file_source": {
                    "format": "csv",
                    "content": csv_content,
                },
            },
            "operations": [
                {"op": "filter", "expr": "col('age') > 30"},
                {"op": "select", "columns": ["name", "age"]},
                {"op": "orderBy", "columns": ["name"], "ascending": [True]},
            ],
            "expected": {"schema": expected_schema, "rows": expected_rows},
        }
    finally:
        # Clean up temp file
        Path(csv_path).unlink(missing_ok=True)


def case_read_parquet(spark: SparkSession) -> dict[str, Any]:
    """Test reading Parquet file and applying operations."""
    # Create data and write to Parquet
    data = [(1, "Alice", "Sales"), (2, "Bob", "Engineering"), (3, "Charlie", "Sales")]
    df_temp = spark.createDataFrame(data, ["id", "name", "department"])

    with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as f:
        parquet_path = f.name

    try:
        # Write to Parquet
        df_temp.write.mode("overwrite").parquet(parquet_path)

        # Read Parquet using PySpark
        df = spark.read.parquet(parquet_path)

        # Apply operations
        out_df = (
            df.filter("department == 'Sales'")
            .select("name", "department")
            .orderBy("name")
        )

        input_schema = schema_to_json(df.schema)
        input_rows = df_to_rows(df)

        expected_schema = schema_to_json(out_df.schema)
        expected_rows = df_to_rows(out_df)

        # For Parquet, we'll embed the original data as CSV-like content for the fixture
        # (Parquet is binary, so we store the source data representation)
        parquet_content = (
            "id,name,department\n1,Alice,Sales\n2,Bob,Engineering\n3,Charlie,Sales\n"
        )

        return {
            "name": "read_parquet",
            "pyspark_version": spark.version,
            "input": {
                "schema": input_schema,
                "rows": input_rows,
                "file_source": {
                    "format": "parquet",
                    "content": parquet_content,  # Source data representation
                },
            },
            "operations": [
                {"op": "filter", "expr": "col('department') == 'Sales'"},
                {"op": "select", "columns": ["name", "department"]},
                {"op": "orderBy", "columns": ["name"], "ascending": [True]},
            ],
            "expected": {"schema": expected_schema, "rows": expected_rows},
        }
    finally:
        # Clean up temp file
        import shutil

        if Path(parquet_path).exists():
            if Path(parquet_path).is_dir():
                shutil.rmtree(parquet_path)
            else:
                Path(parquet_path).unlink(missing_ok=True)


def case_read_json(spark: SparkSession) -> dict[str, Any]:
    """Test reading JSON file and applying operations."""
    # Create JSONL content (one JSON object per line)
    json_content = '{"id":1,"age":25,"name":"Alice"}\n{"id":2,"age":30,"name":"Bob"}\n{"id":3,"age":35,"name":"Charlie"}\n'

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
        f.write(json_content)
        json_path = f.name

    try:
        # Read JSON using PySpark
        df = spark.read.option("multiLine", "false").json(json_path)

        # Apply operations
        out_df = df.filter("age > 30").select("name", "age").orderBy("name")

        input_schema = schema_to_json(df.schema)
        input_rows = df_to_rows(df)

        expected_schema = schema_to_json(out_df.schema)
        expected_rows = df_to_rows(out_df)

        return {
            "name": "read_json",
            "pyspark_version": spark.version,
            "input": {
                "schema": input_schema,
                "rows": input_rows,
                "file_source": {
                    "format": "json",
                    "content": json_content,
                },
            },
            "operations": [
                {"op": "filter", "expr": "col('age') > 30"},
                {"op": "select", "columns": ["name", "age"]},
                {"op": "orderBy", "columns": ["name"], "ascending": [True]},
            ],
            "expected": {"schema": expected_schema, "rows": expected_rows},
        }
    finally:
        # Clean up temp file
        Path(json_path).unlink(missing_ok=True)


def case_when_otherwise(spark: SparkSession) -> dict[str, Any]:
    """Test when().otherwise() conditional expression."""
    from pyspark.sql.functions import when as pyspark_when, col, lit

    data = [(1, 25), (2, 17), (3, 30), (4, 16)]
    df = spark.createDataFrame(data, ["id", "age"])

    # Use when(age >= 18).otherwise("minor") -> should return "adult" or "minor"
    out_df = (
        df.withColumn(
            "status",
            pyspark_when(col("age") >= 18, lit("adult")).otherwise(lit("minor")),
        )
        .select("id", "age", "status")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "when_otherwise",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "status",
                "expr": "when(col('age') >= 18).then('adult').otherwise('minor')",
            },
            {"op": "select", "columns": ["id", "age", "status"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_when_then_otherwise(spark: SparkSession) -> dict[str, Any]:
    """Test when().then().otherwise() conditional expression."""
    from pyspark.sql.functions import when as pyspark_when, col, lit

    data = [(1, 25), (2, 17), (3, 30), (4, 16)]
    df = spark.createDataFrame(data, ["id", "age"])

    # Use when(age >= 18).then("adult").otherwise("minor")
    out_df = (
        df.withColumn(
            "status",
            pyspark_when(col("age") >= 18, lit("adult")).otherwise(lit("minor")),
        )
        .select("id", "age", "status")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "when_then_otherwise",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "status",
                "expr": "when(col('age') >= 18).then('adult').otherwise('minor')",
            },
            {"op": "select", "columns": ["id", "age", "status"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_coalesce(spark: SparkSession) -> dict[str, Any]:
    """Test coalesce() function with nulls."""
    from pyspark.sql.functions import coalesce, col, lit
    from pyspark.sql.types import StructType, StructField, StringType, IntegerType

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("col1", StringType(), True),
            StructField("col2", StringType(), True),
        ]
    )
    data = [
        (1, "A", "X"),
        (2, None, "Y"),
        (3, "C", None),
        (4, None, None),
    ]
    df = spark.createDataFrame(data, schema=schema)

    # coalesce(col1, col2, lit("default"))
    out_df = (
        df.withColumn("result", coalesce(col("col1"), col("col2"), lit("default")))
        .select("id", "col1", "col2", "result")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "coalesce",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "result",
                "expr": "coalesce(col('col1'), col('col2'), lit('default'))",
            },
            {"op": "select", "columns": ["id", "col1", "col2", "result"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_sum(spark: SparkSession) -> dict[str, Any]:
    """Test groupBy().sum() aggregation."""
    data = [
        (1, "Sales", 1000),
        (2, "Sales", 1500),
        (3, "Engineering", 2000),
        (4, "Engineering", 2500),
        (5, "Sales", 1200),
    ]
    df = spark.createDataFrame(data, ["id", "department", "salary"])

    out_df = df.groupBy("department").sum("salary").orderBy("department")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_sum",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["department"]},
            {
                "op": "agg",
                "aggregations": [{"func": "sum", "alias": "sum", "column": "salary"}],
            },
            {"op": "orderBy", "columns": ["department"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_avg(spark: SparkSession) -> dict[str, Any]:
    """Test groupBy().avg() aggregation."""
    data = [
        (1, "Sales", 1000),
        (2, "Sales", 1500),
        (3, "Engineering", 2000),
        (4, "Engineering", 2500),
        (5, "Sales", 1200),
    ]
    df = spark.createDataFrame(data, ["id", "department", "salary"])

    out_df = df.groupBy("department").avg("salary").orderBy("department")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_avg",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["department"]},
            {
                "op": "agg",
                "aggregations": [{"func": "avg", "alias": "avg", "column": "salary"}],
            },
            {"op": "orderBy", "columns": ["department"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_min(spark: SparkSession) -> dict[str, Any]:
    """Test groupBy().min() aggregation."""
    data = [
        (1, "Sales", 1000),
        (2, "Sales", 1500),
        (3, "Engineering", 2000),
        (4, "Engineering", 2500),
        (5, "Sales", 1200),
    ]
    df = spark.createDataFrame(data, ["id", "department", "salary"])

    out_df = df.groupBy("department").min("salary").orderBy("department")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_min",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["department"]},
            {
                "op": "agg",
                "aggregations": [{"func": "min", "alias": "min", "column": "salary"}],
            },
            {"op": "orderBy", "columns": ["department"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_max(spark: SparkSession) -> dict[str, Any]:
    """Test groupBy().max() aggregation."""
    data = [
        (1, "Sales", 1000),
        (2, "Sales", 1500),
        (3, "Engineering", 2000),
        (4, "Engineering", 2500),
        (5, "Sales", 1200),
    ]
    df = spark.createDataFrame(data, ["id", "department", "salary"])

    out_df = df.groupBy("department").max("salary").orderBy("department")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_max",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["department"]},
            {
                "op": "agg",
                "aggregations": [{"func": "max", "alias": "max", "column": "salary"}],
            },
            {"op": "orderBy", "columns": ["department"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_multi_agg(spark: SparkSession) -> dict[str, Any]:
    """Test groupBy with multiple aggregations in one agg() call."""
    from pyspark.sql.functions import count, sum as spark_sum, avg

    data = [
        (1, "Sales", 1000),
        (2, "Sales", 1500),
        (3, "Engineering", 2000),
        (4, "Engineering", 2500),
        (5, "Sales", 1200),
    ]
    df = spark.createDataFrame(data, ["id", "department", "salary"])

    out_df = (
        df.groupBy("department")
        .agg(
            count("*").alias("count"),
            spark_sum("salary").alias("sum(salary)"),
            avg("salary").alias("avg(salary)"),
        )
        .orderBy("department")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_multi_agg",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["department"]},
            {
                "op": "agg",
                "aggregations": [
                    {"func": "count", "alias": "count"},
                    {"func": "sum", "alias": "sum(salary)", "column": "salary"},
                    {"func": "avg", "alias": "avg(salary)", "column": "salary"},
                ],
            },
            {"op": "orderBy", "columns": ["department"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_null_keys(spark: SparkSession) -> dict[str, Any]:
    """Test groupBy with NULL values as grouping keys."""
    from pyspark.sql.types import StructType, StructField, StringType, IntegerType

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("category", StringType(), True),
            StructField("value", IntegerType(), True),
        ]
    )
    data = [
        (1, "A", 10),
        (2, "A", 20),
        (3, None, 30),
        (4, "B", 40),
        (5, None, 50),
        (6, "B", 60),
    ]
    df = spark.createDataFrame(data, schema=schema)

    out_df = df.groupBy("category").sum("value").orderBy("category")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_null_keys",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["category"]},
            {
                "op": "agg",
                "aggregations": [{"func": "sum", "alias": "sum", "column": "value"}],
            },
            {"op": "orderBy", "columns": ["category"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_null_comparison_equality(spark: SparkSession) -> dict[str, Any]:
    """Test null comparison semantics: col == NULL, col != NULL return NULL."""
    from pyspark.sql.types import StructType, StructField, IntegerType
    from pyspark.sql.functions import col, lit

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("value", IntegerType(), True),
        ]
    )
    data = [(1, 10), (2, None), (3, 20), (4, None)]
    df = spark.createDataFrame(data, schema=schema)

    # Test that col == NULL and col != NULL return NULL (not boolean)
    # Use a null column instead of lit(None) for Polars compatibility
    from pyspark.sql.functions import when as pyspark_when

    df_with_null_col = df.withColumn("null_col", lit(None).cast("int"))
    out_df = (
        df_with_null_col.withColumn(
            "eq_null",
            pyspark_when(col("value") == col("null_col"), lit("NULL")).otherwise(
                lit("NOT_NULL")
            ),
        )
        .withColumn(
            "ne_null",
            pyspark_when(col("value") != col("null_col"), lit("NOT_NULL")).otherwise(
                lit("NULL")
            ),
        )
        .select("id", "value", "eq_null", "ne_null")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "null_comparison_equality",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "withColumn", "column": "null_col", "expr": "lit(None)"},
            {
                "op": "withColumn",
                "column": "eq_null",
                "expr": "when(col('value') == col('null_col')).then('NULL').otherwise('NOT_NULL')",
            },
            {
                "op": "withColumn",
                "column": "ne_null",
                "expr": "when(col('value') != col('null_col')).then('NOT_NULL').otherwise('NULL')",
            },
            {"op": "select", "columns": ["id", "value", "eq_null", "ne_null"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_null_comparison_ordering(spark: SparkSession) -> dict[str, Any]:
    """Test null comparison semantics: col > NULL, col < NULL return NULL."""
    from pyspark.sql.types import StructType, StructField, IntegerType
    from pyspark.sql.functions import col, lit

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("value", IntegerType(), True),
        ]
    )
    data = [(1, 10), (2, None), (3, 20), (4, None)]
    df = spark.createDataFrame(data, schema=schema)

    # Test that ordering comparisons with NULL return NULL
    # Note: We test this indirectly by checking that comparisons with null columns return null
    # Direct lit(None) comparisons are not easily representable in Polars
    # Instead, we'll test by comparing with a null column
    from pyspark.sql.functions import when as pyspark_when

    # Create a column that's always null for comparison
    df_with_null_col = df.withColumn("null_col", lit(None).cast("int"))
    out_df = (
        df_with_null_col.withColumn(
            "gt_null",
            pyspark_when(col("value") > col("null_col"), lit("TRUE")).otherwise(
                lit("FALSE_OR_NULL")
            ),
        )
        .withColumn(
            "lt_null",
            pyspark_when(col("value") < col("null_col"), lit("TRUE")).otherwise(
                lit("FALSE_OR_NULL")
            ),
        )
        .select("id", "value", "gt_null", "lt_null")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "null_comparison_ordering",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "withColumn", "column": "null_col", "expr": "lit(None)"},
            {
                "op": "withColumn",
                "column": "gt_null",
                "expr": "when(col('value') > col('null_col')).then('TRUE').otherwise('FALSE_OR_NULL')",
            },
            {
                "op": "withColumn",
                "column": "lt_null",
                "expr": "when(col('value') < col('null_col')).then('TRUE').otherwise('FALSE_OR_NULL')",
            },
            {"op": "select", "columns": ["id", "value", "gt_null", "lt_null"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_null_safe_equality(spark: SparkSession) -> dict[str, Any]:
    """Test null-safe equality: NULL <=> NULL returns True."""
    from pyspark.sql.types import StructType, StructField, IntegerType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("value1", IntegerType(), True),
            StructField("value2", IntegerType(), True),
        ]
    )
    data = [
        (1, 10, 10),
        (2, None, None),
        (3, 20, None),
        (4, None, 30),
        (5, 40, 40),
    ]
    df = spark.createDataFrame(data, schema=schema)

    # Test eqNullSafe: NULL <=> NULL = True, value <=> NULL = False
    out_df = (
        df.withColumn("null_safe_eq", col("value1").eqNullSafe(col("value2")))
        .select("id", "value1", "value2", "null_safe_eq")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "null_safe_equality",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "null_safe_eq",
                "expr": "col('value1').eqNullSafe(col('value2'))",
            },
            {"op": "select", "columns": ["id", "value1", "value2", "null_safe_eq"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_null_in_filter(spark: SparkSession) -> dict[str, Any]:
    """Test filtering with null values: df.filter(col != 1) excludes NULL rows."""
    from pyspark.sql.types import StructType, StructField, IntegerType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("value", IntegerType(), True),
        ]
    )
    data = [(1, 1), (2, 2), (3, None), (4, 1), (5, None)]
    df = spark.createDataFrame(data, schema=schema)

    # filter(col != 1) should exclude rows where value is NULL
    # because NULL != 1 returns NULL (falsy)
    out_df = df.filter(col("value") != 1).select("id", "value").orderBy("id")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "null_in_filter",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "filter", "expr": "col('value') != 1"},
            {"op": "select", "columns": ["id", "value"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_type_coercion_numeric(spark: SparkSession) -> dict[str, Any]:
    """Test type coercion: int vs double comparisons."""
    from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("int_col", IntegerType(), True),
            StructField("double_col", DoubleType(), True),
        ]
    )
    data = [(1, 10, 5.5), (2, 20, 15.7), (3, 30, 25.3)]
    df = spark.createDataFrame(data, schema=schema)

    # Test that int_col > 5.5 coerces int to double
    out_df = (
        df.filter(col("int_col") > 5.5)
        .select("id", "int_col", "double_col")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "type_coercion_numeric",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "filter", "expr": "col('int_col') > 5.5"},
            {"op": "select", "columns": ["id", "int_col", "double_col"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_type_coercion_mixed(spark: SparkSession) -> dict[str, Any]:
    """Test type coercion in arithmetic operations."""
    from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("int_val", IntegerType(), True),
            StructField("double_val", DoubleType(), True),
        ]
    )
    data = [(1, 10, 2.5), (2, 20, 3.5), (3, 30, 4.5)]
    df = spark.createDataFrame(data, schema=schema)

    # Test arithmetic with type coercion: int + double = double
    # We'll use withColumn to see the result type
    out_df = (
        df.withColumn("sum", col("int_val") + col("double_val"))
        .select("id", "int_val", "double_val", "sum")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "type_coercion_mixed",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "sum",
                "expr": "col('int_val') + col('double_val')",
            },
            {"op": "select", "columns": ["id", "int_val", "double_val", "sum"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_groupby_with_nulls(spark: SparkSession) -> dict[str, Any]:
    from pyspark.sql.types import StructType, StructField, StringType, IntegerType

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("category", StringType(), True),
        ]
    )
    data = [(1, "A"), (2, "A"), (3, None), (4, "B"), (5, None)]
    df = spark.createDataFrame(data, schema=schema)

    out_df = df.groupBy("category").count().orderBy("category")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "groupby_with_nulls",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "groupBy", "columns": ["category"]},
            {"op": "agg", "aggregations": [{"func": "count", "alias": "count"}]},
            {"op": "orderBy", "columns": ["category"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_filter_and_or(spark: SparkSession) -> dict[str, Any]:
    """Filter combining AND/OR conditions."""
    from pyspark.sql.types import (
        StructType,
        StructField,
        IntegerType,
        IntegerType as IntType,
    )
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("age", IntegerType(), True),
            StructField("score", IntType(), True),
        ]
    )
    data = [
        (1, 25, 50),
        (2, 35, 60),
        (3, 45, 110),
        (4, 55, 70),
    ]
    df = spark.createDataFrame(data, schema=schema)

    out_df = df.filter(
        (col("age") > 30) & ((col("score") < 100) | (col("age") > 50))
    ).orderBy("id")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "filter_and_or",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "filter",
                "expr": "col('age') > 30 AND (col('score') < 100 OR col('age') > 50)",
            },
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_filter_not(spark: SparkSession) -> dict[str, Any]:
    """Filter using NOT on a condition."""
    from pyspark.sql.types import StructType, StructField, IntegerType, StringType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("flag", StringType(), True),
        ]
    )
    data = [
        (1, "Y"),
        (2, "N"),
        (3, "Y"),
        (4, "N"),
    ]
    df = spark.createDataFrame(data, schema=schema)

    out_df = df.filter(~(col("flag") == "N")).orderBy("id")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "filter_not",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "filter", "expr": "NOT col('flag') == 'N'"},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_with_logical_column(spark: SparkSession) -> dict[str, Any]:
    """withColumn producing a boolean from a complex logical expression."""
    from pyspark.sql.types import StructType, StructField, IntegerType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("age", IntegerType(), True),
            StructField("score", IntegerType(), True),
        ]
    )
    data = [
        (1, 25, 50),
        (2, 35, 60),
        (3, 45, 110),
        (4, 55, 70),
    ]
    df = spark.createDataFrame(data, schema=schema)

    out_df = (
        df.withColumn(
            "is_target",
            (col("age") > 30) & (col("score") < 100),
        )
        .select("id", "age", "score", "is_target")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "with_logical_column",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "is_target",
                "expr": "col('age') > 30 AND col('score') < 100",
            },
            {"op": "select", "columns": ["id", "age", "score", "is_target"]},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_filter_nested(spark: SparkSession) -> dict[str, Any]:
    """Filter with nested logical expressions using parentheses."""
    from pyspark.sql.types import StructType, StructField, IntegerType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("age", IntegerType(), True),
            StructField("score", IntegerType(), True),
            StructField("vip", IntegerType(), True),
        ]
    )
    data = [
        (1, 25, 50, 0),
        (2, 35, 60, 0),
        (3, 45, 110, 1),
        (4, 55, 70, 0),
        (5, 30, 80, 1),
    ]
    df = spark.createDataFrame(data, schema=schema)

    # Test: (age > 30 AND (score < 100 OR vip == 1))
    out_df = df.filter(
        (col("age") > 30) & ((col("score") < 100) | (col("vip") == 1))
    ).orderBy("id")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "filter_nested",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "filter",
                "expr": "(col('age') > 30 AND (col('score') < 100 OR col('vip') == 1))",
            },
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_with_arithmetic_logical_mix(spark: SparkSession) -> dict[str, Any]:
    """withColumn combining arithmetic and logical conditions."""
    from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType
    from pyspark.sql.functions import col

    schema = StructType(
        [
            StructField("id", IntegerType(), True),
            StructField("int_val", IntegerType(), True),
            StructField("double_val", DoubleType(), True),
            StructField("threshold", DoubleType(), True),
        ]
    )
    data = [
        (1, 10, 2.5, 5.0),
        (2, 20, 3.5, 10.0),
        (3, 30, 4.5, 15.0),
        (4, 15, 1.5, 8.0),
    ]
    df = spark.createDataFrame(data, schema=schema)

    # Test: withColumn producing a boolean from arithmetic + logical
    # (int_val + double_val) > threshold
    out_df = (
        df.withColumn(
            "above_threshold",
            (col("int_val") + col("double_val")) > col("threshold"),
        )
        .select("id", "int_val", "double_val", "threshold", "above_threshold")
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)

    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "with_arithmetic_logical_mix",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "above_threshold",
                "expr": "(col('int_val') + col('double_val')) > col('threshold')",
            },
            {
                "op": "select",
                "columns": [
                    "id",
                    "int_val",
                    "double_val",
                    "threshold",
                    "above_threshold",
                ],
            },
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_inner_join(spark: SparkSession) -> dict[str, Any]:
    """Test inner join between employees and departments."""
    employees_data = [
        {"id": 1, "name": "Alice", "dept_id": 10, "salary": 50000},
        {"id": 2, "name": "Bob", "dept_id": 20, "salary": 60000},
        {"id": 3, "name": "Charlie", "dept_id": 10, "salary": 70000},
        {"id": 4, "name": "David", "dept_id": 30, "salary": 55000},
    ]
    departments_data = [
        {"dept_id": 10, "name": "IT", "location": "NYC"},
        {"dept_id": 20, "name": "HR", "location": "LA"},
        {"dept_id": 40, "name": "Finance", "location": "Chicago"},
    ]
    emp_df = spark.createDataFrame(employees_data)
    dept_df = spark.createDataFrame(departments_data)
    out_df = emp_df.join(dept_df, emp_df.dept_id == dept_df.dept_id, "inner")

    input_schema = schema_to_json(emp_df.schema)
    input_rows = df_to_rows(emp_df)
    right_schema = schema_to_json(dept_df.schema)
    right_rows = df_to_rows(dept_df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "inner_join",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "right_input": {"schema": right_schema, "rows": right_rows},
        "operations": [
            {"op": "join", "on": ["dept_id"], "how": "inner"},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_left_join(spark: SparkSession) -> dict[str, Any]:
    """Test left join between employees and departments."""
    employees_data = [
        {"id": 1, "name": "Alice", "dept_id": 10, "salary": 50000},
        {"id": 2, "name": "Bob", "dept_id": 20, "salary": 60000},
        {"id": 3, "name": "Charlie", "dept_id": 10, "salary": 70000},
        {"id": 4, "name": "David", "dept_id": 30, "salary": 55000},
    ]
    departments_data = [
        {"dept_id": 10, "name": "IT", "location": "NYC"},
        {"dept_id": 20, "name": "HR", "location": "LA"},
        {"dept_id": 40, "name": "Finance", "location": "Chicago"},
    ]
    emp_df = spark.createDataFrame(employees_data)
    dept_df = spark.createDataFrame(departments_data)
    out_df = emp_df.join(dept_df, emp_df.dept_id == dept_df.dept_id, "left")
    out_df = out_df.orderBy("id")

    input_schema = schema_to_json(emp_df.schema)
    input_rows = df_to_rows(emp_df)
    right_schema = schema_to_json(dept_df.schema)
    right_rows = df_to_rows(dept_df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "left_join",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "right_input": {"schema": right_schema, "rows": right_rows},
        "operations": [
            {"op": "join", "on": ["dept_id"], "how": "left"},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_right_join(spark: SparkSession) -> dict[str, Any]:
    """Test right join between employees and departments."""
    employees_data = [
        {"id": 1, "name": "Alice", "dept_id": 10, "salary": 50000},
        {"id": 2, "name": "Bob", "dept_id": 20, "salary": 60000},
        {"id": 3, "name": "Charlie", "dept_id": 10, "salary": 70000},
        {"id": 4, "name": "David", "dept_id": 30, "salary": 55000},
    ]
    departments_data = [
        {"dept_id": 10, "name": "IT", "location": "NYC"},
        {"dept_id": 20, "name": "HR", "location": "LA"},
        {"dept_id": 40, "name": "Finance", "location": "Chicago"},
    ]
    emp_df = spark.createDataFrame(employees_data)
    dept_df = spark.createDataFrame(departments_data)
    out_df = emp_df.join(dept_df, emp_df.dept_id == dept_df.dept_id, "right")
    out_df = out_df.orderBy("id")

    input_schema = schema_to_json(emp_df.schema)
    input_rows = df_to_rows(emp_df)
    right_schema = schema_to_json(dept_df.schema)
    right_rows = df_to_rows(dept_df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "right_join",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "right_input": {"schema": right_schema, "rows": right_rows},
        "operations": [
            {"op": "join", "on": ["dept_id"], "how": "right"},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_outer_join(spark: SparkSession) -> dict[str, Any]:
    """Test outer join between employees and departments."""
    employees_data = [
        {"id": 1, "name": "Alice", "dept_id": 10, "salary": 50000},
        {"id": 2, "name": "Bob", "dept_id": 20, "salary": 60000},
        {"id": 3, "name": "Charlie", "dept_id": 10, "salary": 70000},
        {"id": 4, "name": "David", "dept_id": 30, "salary": 55000},
    ]
    departments_data = [
        {"dept_id": 10, "name": "IT", "location": "NYC"},
        {"dept_id": 20, "name": "HR", "location": "LA"},
        {"dept_id": 40, "name": "Finance", "location": "Chicago"},
    ]
    emp_df = spark.createDataFrame(employees_data)
    dept_df = spark.createDataFrame(departments_data)
    out_df = emp_df.join(dept_df, emp_df.dept_id == dept_df.dept_id, "outer")
    out_df = out_df.orderBy("id")

    input_schema = schema_to_json(emp_df.schema)
    input_rows = df_to_rows(emp_df)
    right_schema = schema_to_json(dept_df.schema)
    right_rows = df_to_rows(dept_df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "outer_join",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "right_input": {"schema": right_schema, "rows": right_rows},
        "operations": [
            {"op": "join", "on": ["dept_id"], "how": "outer"},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_string_upper_lower(spark: SparkSession) -> dict[str, Any]:
    """Test upper() and lower() string functions."""
    from pyspark.sql.functions import col, lower, upper

    data = [(1, "Alice", "Sales"), (2, "Bob", "Engineering"), (3, "Charlie", "SALES")]
    df = spark.createDataFrame(data, ["id", "name", "dept"])
    out_df = (
        df.withColumn("name_upper", upper(col("name")))
        .withColumn("dept_lower", lower(col("dept")))
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "string_upper_lower",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "withColumn", "column": "name_upper", "expr": "upper(col('name'))"},
            {"op": "withColumn", "column": "dept_lower", "expr": "lower(col('dept'))"},
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_string_substring(spark: SparkSession) -> dict[str, Any]:
    """Test substring() string function."""
    from pyspark.sql.functions import col, substring

    data = [(1, "Alice"), (2, "Bob"), (3, "Charlie")]
    df = spark.createDataFrame(data, ["id", "name"])
    out_df = df.withColumn("name_prefix", substring(col("name"), 1, 3)).orderBy("id")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "string_substring",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "name_prefix",
                "expr": "substring(col('name'), 1, 3)",
            },
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_string_concat(spark: SparkSession) -> dict[str, Any]:
    """Test concat() and concat_ws() string functions."""
    from pyspark.sql.functions import col, concat, concat_ws, lit

    data = [(1, "Alice", "Smith"), (2, "Bob", "Jones")]
    df = spark.createDataFrame(data, ["id", "first", "last"])
    out_df = (
        df.withColumn("full_name", concat(col("first"), lit(" "), col("last")))
        .withColumn("initials", concat_ws(".", col("first"), col("last")))
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "string_concat",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "withColumn",
                "column": "full_name",
                "expr": "concat(col('first'), lit(' '), col('last'))",
            },
            {
                "op": "withColumn",
                "column": "initials",
                "expr": "concat_ws('.', col('first'), col('last'))",
            },
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_row_number_window(spark: SparkSession) -> dict[str, Any]:
    """Test row_number() window: partition by dept, order by salary desc."""
    from pyspark.sql.functions import col, row_number
    from pyspark.sql.window import Window

    data = [
        (1, "Alice", "Sales", 1000),
        (2, "Bob", "Sales", 1500),
        (3, "Charlie", "Engineering", 2000),
        (4, "David", "Engineering", 2500),
        (5, "Eve", "Sales", 1200),
    ]
    df = spark.createDataFrame(data, ["id", "name", "dept", "salary"])
    window = Window.partitionBy("dept").orderBy(col("salary").desc())
    out_df = df.withColumn("rn", row_number().over(window)).orderBy("dept", "rn")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "row_number_window",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "orderBy",
                "columns": ["dept", "salary"],
                "ascending": [True, False],
            },
            {
                "op": "window",
                "column": "rn",
                "func": "row_number",
                "partition_by": ["dept"],
                "order_by": [{"col": "salary", "asc": False}],
            },
            {"op": "orderBy", "columns": ["dept", "rn"], "ascending": [True, True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_rank_window(spark: SparkSession) -> dict[str, Any]:
    """Test rank() window: partition by dept, order by salary desc."""
    from pyspark.sql.functions import col, rank
    from pyspark.sql.window import Window

    data = [
        (1, "Alice", "Sales", 1000),
        (2, "Bob", "Sales", 1500),
        (3, "Charlie", "Sales", 1500),
        (4, "David", "Engineering", 2000),
    ]
    df = spark.createDataFrame(data, ["id", "name", "dept", "salary"])
    window = Window.partitionBy("dept").orderBy(col("salary").desc())
    out_df = df.withColumn("rank", rank().over(window)).orderBy("dept", "rank")

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "rank_window",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {
                "op": "orderBy",
                "columns": ["dept", "salary"],
                "ascending": [True, False],
            },
            {
                "op": "window",
                "column": "rank",
                "func": "rank",
                "partition_by": ["dept"],
                "order_by": [{"col": "salary", "asc": False}],
            },
            {"op": "orderBy", "columns": ["dept", "rank"], "ascending": [True, True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def case_lag_lead_window(spark: SparkSession) -> dict[str, Any]:
    """Test lag and lead window functions."""
    from pyspark.sql.functions import col, lag, lead
    from pyspark.sql.window import Window

    data = [(1, "A", 10), (2, "A", 20), (3, "A", 30), (4, "B", 40), (5, "B", 50)]
    df = spark.createDataFrame(data, ["id", "dept", "val"])
    window = Window.partitionBy("dept").orderBy("id")
    out_df = (
        df.withColumn("lag_val", lag(col("val"), 1).over(window))
        .withColumn("lead_val", lead(col("val"), 1).over(window))
        .orderBy("id")
    )

    input_schema = schema_to_json(df.schema)
    input_rows = df_to_rows(df)
    expected_schema = schema_to_json(out_df.schema)
    expected_rows = df_to_rows(out_df)

    return {
        "name": "lag_lead_window",
        "pyspark_version": spark.version,
        "input": {"schema": input_schema, "rows": input_rows},
        "operations": [
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
            {
                "op": "window",
                "column": "lag_val",
                "func": "lag",
                "partition_by": ["dept"],
                "order_by": [{"col": "id", "asc": True}],
                "value_column": "val",
            },
            {
                "op": "window",
                "column": "lead_val",
                "func": "lead",
                "partition_by": ["dept"],
                "order_by": [{"col": "id", "asc": True}],
                "value_column": "val",
            },
            {"op": "orderBy", "columns": ["id"], "ascending": [True]},
        ],
        "expected": {"schema": expected_schema, "rows": expected_rows},
    }


def main() -> None:
    spark = SparkSession.builder.appName("robin_sparkless_parity_gen").getOrCreate()

    out_dir = Path("tests/fixtures")
    out_dir.mkdir(parents=True, exist_ok=True)

    fixtures: List[Dict[str, Any]] = [
        case_filter_age_gt_30(spark),
        case_groupby_count(spark),
        case_groupby_with_nulls(spark),
        case_filter_and_or(spark),
        case_filter_not(spark),
        case_filter_nested(spark),
        case_read_csv(spark),
        case_read_parquet(spark),
        case_read_json(spark),
        case_when_otherwise(spark),
        case_when_then_otherwise(spark),
        case_coalesce(spark),
        case_groupby_sum(spark),
        case_groupby_avg(spark),
        case_groupby_min(spark),
        case_groupby_max(spark),
        case_groupby_multi_agg(spark),
        case_groupby_null_keys(spark),
        case_null_comparison_equality(spark),
        case_null_comparison_ordering(spark),
        case_null_safe_equality(spark),
        case_null_in_filter(spark),
        case_type_coercion_numeric(spark),
        case_type_coercion_mixed(spark),
        case_with_logical_column(spark),
        case_with_arithmetic_logical_mix(spark),
        case_inner_join(spark),
        case_left_join(spark),
        case_right_join(spark),
        case_outer_join(spark),
        case_string_upper_lower(spark),
        case_string_substring(spark),
        case_string_concat(spark),
        case_row_number_window(spark),
        case_rank_window(spark),
        case_lag_lead_window(spark),
    ]

    for fx in fixtures:
        path = out_dir / f"{fx['name']}.json"
        path.write_text(json.dumps(fx, indent=2))
        print(f"wrote {path}")

    spark.stop()


if __name__ == "__main__":
    main()
