import os

import lance
import pyarrow as pa
import pytest
from packaging.version import Version
from pytest_lazy_fixtures import lf as lazy_fixture

import ray
from ray._common.test_utils import wait_for_condition
from ray.data import Schema
from ray.data._internal.datasource.lance_datasink import (
    _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
    LanceDatasink,
    _write_fragment,
)
from ray.data.datasource import SaveMode
from ray.data.datasource.path_util import _unwrap_protocol

# Skip tests for older pylance versions (<=0.3.19) due to incompatible lance API changes with pyarrow v9.0.0
pytestmark = pytest.mark.skipif(
    Version(lance.__version__) <= Version("0.3.19"),
    reason=f"pylance {lance.__version__} <= 0.3.19; API incompatible",
)


@pytest.mark.parametrize(
    "fs,data_path",
    [
        (None, lazy_fixture("local_path")),
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
        (
            lazy_fixture("s3_fs_with_space"),
            lazy_fixture("s3_path_with_space"),
        ),  # Path contains space.
        (
            lazy_fixture("s3_fs_with_anonymous_crendential"),
            lazy_fixture("s3_path_with_anonymous_crendential"),
        ),
    ],
)
@pytest.mark.parametrize(
    "batch_size",
    [None, 100],
)
def test_lance_read_basic(fs, data_path, batch_size):
    df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
    setup_data_path = _unwrap_protocol(data_path)
    path = os.path.join(setup_data_path, "test.lance")
    lance.write_dataset(df1, path)

    ds_lance = lance.dataset(path)
    assert ds_lance is not None
    df2 = pa.table(
        {
            "one": [1, 2, 3, 4, 5, 6],
            "three": [4, 5, 8, 9, 12, 13],
            "four": ["u", "v", "w", "x", "y", "z"],
        }
    )
    ds_lance.merge(df2, "one")

    if batch_size is None:
        ds = ray.data.read_lance(path)
    else:
        ds = ray.data.read_lance(path, scanner_options={"batch_size": batch_size})

    # Test metadata-only ops.
    assert ds.count() == 6
    assert ds.schema() == Schema(
        pa.schema(
            {
                "one": pa.int64(),
                "two": pa.string(),
                "three": pa.int64(),
                "four": pa.string(),
            }
        )
    )

    # Test read.
    values = [[s["one"], s["two"]] for s in ds.take_all()]
    assert sorted(values) == [
        [1, "a"],
        [2, "b"],
        [3, "c"],
        [4, "e"],
        [5, "f"],
        [6, "g"],
    ]

    # Test column projection.
    ds = ray.data.read_lance(path, columns=["one"])
    values = [s["one"] for s in ds.take_all()]
    assert sorted(values) == [1, 2, 3, 4, 5, 6]
    assert ds.schema().names == ["one"]


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_read_with_scanner_fragments(data_path):
    table = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
    setup_data_path = _unwrap_protocol(data_path)
    path = os.path.join(setup_data_path, "test.lance")
    dataset = lance.write_dataset(table, path, max_rows_per_file=2)
    assert dataset is not None

    fragments = dataset.get_fragments()
    ds = ray.data.read_lance(path, scanner_options={"fragments": fragments[:1]})
    values = [[s["one"], s["two"]] for s in ds.take_all()]
    assert values == [
        [2, "b"],
        [1, "a"],
    ]


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_read_many_files(data_path):
    setup_data_path = _unwrap_protocol(data_path)
    path = os.path.join(setup_data_path, "test.lance")
    num_rows = 1024
    data = pa.table({"id": pa.array(range(num_rows))})
    lance.write_dataset(data, path, max_rows_per_file=1)

    def test_lance():
        ds = ray.data.read_lance(path)
        return ds.count() == num_rows

    wait_for_condition(test_lance, timeout=10)


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_write(data_path):
    schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])

    ray.data.range(10).map(
        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
    ).write_lance(data_path, schema=schema)

    ds = lance.dataset(data_path)
    assert ds is not None
    ds.count_rows() == 10
    assert ds.schema.names == schema.names
    # The schema is platform-dependent, because numpy uses int32 on Windows.
    # So we observe the schema that is written and use that.
    schema = ds.schema

    tbl = ds.to_table()
    assert sorted(tbl["id"].to_pylist()) == list(range(10))
    assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(10)}

    ray.data.range(10).map(
        lambda x: {"id": x["id"] + 10, "str": f"str-{x['id'] + 10}"}
    ).write_lance(data_path, mode=SaveMode.APPEND)

    ds = lance.dataset(data_path)
    assert ds is not None
    ds.count_rows() == 20
    tbl = ds.to_table()
    assert sorted(tbl["id"].to_pylist()) == list(range(20))
    assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(20)}

    ray.data.range(10).map(
        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
    ).write_lance(data_path, schema=schema, mode=SaveMode.OVERWRITE)

    ds = lance.dataset(data_path)
    assert ds is not None
    ds.count_rows() == 10
    assert ds.schema == schema


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_write_create_errors_if_exists(data_path):
    table_path = os.path.join(data_path, "my_table")
    ds = ray.data.range(10)

    # First CREATE succeeds on an empty destination.
    ds.write_lance(table_path, mode=SaveMode.CREATE)
    assert lance.dataset(table_path).count_rows() == 10

    # A second CREATE must error instead of silently overwriting.
    with pytest.raises(ValueError, match="already exists"):
        ray.data.range(5).write_lance(table_path, mode=SaveMode.CREATE)

    # Existing data is untouched.
    assert lance.dataset(table_path).count_rows() == 10

    # CREATE is also the default mode, so it must guard too.
    with pytest.raises(ValueError, match="already exists"):
        ray.data.range(5).write_lance(table_path)
    assert lance.dataset(table_path).count_rows() == 10

    # OVERWRITE replaces the existing data.
    ray.data.range(5).write_lance(table_path, mode=SaveMode.OVERWRITE)
    assert lance.dataset(table_path).count_rows() == 5


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_write_append_errors_if_missing(data_path):
    table_path = os.path.join(data_path, "missing_table")
    # APPEND surfaces Lance's own "not found" error. We don't pin the message,
    # since it can change across Lance versions.
    expected_errors: tuple[type[Exception], ...] = (
        ValueError,
        OSError,
        FileNotFoundError,
    )
    with pytest.raises(expected_errors):
        ray.data.range(5).write_lance(table_path, mode=SaveMode.APPEND)
    assert not os.path.exists(table_path)


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_write_min_rows_per_file(data_path):
    schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])

    ray.data.range(10).map(
        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
    ).write_lance(data_path, schema=schema, min_rows_per_file=100)

    ds = lance.dataset(data_path)
    assert ds is not None
    assert ds.count_rows() == 10
    assert ds.schema == schema

    assert len(ds.get_fragments()) == 1


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_write_max_rows_per_file(data_path):
    schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])

    ray.data.range(10).map(
        lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
    ).write_lance(data_path, schema=schema, max_rows_per_file=1)

    ds = lance.dataset(data_path)
    assert ds is not None
    assert ds.count_rows() == 10
    assert ds.schema == schema

    assert len(ds.get_fragments()) == 10


