"""Tests for DataFrame operations."""
import pytest
from singlespark import SparkSession, Row
from singlespark.sql import functions as F
from singlespark.sql.types import (
    StructType, StructField, StringType, IntegerType, LongType, DoubleType
)


@pytest.fixture
def spark():
    s = SparkSession.builder.appName("test_df").create()
    yield s
    s.stop()


@pytest.fixture
def people(spark):
    data = [("Alice", 30, "Engineering"), ("Bob", 25, "Sales"), ("Carol", 35, "Engineering")]
    return spark.createDataFrame(data, ["name", "age", "dept"])


# ---------------------------------------------------------------------------
# Schema
# ---------------------------------------------------------------------------

def test_columns(people):
    assert people.columns == ["name", "age", "dept"]


def test_dtypes(people):
    names = [name for name, _ in people.dtypes]
    assert names == ["name", "age", "dept"]


def test_schema(people):
    schema = people.schema
    assert schema.fieldNames() == ["name", "age", "dept"]


def test_print_schema(people, capsys):
    people.printSchema()
    out = capsys.readouterr().out
    assert "name" in out


# ---------------------------------------------------------------------------
# createDataFrame
# ---------------------------------------------------------------------------

def test_create_from_list_of_tuples(spark):
    df = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "val"])
    assert df.count() == 2
    assert df.columns == ["id", "val"]


def test_create_from_list_of_dicts(spark):
    df = spark.createDataFrame([{"x": 1, "y": 2}, {"x": 3, "y": 4}])
    assert df.count() == 2


def test_create_with_struct_schema(spark):
    schema = StructType([
        StructField("name", StringType()),
        StructField("age", IntegerType()),
    ])
    df = spark.createDataFrame([("Alice", 30)], schema)
    assert df.columns == ["name", "age"]


def test_create_from_rows(spark):
    rows = [Row(name="Alice", age=30), Row(name="Bob", age=25)]
    df = spark.createDataFrame(rows)
    assert df.count() == 2


# ---------------------------------------------------------------------------
# select / filter
# ---------------------------------------------------------------------------

def test_select_by_name(people):
    df = people.select("name", "age")
    assert df.columns == ["name", "age"]


def test_select_by_col(people):
    df = people.select(F.col("name").alias("n"), F.col("age"))
    assert "n" in df.columns


def test_filter_column_expr(people):
    result = people.filter(F.col("age") > 28)
    assert result.count() == 2


def test_filter_string_condition(people):
    result = people.filter("age > 28")
    assert result.count() == 2


def test_where_alias(people):
    result = people.where(F.col("name") == "Alice")
    assert result.count() == 1


def test_filter_chained(people):
    result = people.filter(F.col("age") > 24).filter(F.col("dept") == "Engineering")
    assert result.count() == 2


# ---------------------------------------------------------------------------
# withColumn / withColumnRenamed / drop
# ---------------------------------------------------------------------------

def test_with_column_new(people):
    df = people.withColumn("senior", F.col("age") >= 30)
    assert "senior" in df.columns


def test_with_column_replace(people):
    df = people.withColumn("age", F.col("age") * 2)
    first_age = df.filter(F.col("name") == "Alice").collect()[0]["age"]
    assert first_age == 60


def test_with_column_renamed(people):
    df = people.withColumnRenamed("name", "full_name")
    assert "full_name" in df.columns
    assert "name" not in df.columns


def test_drop(people):
    df = people.drop("dept")
    assert "dept" not in df.columns
    assert "name" in df.columns


# ---------------------------------------------------------------------------
# groupBy / agg
# ---------------------------------------------------------------------------

def test_group_by_count(people):
    result = people.groupBy("dept").count().collect()
    counts = {r["dept"]: r["count"] for r in result}
    assert counts["Engineering"] == 2
    assert counts["Sales"] == 1


def test_group_by_agg_sum(people):
    result = people.groupBy("dept").agg(F.sum("age").alias("total_age")).collect()
    eng_row = next(r for r in result if r["dept"] == "Engineering")
    assert eng_row["total_age"] == 65  # 30 + 35


def test_group_by_agg_avg(people):
    result = people.groupBy("dept").agg(F.avg("age").alias("avg_age")).collect()
    eng_row = next(r for r in result if r["dept"] == "Engineering")
    assert eng_row["avg_age"] == 32.5


def test_global_agg(people):
    result = people.agg(F.count("*").alias("n"), F.max("age").alias("oldest")).collect()
    assert result[0]["n"] == 3
    assert result[0]["oldest"] == 35


# ---------------------------------------------------------------------------
# join
# ---------------------------------------------------------------------------

