"""Tests for Spark SQL via DuckDB."""
import pytest
from singlespark import SparkSession
from singlespark.sql import functions as F


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


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


def test_sql_select(spark, people_view):
    result = spark.sql("SELECT name FROM people WHERE age > 28")
    names = [r["name"] for r in result.collect()]
    assert "Alice" in names
    assert "Carol" in names
    assert "Bob" not in names


def test_sql_count(spark, people_view):
    result = spark.sql("SELECT count(*) as cnt FROM people")
    assert result.collect()[0]["cnt"] == 3


def test_sql_group_by(spark, people_view):
    result = spark.sql("SELECT dept, count(*) as n FROM people GROUP BY dept ORDER BY dept")
    rows = {r["dept"]: r["n"] for r in result.collect()}
    assert rows["Engineering"] == 2
    assert rows["Sales"] == 1


def test_sql_join(spark):
    left = spark.createDataFrame([(1, "Alice"), (2, "Bob")], ["id", "name"])
    right = spark.createDataFrame([(1, "Eng"), (2, "Sales")], ["id", "dept"])
    left.createOrReplaceTempView("left_t")
    right.createOrReplaceTempView("right_t")
    result = spark.sql("SELECT l.name, r.dept FROM left_t l JOIN right_t r ON l.id = r.id")
    assert result.count() == 2


def test_sql_window_function(spark):
    df = spark.createDataFrame(
        [("Alice", 30, "Eng"), ("Bob", 25, "Eng"), ("Carol", 35, "Sales")],
        ["name", "age", "dept"]
    )
    df.createOrReplaceTempView("emp")
    result = spark.sql("""
        SELECT name, age, dept,
               ROW_NUMBER() OVER (PARTITION BY dept ORDER BY age) as rn
        FROM emp
    """)
    assert result.count() == 3


def test_create_or_replace_temp_view(spark):
    df = spark.createDataFrame([(1,), (2,)], ["x"])
    df.createOrReplaceTempView("nums")
    result = spark.sql("SELECT * FROM nums WHERE x > 1")
    assert result.count() == 1


def test_create_temp_view_duplicate_raises(spark):
    df = spark.createDataFrame([(1,)], ["x"])
    df.createOrReplaceTempView("nums2")
    with pytest.raises(Exception):
        df.createTempView("nums2")


def test_spark_table(spark):
    df = spark.createDataFrame([(1, "a")], ["id", "val"])
    df.createOrReplaceTempView("mytable")
    result = spark.table("mytable")
    assert result.count() == 1


def test_select_expr(spark):
    df = spark.createDataFrame([(1, 2), (3, 4)], ["a", "b"])
    result = df.selectExpr("a + b AS total", "a * b AS product")
    assert "total" in result.columns
    assert "product" in result.columns
    # Verify the computed values
    rows = result.collect()
    totals = sorted(r["total"] for r in rows)
    assert totals == [3, 7]  # 1+2=3, 3+4=7


def test_sql_aggregations(spark):
    df = spark.createDataFrame([(1, 10.0), (2, 20.0), (3, 30.0)], ["id", "val"])
    df.createOrReplaceTempView("vals")
    result = spark.sql("SELECT avg(val) as avg_val, sum(val) as sum_val FROM vals").collect()[0]
    assert result["avg_val"] == 20.0
    assert result["sum_val"] == 60.0