@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
def test_lance_read_with_version(data_path):
    # Write an initial dataset (version 1)
    df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
    setup_data_path = _unwrap_protocol(data_path)
    path = os.path.join(setup_data_path, "test_version.lance")
    lance.write_dataset(df1, path)

    # Merge new data to create a later version (latest)
    ds_lance = lance.dataset(path)
    assert ds_lance is not None
    # Get the initial version
    initial_version = ds_lance.version

    df2 = pa.table(
        {
            "one": [1, 2, 3, 4, 5, 6],
            "three": [4, 5, 8, 9, 12, 13],
            "four": ["u", "v", "w", "x", "y", "z"],
        }
    )
    ds_lance.merge(df2, "one")

    # Default read should return the latest (merged) dataset.
    ds_latest = ray.data.read_lance(path)

    assert ds_latest.count() == 6
    # Latest dataset should contain merged columns
    assert "three" in ds_latest.schema().names

    # Read the initial version and ensure it contains the original columns
    ds_prev = ray.data.read_lance(path, version=initial_version)
    assert ds_prev.count() == 6
    assert ds_prev.schema().names == ["one", "two"]

    values_prev = [[s["one"], s["two"]] for s in ds_prev.take_all()]
    assert sorted(values_prev) == [
        [1, "a"],
        [2, "b"],
        [3, "c"],
        [4, "e"],
        [5, "f"],
        [6, "g"],
    ]


