# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Lance Authors

from pathlib import Path

import lance
import numpy as np
import pyarrow as pa
import pytest
from lance.sampler import ShardedBatchSampler, ShardedFixedBatchSampler, maybe_sample

TEST_CONFIG = {
    "total_rows": 1000,
    "batch_size": 250,
    "world_size": 4,
    "vec_dim": 32,
    "test_port": "29501",
    "master_addr": "127.0.0.1",
    "seed": 42,
    "test_shard_ratio": 0.5,
    "max_takes_factor": 0.1,
}


@pytest.fixture
def sample_dataset_path(tmp_path):
    data = pa.Table.from_arrays(
        [
            pa.array(range(TEST_CONFIG["total_rows"])),
            pa.array(np.random.rand(TEST_CONFIG["total_rows"])),
            pa.array([f"text_{i}" for i in range(TEST_CONFIG["total_rows"])]),
        ],
        names=["id", "value", "text"],
    )

    dataset_path = tmp_path / "test_dataset.lance"
    lance.write_dataset(data, dataset_path)
    return dataset_path


@pytest.fixture
def sample_dataset(sample_dataset_path) -> lance.LanceDataset:
    return lance.dataset(sample_dataset_path)


def test_consecutive_index_blocks():
    sampler = ShardedFixedBatchSampler(
        rank=0,
        world_size=TEST_CONFIG["world_size"],
        total_num_rows=TEST_CONFIG["total_rows"],
        batch_size=TEST_CONFIG["batch_size"],
    )

    batches = list(sampler)
    expected_size = TEST_CONFIG["total_rows"] // (
        TEST_CONFIG["world_size"] * TEST_CONFIG["batch_size"]
    )
    assert len(batches) == expected_size
    assert batches[0] == list(range(TEST_CONFIG["batch_size"]))


def _distributed_test_worker(rank, world_size, dataset_path):
    import os

    import torch

    os.environ.update(
        {
            "MASTER_ADDR": TEST_CONFIG["master_addr"],
            "MASTER_PORT": TEST_CONFIG["test_port"],
            "CUDA_VISIBLE_DEVICES": ",".join(
                map(str, range(torch.cuda.device_count()))
            ),
        }
    )

    try:
        if torch.cuda.is_available():
            torch.cuda.set_device(rank % torch.cuda.device_count())

        backend = "nccl" if torch.cuda.is_available() else "gloo"
        torch.distributed.init_process_group(
            backend=backend, world_size=world_size, rank=rank
        )

        dataset = lance.dataset(dataset_path)
        assert len(dataset) == TEST_CONFIG["total_rows"]

        sampler = ShardedBatchSampler(
            rank=rank,
            world_size=world_size,
            total_num_rows=TEST_CONFIG["total_rows"],
            batch_size=TEST_CONFIG["batch_size"],
        )

        class DatasetAdapter(torch.utils.data.Dataset):
            def __init__(self, dataset):
                self.dataset = dataset

            def __getitem__(self, index):
                return self.dataset.take([index], ["id", "value"]).to_pylist()[0]

            def __len__(self):
                return len(self.dataset)

        def collate_fn(batch):
            return {
                "ids": torch.tensor([x["id"] for x in batch], dtype=torch.long),
                "values": torch.tensor(
                    [x["value"] for x in batch], dtype=torch.float32
                ),
            }

        dataloader = torch.utils.data.DataLoader(
            DatasetAdapter(dataset),
            batch_sampler=sampler,
            collate_fn=collate_fn,
            num_workers=0,
        )

        total = 0
        for batch_indices, batch_data in zip(sampler, dataloader):
            current_size = batch_data["ids"].size(0)
            assert current_size == TEST_CONFIG["batch_size"]
            assert batch_data["ids"].tolist() == list(batch_indices)
            total += current_size

        expected_total = TEST_CONFIG["total_rows"] // world_size
        assert total == expected_total

    finally:
        if torch.distributed.is_initialized():
            torch.distributed.destroy_process_group()


@pytest.mark.cuda
def test_pytorch_integration(sample_dataset_path):
    import torch

    test_world_sizes = [1, 2] if torch.cuda.device_count() >= 2 else [1]
    for ws in test_world_sizes:
        torch.multiprocessing.spawn(
            _distributed_test_worker,
            args=(ws, str(sample_dataset_path)),
            nprocs=ws,
            join=True,
        )


