import io
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any

import pytest

import polars as pl
from polars._typing import EngineType
from polars.testing import assert_frame_equal

SINKS = [
    (pl.scan_ipc, pl.LazyFrame.sink_ipc),
    (pl.scan_parquet, pl.LazyFrame.sink_parquet),
    (pl.scan_csv, pl.LazyFrame.sink_csv),
    (pl.scan_ndjson, pl.LazyFrame.sink_ndjson),
]


@pytest.mark.parametrize(("scan", "sink"), SINKS)
@pytest.mark.parametrize("engine", ["in-memory", "streaming"])
@pytest.mark.write_disk
def test_mkdir(tmp_path: Path, scan: Any, sink: Any, engine: EngineType) -> None:
    df = pl.DataFrame(
        {
            "a": [1, 2, 3],
        }
    )

    with pytest.raises(FileNotFoundError):
        sink(df.lazy(), tmp_path / "a" / "b" / "c" / "file", engine=engine)

    f = tmp_path / "a" / "b" / "c" / "file2"
    sink(df.lazy(), f, mkdir=True)

    assert_frame_equal(scan(f).collect(), df)


def test_write_mkdir(tmp_path: Path) -> None:
    df = pl.DataFrame(
        {
            "a": [1, 2, 3],
        }
    )

    with pytest.raises(FileNotFoundError):
        df.write_parquet(tmp_path / "a" / "b" / "c" / "file")

    f = tmp_path / "a" / "b" / "c" / "file2"
    df.write_parquet(f, mkdir=True)

    assert_frame_equal(pl.read_parquet(f), df)


@pytest.mark.parametrize(("scan", "sink"), SINKS)
@pytest.mark.parametrize("engine", ["in-memory", "streaming"])
@pytest.mark.write_disk
def test_lazy_sinks(tmp_path: Path, scan: Any, sink: Any, engine: EngineType) -> None:
    df = pl.DataFrame({"a": [1, 2, 3]})
    lf1 = sink(df.lazy(), tmp_path / "a", lazy=True)
    lf2 = sink(df.lazy(), tmp_path / "b", lazy=True)

    assert not Path(tmp_path / "a").exists()
    assert not Path(tmp_path / "b").exists()

    pl.collect_all([lf1, lf2], engine=engine)

    assert_frame_equal(scan(tmp_path / "a").collect(), df)
    assert_frame_equal(scan(tmp_path / "b").collect(), df)


@pytest.mark.parametrize(
    "sink",
    [
        pl.LazyFrame.sink_ipc,
        pl.LazyFrame.sink_parquet,
        pl.LazyFrame.sink_csv,
        pl.LazyFrame.sink_ndjson,
    ],
)
@pytest.mark.write_disk
def test_double_lazy_error(sink: Any) -> None:
    df = pl.DataFrame({})

    with pytest.raises(
        pl.exceptions.InvalidOperationError,
        match="cannot create a sink on top of another sink",
    ):
        sink(sink(df.lazy(), "a", lazy=True), "b")


@pytest.mark.parametrize(("scan", "sink"), SINKS)
def test_sink_to_memory(sink: Any, scan: Any) -> None:
    df = pl.DataFrame(
        {
            "a": [5, 10, 1996],
        }
    )

    f = io.BytesIO()
    sink(df.lazy(), f)

    f.seek(0)
    assert_frame_equal(
        scan(f).collect(),
        df,
    )


@pytest.mark.parametrize(("scan", "sink"), SINKS)
@pytest.mark.write_disk
def test_sink_to_file(tmp_path: Path, sink: Any, scan: Any) -> None:
    df = pl.DataFrame(
        {
            "a": [5, 10, 1996],
        }
    )

    with (tmp_path / "f").open("w+") as f:
        sink(df.lazy(), f, sync_on_close="all")
        f.seek(0)
        assert_frame_equal(
            scan(f).collect(),
            df,
        )


@pytest.mark.parametrize(("scan", "sink"), SINKS)
def test_sink_empty(sink: Any, scan: Any) -> None:
    df = pl.LazyFrame(data={"col1": ["a"]})

    df_empty = pl.LazyFrame(
        data={"col1": []},
        schema={"col1": str},
    )

    expected = df_empty.join(df, how="cross").collect()
    expected_schema = expected.schema

    kwargs = {}
    if scan == pl.scan_ndjson:
        kwargs["schema"] = expected_schema

    # right empty
    f = io.BytesIO()
    sink(df.join(df_empty, how="cross"), f)
    f.seek(0)
    assert_frame_equal(scan(f, **kwargs), expected.lazy())

    # left empty
    f.seek(0)
    sink(df_empty.join(df, how="cross"), f)
    f.truncate()
    f.seek(0)
    assert_frame_equal(scan(f, **kwargs), expected.lazy())

    # both empty
    f.seek(0)
    sink(df_empty.join(df_empty, how="cross"), f)
    f.truncate()
    f.seek(0)
    assert_frame_equal(scan(f, **kwargs), expected.lazy())


@pytest.mark.parametrize(("scan", "sink"), SINKS)
def test_sink_boolean_panic_25806(sink: Any, scan: Any) -> None:
    df = pl.select(bool=pl.repeat(True, 300_000))

    f = io.BytesIO()
    sink(df.lazy(), f)

    assert_frame_equal(scan(f).collect(), df)


def test_collect_all_lazy() -> None:
    with TemporaryDirectory() as tmpdir:
        tmp_path = Path(tmpdir)

        a = pl.LazyFrame({"a": [1, 2, 3, 4, 5, 6]})
        b = a.filter(pl.col("a") % 2 == 0).sink_csv(tmp_path / "b.csv", lazy=True)
        c = a.filter(pl.col("a") % 3 == 0).sink_csv(tmp_path / "c.csv", lazy=True)
        d = a.sink_csv(tmp_path / "a.csv", lazy=True)

        q = pl.collect_all([d, b, c], lazy=True)

        assert q._ldf._node_name() == "SinkMultiple"  # type: ignore[attr-defined]
        q.collect()
        df_a = pl.read_csv(tmp_path / "a.csv")
        df_b = pl.read_csv(tmp_path / "b.csv")
        df_c = pl.read_csv(tmp_path / "c.csv")

        assert_frame_equal(df_a, pl.DataFrame({"a": [1, 2, 3, 4, 5, 6]}))
        assert_frame_equal(df_b, pl.DataFrame({"a": [2, 4, 6]}))
        assert_frame_equal(df_c, pl.DataFrame({"a": [3, 6]}))

    with pytest.raises(ValueError, match="all LazyFrames must end with a sink to use"):
        pl.collect_all([a, a], lazy=True)