def test_inner_join(spark):
    left = spark.createDataFrame([(1, "Alice"), (2, "Bob")], ["id", "name"])
    right = spark.createDataFrame([(1, "Eng"), (3, "Sales")], ["id", "dept"])
    result = left.join(right, on="id", how="inner")
    assert result.count() == 1


def test_left_join(spark):
    left = spark.createDataFrame([(1, "Alice"), (2, "Bob")], ["id", "name"])
    right = spark.createDataFrame([(1, "Eng")], ["id", "dept"])
    result = left.join(right, on="id", how="left")
    assert result.count() == 2


def test_cross_join(spark):
    a = spark.createDataFrame([(1,), (2,)], ["x"])
    b = spark.createDataFrame([(3,), (4,)], ["y"])
    result = a.crossJoin(b)
    assert result.count() == 4


# ---------------------------------------------------------------------------
# orderBy / sort
# ---------------------------------------------------------------------------

def test_order_by_asc(people):
    result = people.orderBy("age").collect()
    ages = [r["age"] for r in result]
    assert ages == sorted(ages)


def test_order_by_desc(people):
    result = people.orderBy(F.col("age").desc()).collect()
    ages = [r["age"] for r in result]
    assert ages == sorted(ages, reverse=True)


def test_sort_alias(people):
    result = people.sort("age", ascending=False).collect()
    ages = [r["age"] for r in result]
    assert ages == sorted(ages, reverse=True)


# ---------------------------------------------------------------------------
# Set operations
# ---------------------------------------------------------------------------

def test_union(spark):
    a = spark.createDataFrame([(1, "a")], ["x", "y"])
    b = spark.createDataFrame([(2, "b")], ["x", "y"])
    result = a.union(b)
    assert result.count() == 2


def test_distinct(spark):
    df = spark.createDataFrame([(1, "a"), (1, "a"), (2, "b")], ["x", "y"])
    assert df.distinct().count() == 2


def test_limit(people):
    assert people.limit(2).count() == 2


# ---------------------------------------------------------------------------
# Actions
# ---------------------------------------------------------------------------

def test_collect_returns_rows(people):
    rows = people.collect()
    assert len(rows) == 3
    assert all(isinstance(r, Row) for r in rows)


def test_first(people):
    row = people.first()
    assert isinstance(row, Row)


def test_count(people):
    assert people.count() == 3


def test_show(people, capsys):
    people.show(2)
    out = capsys.readouterr().out
    assert len(out) > 0


def test_to_pandas(people):
    pdf = people.toPandas()
    assert len(pdf) == 3
    assert list(pdf.columns) == ["name", "age", "dept"]


def test_to_polars(people):
    import polars as pl
    pdf = people.toPolars()
    assert isinstance(pdf, pl.DataFrame)


# ---------------------------------------------------------------------------
# select expr
# ---------------------------------------------------------------------------

def test_select_expr(people):
    result = people.selectExpr("name", "age * 2 AS double_age")
    assert "double_age" in result.columns
    alice = result.filter(F.col("name") == "Alice").collect()[0]
    assert alice["double_age"] == 60


# ---------------------------------------------------------------------------
# Null handling
# ---------------------------------------------------------------------------

def test_dropna(spark):
    df = spark.createDataFrame([(1, None), (2, "b"), (None, "c")], ["x", "y"])
    assert df.dropna().count() == 1


def test_fillna(spark):
    df = spark.createDataFrame([(1, None), (2, "b")], ["x", "y"])
    filled = df.fillna("N/A")
    rows = filled.collect()
    assert any(r["y"] == "N/A" for r in rows)


# ---------------------------------------------------------------------------
# when / otherwise
# ---------------------------------------------------------------------------

def test_when_otherwise(people):
    df = people.withColumn("category",
        F.when(F.col("age") >= 30, "senior").otherwise("junior"))
    seniors = df.filter(F.col("category") == "senior").count()
    assert seniors == 2


def test_when_chained(people):
    df = people.withColumn("tier",
        F.when(F.col("age") > 33, "top")
         .when(F.col("age") > 27, "mid")
         .otherwise("low"))
    tiers = {r["name"]: r["tier"] for r in df.collect()}
    assert tiers["Alice"] == "mid"
    assert tiers["Bob"] == "low"
    assert tiers["Carol"] == "top"


# ---------------------------------------------------------------------------
# Column attribute access
# ---------------------------------------------------------------------------

def test_getitem_string(people):
    col = people["name"]
    assert col is not None


def test_getattr_column(people):
    col = people.name
    assert col is not None


def test_getitem_filter(people):
    result = people[F.col("age") > 28]
    assert result.count() == 2