@pytest.fixture
def mock_lance_write(monkeypatch):
    captured = {}

    class _FakeLanceDatasink:
        def __init__(self, path, **kwargs):
            captured["path"] = path
            captured["kwargs"] = kwargs

    def _fake_write_datasink(self, datasink, **kwargs):
        captured["datasink"] = datasink
        captured["write_kwargs"] = kwargs

    monkeypatch.setattr(ray.data.dataset, "LanceDatasink", _FakeLanceDatasink)
    monkeypatch.setattr(ray.data.Dataset, "write_datasink", _fake_write_datasink)

    return captured, _FakeLanceDatasink


def test_write_lance_passes_namespace_args(mock_lance_write):
    captured, fake_lance_datasink_cls = mock_lance_write
    table_id = ["db", "table"]
    namespace_impl = "dir"
    namespace_properties = {"path": "/tmp/ns"}

    ds = ray.data.range(1)
    ds.write_lance(
        "/tmp/lance-namespace-test",
        table_id=table_id,
        namespace_impl=namespace_impl,
        namespace_properties=namespace_properties,
    )

    assert captured["path"] == "/tmp/lance-namespace-test"
    assert captured["kwargs"]["table_id"] == table_id
    assert captured["kwargs"]["namespace_impl"] == namespace_impl
    assert captured["kwargs"]["namespace_properties"] == namespace_properties
    assert isinstance(captured["datasink"], fake_lance_datasink_cls)


@pytest.mark.parametrize("mode", [SaveMode.APPEND, SaveMode.OVERWRITE])
def test_lance_namespace_write_rejects_non_create_mode(monkeypatch, mode):
    class _FakeNamespace:
        pass

    monkeypatch.setattr(
        "ray.data._internal.datasource.lance_datasink.get_or_create_namespace",
        lambda namespace_impl, namespace_properties: _FakeNamespace(),
    )

    with pytest.raises(ValueError, match="Namespace writes currently only support"):
        LanceDatasink(
            uri="/tmp/lance-namespace-test",
            mode=mode,
            table_id=["db", "table"],
            namespace_impl="dir",
            namespace_properties={"path": "/tmp/ns"},
        )


@pytest.mark.parametrize(
    "max_attempts,expected_blocks_consumed_before_write",
    [(1, 1), (2, 3)],
)
def test_write_fragment_only_materializes_stream_when_retrying(
    monkeypatch, max_attempts, expected_blocks_consumed_before_write
):
    import lance.fragment

    consumed = {"count": 0}
    blocks = [pa.table({"id": [i]}) for i in range(3)]

    def block_stream():
        for block in blocks:
            consumed["count"] += 1
            yield block

    def fake_write_fragments(reader, uri, **kwargs):
        assert consumed["count"] == expected_blocks_consumed_before_write
        return []

    monkeypatch.setattr(lance.fragment, "write_fragments", fake_write_fragments)

    _write_fragment(
        block_stream(),
        "/tmp/lance-materialization-test",
        retry_params={
            "description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
            "match": [],
            "max_attempts": max_attempts,
            "max_backoff_s": 0,
        },
    )


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))