from __future__ import annotations

import sys

import pytest

import narwhals as nw


def test_polars(monkeypatch: pytest.MonkeyPatch) -> None:
    pytest.importorskip("polars")
    import polars as pl

    monkeypatch.delitem(sys.modules, "pandas", raising=False)
    monkeypatch.delitem(sys.modules, "numpy", raising=False)
    monkeypatch.delitem(sys.modules, "pyarrow", raising=False)
    monkeypatch.delitem(sys.modules, "typing_extensions", raising=False)
    monkeypatch.delitem(sys.modules, "duckdb", raising=False)
    monkeypatch.delitem(sys.modules, "dask", raising=False)
    monkeypatch.delitem(sys.modules, "ibis", raising=False)
    monkeypatch.delitem(sys.modules, "pyspark", raising=False)
    df = pl.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]})
    nw.from_native(df, eager_only=True).group_by("a").agg(nw.col("b").mean()).filter(
        nw.col("a") > 1
    )
    assert "polars" in sys.modules
    assert "pandas" not in sys.modules
    assert "numpy" not in sys.modules
    assert "pyarrow" not in sys.modules
    assert "dask" not in sys.modules
    assert "ibis" not in sys.modules
    assert "pyspark" not in sys.modules
    assert "duckdb" not in sys.modules
    assert "typing_extensions" not in sys.modules


def test_pandas(monkeypatch: pytest.MonkeyPatch) -> None:
    pytest.importorskip("pandas")
    import pandas as pd

    monkeypatch.delitem(sys.modules, "polars", raising=False)
    monkeypatch.delitem(sys.modules, "pyarrow", raising=False)
    monkeypatch.delitem(sys.modules, "duckdb", raising=False)
    monkeypatch.delitem(sys.modules, "dask", raising=False)
    monkeypatch.delitem(sys.modules, "ibis", raising=False)
    monkeypatch.delitem(sys.modules, "pyspark", raising=False)
    df = pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]})
    nw.from_native(df, eager_only=True).group_by("a").agg(nw.col("b").mean()).filter(
        nw.col("a") > 1
    )
    assert "polars" not in sys.modules
    assert "pandas" in sys.modules
    assert "numpy" in sys.modules
    assert "pyarrow" not in sys.modules
    assert "dask" not in sys.modules
    assert "ibis" not in sys.modules
    assert "pyspark" not in sys.modules
    assert "duckdb" not in sys.modules


def test_dask(monkeypatch: pytest.MonkeyPatch) -> None:
    pytest.importorskip("dask")
    import dask.dataframe as dd
    import pandas as pd

    monkeypatch.delitem(sys.modules, "polars", raising=False)
    monkeypatch.delitem(sys.modules, "pyarrow", raising=False)
    monkeypatch.delitem(sys.modules, "duckdb", raising=False)
    monkeypatch.delitem(sys.modules, "pyspark", raising=False)
    df = dd.from_pandas(pd.DataFrame({"a": [1, 1, 2], "b": [4, 5, 6]}))
    nw.from_native(df).group_by("a").agg(nw.col("b").mean()).filter(nw.col("a") > 1)
    assert "polars" not in sys.modules
    assert "pandas" in sys.modules
    assert "numpy" in sys.modules
    assert "dask" in sys.modules
    assert "pyspark" not in sys.modules
    assert "duckdb" not in sys.modules


def test_pyarrow(monkeypatch: pytest.MonkeyPatch) -> None:
    pytest.importorskip("pyarrow")
    import pyarrow as pa

    monkeypatch.delitem(sys.modules, "polars", raising=False)
    monkeypatch.delitem(sys.modules, "pandas", raising=False)
    monkeypatch.delitem(sys.modules, "duckdb", raising=False)
    monkeypatch.delitem(sys.modules, "dask", raising=False)
    monkeypatch.delitem(sys.modules, "ibis", raising=False)
    monkeypatch.delitem(sys.modules, "pyspark", raising=False)
    df = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})
    nw.from_native(df).group_by("a").agg(nw.col("b").mean()).filter(nw.col("a") > 1)
    assert "polars" not in sys.modules
    assert "pandas" not in sys.modules
    assert "pyarrow" in sys.modules
    assert "dask" not in sys.modules
    assert "ibis" not in sys.modules
    assert "pyspark" not in sys.modules
    assert "duckdb" not in sys.modules
