from __future__ import annotations

from typing import TYPE_CHECKING
from typing import Any
from typing import Literal
from typing import Mapping

import pandas as pd
import pytest

import narwhals as nw
import narwhals.stable.v1 as nw_v1
from narwhals.utils import Implementation
from tests.utils import PANDAS_VERSION
from tests.utils import Constructor
from tests.utils import ConstructorEager
from tests.utils import assert_equal_data

pytest.importorskip("polars")
import polars as pl

if TYPE_CHECKING:
    from types import ModuleType

data: Mapping[str, Any] = {"a": [1, 2, 3], "b": [4.5, 6.7, 8.9], "z": ["x", "y", "w"]}
TEST_EAGER_BACKENDS = [
    Implementation.POLARS,
    Implementation.PANDAS,
    Implementation.PYARROW,
    "polars",
    "pandas",
    "pyarrow",
]


@pytest.mark.parametrize("backend", TEST_EAGER_BACKENDS)
@pytest.mark.parametrize("nw_namespace", [nw, nw_v1])
def test_read_csv(
    tmpdir: pytest.TempdirFactory,
    backend: Implementation | str,
    nw_namespace: ModuleType,
) -> None:
    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.csv")  # type: ignore[operator]
    df_pl.write_csv(filepath)
    result = nw_namespace.read_csv(filepath, backend=backend)
    assert_equal_data(result, data)
    assert isinstance(result, nw_namespace.DataFrame)


@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="too old for pyarrow")
def test_read_csv_kwargs(tmpdir: pytest.TempdirFactory) -> None:
    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.csv")  # type: ignore[operator]
    df_pl.write_csv(filepath)
    result = nw.read_csv(filepath, backend=pd, engine="pyarrow")
    assert_equal_data(result, data)


@pytest.mark.parametrize("nw_namespace", [nw, nw_v1])
def test_scan_csv(
    tmpdir: pytest.TempdirFactory,
    constructor: Constructor,
    nw_namespace: ModuleType,
) -> None:
    kwargs: dict[str, Any]
    if "sqlframe" in str(constructor):
        from sqlframe.duckdb import DuckDBSession

        kwargs = {
            "session": DuckDBSession(),
            "inferSchema": True,
            "header": True,
        }
    elif "pyspark" in str(constructor):
        from pyspark.sql import SparkSession

        kwargs = {
            "session": (
                SparkSession.builder.appName("unit-tests")  # pyright: ignore[reportAttributeAccessIssue]
                .master("local[1]")
                .config("spark.ui.enabled", "false")
                .config("spark.default.parallelism", "1")
                .config("spark.sql.shuffle.partitions", "2")
                .config("spark.sql.session.timeZone", "UTC")
                .getOrCreate()
            ),
            "inferSchema": True,
            "header": True,
        }
    else:
        kwargs = {}

    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.csv")  # type: ignore[operator]
    df_pl.write_csv(filepath)
    df = nw_namespace.from_native(constructor(data))
    backend = nw_namespace.get_native_namespace(df)
    result = nw_namespace.scan_csv(filepath, backend=backend, **kwargs)
    assert_equal_data(result, data)
    assert isinstance(result, nw_namespace.LazyFrame)


@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="too old for pyarrow")
def test_scan_csv_kwargs(tmpdir: pytest.TempdirFactory) -> None:
    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.csv")  # type: ignore[operator]
    df_pl.write_csv(filepath)
    result = nw.scan_csv(filepath, backend=pd, engine="pyarrow")
    assert_equal_data(result, data)


@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="too old for pyarrow")
@pytest.mark.parametrize("nw_namespace", [nw, nw_v1])
def test_read_parquet(
    tmpdir: pytest.TempdirFactory,
    constructor_eager: ConstructorEager,
    nw_namespace: ModuleType,
) -> None:
    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.parquet")  # type: ignore[operator]
    df_pl.write_parquet(filepath)
    df = nw_namespace.from_native(constructor_eager(data))
    backend = nw_namespace.get_native_namespace(df)
    result = nw_namespace.read_parquet(filepath, backend=backend)
    assert_equal_data(result, data)
    assert isinstance(result, nw_namespace.DataFrame)


@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="too old for pyarrow")
def test_read_parquet_kwargs(tmpdir: pytest.TempdirFactory) -> None:
    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.parquet")  # type: ignore[operator]
    df_pl.write_parquet(filepath)
    result = nw.read_parquet(filepath, backend=pd, engine="pyarrow")
    assert_equal_data(result, data)


@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="too old for pyarrow")
@pytest.mark.parametrize("nw_namespace", [nw, nw_v1])
def test_scan_parquet(
    tmpdir: pytest.TempdirFactory,
    constructor: Constructor,
    nw_namespace: ModuleType,
) -> None:
    kwargs: dict[str, Any]
    if "sqlframe" in str(constructor):
        from sqlframe.duckdb import DuckDBSession

        kwargs = {"session": DuckDBSession(), "inferSchema": True}
    elif "pyspark" in str(constructor):
        from pyspark.sql import SparkSession

        kwargs = {
            "session": (
                SparkSession.builder.appName("unit-tests")  # pyright: ignore[reportAttributeAccessIssue]
                .master("local[1]")
                .config("spark.ui.enabled", "false")
                .config("spark.default.parallelism", "1")
                .config("spark.sql.shuffle.partitions", "2")
                .config("spark.sql.session.timeZone", "UTC")
                .getOrCreate()
            ),
            "inferSchema": True,
        }
    else:
        kwargs = {}
    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.parquet")  # type: ignore[operator]
    df_pl.write_parquet(filepath)
    df = nw_namespace.from_native(constructor(data))
    backend = nw_namespace.get_native_namespace(df)
    result = nw_namespace.scan_parquet(filepath, backend=backend, **kwargs)
    assert_equal_data(result, data)
    assert isinstance(result, nw_namespace.LazyFrame)


@pytest.mark.skipif(PANDAS_VERSION < (1, 5), reason="too old for pyarrow")
def test_scan_parquet_kwargs(tmpdir: pytest.TempdirFactory) -> None:
    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.parquet")  # type: ignore[operator]
    df_pl.write_parquet(filepath)
    result = nw.scan_parquet(filepath, backend=pd, engine="pyarrow")
    assert_equal_data(result, data)


@pytest.mark.parametrize("spark_like_backend", ["pyspark", "sqlframe"])
@pytest.mark.parametrize("scan_method", ["scan_csv", "scan_parquet"])
@pytest.mark.parametrize("nw_namespace", [nw, nw_v1])
def test_scan_fail_spark_like_without_session(
    tmpdir: pytest.TempdirFactory,
    spark_like_backend: str,
    scan_method: Literal["scan_csv", "scan_parquet"],
    nw_namespace: ModuleType,
) -> None:
    _ = pytest.importorskip(spark_like_backend)

    df_pl = pl.DataFrame(data)
    filepath = str(tmpdir / "file.parquet")  # type: ignore[operator]
    df_pl.write_parquet(filepath)

    with pytest.raises(
        ValueError,
        match="Spark like backends require a session object to be passed in `kwargs`.",
    ):
        getattr(nw_namespace, scan_method)("file.parquet", backend=spark_like_backend)
