from __future__ import annotations

from collections.abc import Callable
from typing import TYPE_CHECKING

import pytest

import daft

if TYPE_CHECKING:
    from daft.dataframe.dataframe import DataFrame


def _assert_sql_raise(sql: str, bindings: dict[str, DataFrame], error_msg: str) -> None:
    with pytest.raises(Exception) as excinfo:
        daft.sql(sql, **bindings).collect()

    assert error_msg in str(excinfo.value)


def _assert_df_op_raise(func: Callable, error_msg: str) -> None:
    with pytest.raises(Exception) as excinfo:
        func()

    assert error_msg in str(excinfo.value)


def test_div_floor():
    df = daft.from_pydict({"A": [1, 2, 3, 4], "B": [1.5, 2.5, 3.5, 4.5], "C": [4, 5, 6, 7]})

    actual1 = daft.sql("SELECT (A // 1) AS div_floor FROM df").collect().to_pydict()
    actual2 = df.select((daft.col("A") // 1).alias("div_floor")).collect().to_pydict()
    expected = {
        "div_floor": [1, 2, 3, 4],
    }
    assert actual1 == actual2 == expected

    actual1 = daft.sql("SELECT (A // B) AS div_floor FROM df").collect().to_pydict()
    actual2 = df.select((daft.col("A") // daft.col("B")).alias("div_floor")).collect().to_pydict()
    expected = {
        "div_floor": [0, 0, 0, 0],
    }
    assert actual1 == actual2 == expected

    actual1 = daft.sql("SELECT (B // A) AS div_floor FROM df").collect().to_pydict()
    actual2 = df.select((daft.col("B") // daft.col("A")).alias("div_floor")).collect().to_pydict()
    expected = {
        "div_floor": [1.0, 1.0, 1.0, 1.0],
    }
    assert actual1 == actual2 == expected

    actual1 = daft.sql("SELECT (C // A) AS div_floor FROM df").collect().to_pydict()
    actual2 = df.select((daft.col("C") // daft.col("A")).alias("div_floor")).collect().to_pydict()
    expected = {
        "div_floor": [4, 2, 2, 1],
    }
    assert actual1 == actual2 == expected


def test_unsupported_div_floor():
    df = daft.from_pydict({"A": [1, 2, 3, 4], "B": [1.5, 2.5, 3.5, 4.5], "C": [True, False, True, True]})

    bindings = {"df": df}

    _assert_sql_raise("SELECT A // C FROM df", bindings, "TypeError Cannot perform floor divide on types: Int64, Bool")

    _assert_sql_raise("SELECT C // A FROM df", bindings, "TypeError Cannot perform floor divide on types: Bool, Int64")

    _assert_sql_raise(
        "SELECT B // C FROM df", bindings, "TypeError Cannot perform floor divide on types: Float64, Bool"
    )

    _assert_sql_raise(
        "SELECT B // C FROM df", bindings, "TypeError Cannot perform floor divide on types: Float64, Bool"
    )

    _assert_df_op_raise(
        lambda: df.select(daft.col("A") // daft.col("C")).collect(),
        "Cannot perform floor divide on types: Int64, Bool",
    )

    _assert_df_op_raise(
        lambda: df.select(daft.col("C") // daft.col("A")).collect(),
        "Cannot perform floor divide on types: Bool, Int64",
    )

    _assert_df_op_raise(
        lambda: df.select(daft.col("B") // daft.col("C")).collect(),
        "Cannot perform floor divide on types: Float64, Bool",
    )

    _assert_df_op_raise(
        lambda: df.select(daft.col("C") // daft.col("B")).collect(),
        "Cannot perform floor divide on types: Bool, Float64",
    )