def test_data_stream_without_filter(sample_dataset):
    """Validate direct data loading without filters."""
    sampler = ShardedFixedBatchSampler(0, 4)
    batches = list(sampler(sample_dataset, batch_size=250, columns=["id", "value"]))

    # Data integrity checks
    batch = batches[0]
    assert batch.num_rows == 250, "Batch should contain 250 records"
    assert batch.column_names == ["id", "value"], "Should load specified columns"

    # Consecutive ID validation
    ids = batch["id"].to_numpy()
    assert np.array_equal(ids, np.arange(0, 250)), "IDs should be sequential 0-249"


def test_filtered_data_handling(sample_dataset):
    """Test filtered data processing with sharding."""
    # Apply ID filter and load data
    sampler = ShardedFixedBatchSampler(0, 4)
    batches = list(
        sampler(sample_dataset, batch_size=100, filter="id < 500", columns=["id"])
    )

    # Aggregated results validation
    all_ids = []
    for batch in batches:
        all_ids.extend(batch["id"].to_numpy().tolist())

    # Filter and sharding assertions
    assert all(id_val < 500 for id_val in all_ids), "Should respect ID filter"
    assert all(id_val % 4 == 0 for id_val in all_ids), "Should keep rank 0 shard"


def test_randomization_effect():
    """Verify epoch-based randomization behavior."""
    # Initialize randomized sampler
    sampler = ShardedFixedBatchSampler(
        rank=0,
        world_size=4,
        total_num_rows=2000,
        batch_size=250,
        randomize=True,
        seed=42,
    )

    assert len(list(sampler)) > 1

    # Cross-epoch comparison
    sampler.set_epoch(1)
    epoch1 = list(sampler)
    sampler.set_epoch(2)
    epoch2 = list(sampler)

    assert epoch1 != epoch2, "Different epochs should produce different orders"


def test_edge_cases():
    """Validate handling of partial batches and data boundaries."""

    sampler = ShardedFixedBatchSampler(
        rank=3, world_size=4, batch_size=250, total_num_rows=1000
    )
    batches = list(sampler)
    assert len(batches) == 1, "Should handle partial batch"
    assert batches[0] == list(range(750, 1000)), "Last rank should get 750-999"

    sampler = ShardedFixedBatchSampler(
        rank=0, world_size=2, batch_size=128, total_num_rows=500
    )
    batches = list(sampler)
    # rank 0: 0~249, rank 1: 250~499
    # rank 0: [0-127], [128-249]
    assert batches[0] == list(range(0, 128))
    assert batches[1] == list(range(128, 250))

    # total_num_rows < batch_size
    sampler = ShardedFixedBatchSampler(
        rank=0, world_size=1, batch_size=250, total_num_rows=100
    )
    batches = list(sampler)
    assert len(batches) == 1
    assert batches[0] == list(range(0, 100))

    # total_num_rows < world_size
    sampler = ShardedFixedBatchSampler(
        rank=2, world_size=4, batch_size=10, total_num_rows=2
    )
    batches = list(sampler)
    assert len(batches) == 0, "No data for this rank"

    # batch_size=1
    sampler = ShardedFixedBatchSampler(
        rank=0, world_size=2, batch_size=1, total_num_rows=4
    )
    batches = list(sampler)
    assert batches == [[0], [1]]

    # world_size=1
    sampler = ShardedFixedBatchSampler(
        rank=0, world_size=1, batch_size=3, total_num_rows=5
    )
    batches = list(sampler)
    assert batches == [list(range(0, 3)), list(range(3, 5))]


# We use + 97 to test case where num_rows and chunk_size aren't exactly aligned.
@pytest.mark.parametrize("nrows", [10, 10240, 10240 + 97, 10240 + 1024])
def test_sample_dataset(tmp_path: Path, nrows: int):
    # nrows of 32-d vectors.
    data = np.random.random(nrows * 32).astype("f")

    fsl = pa.FixedSizeListArray.from_arrays(data, 32)
    tbl = pa.Table.from_arrays([fsl], ["vec"])

    ds = lance.write_dataset(tbl, tmp_path / "data.lance")

    # Simple path
    simple_scan = list(maybe_sample(ds, 128, ["vec"]))
    assert len(simple_scan) == 1
    assert isinstance(simple_scan[0], pa.RecordBatch)
    assert simple_scan[0].schema == pa.schema([pa.field("vec", fsl.type)])
    assert simple_scan[0].num_rows == min(nrows, 128)

    # Random path.
    large_scan = list(maybe_sample(ds, 128, ["vec"], max_takes=32))
    assert len(large_scan) == 1
    assert isinstance(large_scan[0], pa.RecordBatch)
    assert large_scan[0].schema == pa.schema([pa.field("vec", fsl.type)])
    assert large_scan[0].num_rows == min(nrows, 128)
