import logging
import os
import posixpath
import urllib.parse
from functools import partial
from unittest.mock import patch

import pandas as pd
import pytest
from pyarrow.fs import LocalFileSystem
from pytest_lazy_fixtures import lf as lazy_fixture

from ray.data.datasource import (
    BaseFileMetadataProvider,
    DefaultFileMetadataProvider,
    FileMetadataProvider,
)
from ray.data.datasource.file_based_datasource import (
    FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD,
)
from ray.data.datasource.file_meta_provider import (
    _get_file_infos_common_path_prefix,
    _get_file_infos_parallel,
    _get_file_infos_serial,
)
from ray.data.datasource.path_util import (
    _resolve_paths_and_filesystem,
    _unwrap_protocol,
)
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.test_partitioning import PathPartitionEncoder
from ray.tests.conftest import *  # noqa


def df_to_csv(dataframe, path, **kwargs):
    dataframe.to_csv(path, **kwargs)


def _get_file_sizes_bytes(paths, fs):
    from pyarrow.fs import FileType

    file_sizes = []
    for path in paths:
        file_info = fs.get_file_info(path)
        if file_info.type == FileType.File:
            file_sizes.append(file_info.size)
        else:
            raise FileNotFoundError(path)
    return file_sizes


def test_file_metadata_providers_not_implemented():
    meta_provider = FileMetadataProvider()
    with pytest.raises(NotImplementedError):
        meta_provider(["/foo/bar.csv"])
    meta_provider = BaseFileMetadataProvider()
    with pytest.raises(NotImplementedError):
        meta_provider(["/foo/bar.csv"], rows_per_file=None, file_sizes=[None])
    with pytest.raises(NotImplementedError):
        meta_provider.expand_paths(["/foo/bar.csv"], None)


@pytest.mark.parametrize(
    "fs,data_path,endpoint_url",
    [
        (None, lazy_fixture("local_path"), None),
        (lazy_fixture("local_fs"), lazy_fixture("local_path"), None),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path"), lazy_fixture("s3_server")),
        (
            lazy_fixture("s3_fs_with_space"),
            lazy_fixture("s3_path_with_space"),
            lazy_fixture("s3_server"),
        ),  # Path contains space.
        (
            lazy_fixture("s3_fs_with_special_chars"),
            lazy_fixture("s3_path_with_special_chars"),
            lazy_fixture("s3_server"),
        ),
    ],
)
def test_default_file_metadata_provider(
    propagate_logs, caplog, fs, data_path, endpoint_url
):
    storage_options = (
        {}
        if endpoint_url is None
        else dict(client_kwargs=dict(endpoint_url=endpoint_url))
    )

    path_module = os.path if urllib.parse.urlparse(data_path).scheme else posixpath
    path1 = path_module.join(data_path, "test1.csv")
    path2 = path_module.join(data_path, "test2.csv")
    paths = [path1, path2]
    paths, fs = _resolve_paths_and_filesystem(paths, fs)

    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
    df1.to_csv(path1, index=False, storage_options=storage_options)
    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
    df2.to_csv(path2, index=False, storage_options=storage_options)

    meta_provider = DefaultFileMetadataProvider()
    with caplog.at_level(logging.WARNING), patch(
        "ray.data.datasource.file_meta_provider._get_file_infos_serial",
        wraps=_get_file_infos_serial,
    ) as mock_get:
        file_paths, file_sizes = map(list, zip(*meta_provider.expand_paths(paths, fs)))
    mock_get.assert_called_once_with(paths, fs, False)
    # No warning should be logged.
    assert len(caplog.text) == 0
    assert file_paths == paths
    expected_file_sizes = _get_file_sizes_bytes(paths, fs)
    assert file_sizes == expected_file_sizes

    meta = meta_provider(
        paths,
        rows_per_file=3,
        file_sizes=file_sizes,
    )
    assert meta.size_bytes == sum(expected_file_sizes)
    assert meta.num_rows == 6
    assert len(paths) == 2
    assert all(path in meta.input_files for path in paths)


@pytest.mark.parametrize(
    "fs,data_path,endpoint_url",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path"), None),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path"), lazy_fixture("s3_server")),
    ],
)
def test_default_metadata_provider_ignore_missing(fs, data_path, endpoint_url):
    storage_options = (
        {}
        if endpoint_url is None
        else dict(client_kwargs=dict(endpoint_url=endpoint_url))
    )

    path1 = os.path.join(data_path, "test1.csv")
    path2 = os.path.join(data_path, "test2.csv")
    paths = [path1, path2]
    paths_with_missing = paths + [os.path.join(data_path, "missing.csv")]
    paths_with_missing, fs = _resolve_paths_and_filesystem(paths_with_missing, fs)

    df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
    df1.to_csv(path1, index=False, storage_options=storage_options)
    df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
    df2.to_csv(path2, index=False, storage_options=storage_options)

    meta_provider = DefaultFileMetadataProvider()
    file_paths, _ = map(
        list,
        zip(
            *meta_provider.expand_paths(
                paths_with_missing, fs, ignore_missing_paths=True
            )
        ),
    )

    assert len(file_paths) == 2


@pytest.mark.parametrize(
    "fs,data_path,endpoint_url",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path"), None),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path"), lazy_fixture("s3_server")),
    ],
)
def test_default_file_metadata_provider_many_files_basic(
    propagate_logs,
    caplog,
    fs,
    data_path,
    endpoint_url,
):
    if endpoint_url is None:
        storage_options = {}
    else:
        storage_options = dict(client_kwargs=dict(endpoint_url=endpoint_url))

    paths = []
    dfs = []
    num_dfs = 4 * FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
    for i in range(num_dfs):
        df = pd.DataFrame({"one": list(range(i * 3, (i + 1) * 3))})
        dfs.append(df)
        path = os.path.join(data_path, f"test_{i}.csv")
        if i % 4 == 0:
            # Append same path multiple times to test duplicated paths.
            paths.extend([path, path, path])
        else:
            paths.append(path)
        df.to_csv(path, index=False, storage_options=storage_options)
    paths, fs = _resolve_paths_and_filesystem(paths, fs)

    meta_provider = DefaultFileMetadataProvider()
    if isinstance(fs, LocalFileSystem):
        patcher = patch(
            "ray.data.datasource.file_meta_provider._get_file_infos_serial",
            wraps=_get_file_infos_serial,
        )
    else:
        patcher = patch(
            "ray.data.datasource.file_meta_provider._get_file_infos_common_path_prefix",
            wraps=_get_file_infos_common_path_prefix,
        )
    with caplog.at_level(logging.WARNING), patcher as mock_get:
        file_paths, file_sizes = map(list, zip(*meta_provider.expand_paths(paths, fs)))
    if isinstance(fs, LocalFileSystem):
        mock_get.assert_called_once_with(paths, fs, False)
    else:
        mock_get.assert_called_once_with(paths, _unwrap_protocol(data_path), fs, False)
    # No warning should be logged.
    assert len(caplog.text) == 0
    assert file_paths == paths
    expected_file_sizes = _get_file_sizes_bytes(paths, fs)
    assert file_sizes == expected_file_sizes


@pytest.mark.parametrize(
    "fs,data_path,endpoint_url",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path"), None),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path"), lazy_fixture("s3_server")),
    ],
)
def test_default_file_metadata_provider_many_files_partitioned(
    propagate_logs,
    caplog,
    fs,
    data_path,
    endpoint_url,
    write_partitioned_df,
):
    if endpoint_url is None:
        storage_options = {}
    else:
        storage_options = dict(client_kwargs=dict(endpoint_url=endpoint_url))

    partition_keys = ["one"]
    partition_path_encoder = PathPartitionEncoder.of(
        base_dir=data_path,
        field_names=partition_keys,
        filesystem=fs,
    )
    paths = []
    dfs = []
    num_dfs = FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
    for i in range(num_dfs):
        df = pd.DataFrame(
            {"one": [1, 1, 1, 3, 3, 3], "two": list(range(6 * i, 6 * (i + 1)))}
        )
        df_paths = write_partitioned_df(
            df,
            partition_keys,
            partition_path_encoder,
            partial(df_to_csv, storage_options=storage_options, index=False),
            file_name_suffix=i,
        )
        dfs.append(df)
        paths.extend(df_paths)
    paths, fs = _resolve_paths_and_filesystem(paths, fs)
    partitioning = partition_path_encoder.scheme

    meta_provider = DefaultFileMetadataProvider()
    if isinstance(fs, LocalFileSystem):
        patcher = patch(
            "ray.data.datasource.file_meta_provider._get_file_infos_serial",
            wraps=_get_file_infos_serial,
        )
    else:
        patcher = patch(
            "ray.data.datasource.file_meta_provider._get_file_infos_common_path_prefix",
            wraps=_get_file_infos_common_path_prefix,
        )
    with caplog.at_level(logging.WARNING), patcher as mock_get:
        file_paths, file_sizes = map(
            list, zip(*meta_provider.expand_paths(paths, fs, partitioning))
        )
    if isinstance(fs, LocalFileSystem):
        mock_get.assert_called_once_with(paths, fs, False)
    else:
        mock_get.assert_called_once_with(
            paths, _unwrap_protocol(partitioning.base_dir), fs, False
        )
    assert len(caplog.text) == 0
    assert file_paths == paths
    expected_file_sizes = _get_file_sizes_bytes(paths, fs)
    assert file_sizes == expected_file_sizes


@pytest.mark.parametrize(
    "fs,data_path,endpoint_url",
    [
        (lazy_fixture("local_fs"), lazy_fixture("local_path"), None),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path"), lazy_fixture("s3_server")),
    ],
)
def test_default_file_metadata_provider_many_files_diff_dirs(
    ray_start_regular,
    propagate_logs,
    caplog,
    fs,
    data_path,
    endpoint_url,
):
    if endpoint_url is None:
        storage_options = {}
    else:
        storage_options = dict(client_kwargs=dict(endpoint_url=endpoint_url))

    dir1 = os.path.join(data_path, "dir1")
    dir2 = os.path.join(data_path, "dir2")
    if fs is None:
        os.mkdir(dir1)
        os.mkdir(dir2)
    else:
        fs.create_dir(_unwrap_protocol(dir1))
        fs.create_dir(_unwrap_protocol(dir2))

    paths = []
    dfs = []
    num_dfs = 2 * FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
    for i, dir_path in enumerate([dir1, dir2]):
        for j in range(num_dfs * i, num_dfs * (i + 1)):
            df = pd.DataFrame({"one": list(range(3 * j, 3 * (j + 1)))})
            dfs.append(df)
            path = os.path.join(dir_path, f"test_{j}.csv")
            paths.append(path)
            df.to_csv(path, index=False, storage_options=storage_options)
    paths, fs = _resolve_paths_and_filesystem(paths, fs)

    meta_provider = DefaultFileMetadataProvider()
    if isinstance(fs, LocalFileSystem):
        patcher = patch(
            "ray.data.datasource.file_meta_provider._get_file_infos_serial",
            wraps=_get_file_infos_serial,
        )
    else:
        patcher = patch(
            "ray.data.datasource.file_meta_provider._get_file_infos_parallel",
            wraps=_get_file_infos_parallel,
        )
    with caplog.at_level(logging.WARNING), patcher as mock_get:
        file_paths, file_sizes = map(list, zip(*meta_provider.expand_paths(paths, fs)))

    mock_get.assert_called_once_with(paths, fs, False)
    if isinstance(fs, LocalFileSystem):
        # No warning should be logged.
        assert len(caplog.text) == 0
    else:
        # Many files with different directories on cloud storage should log warning.
        assert "common parent directory" in caplog.text
    assert file_paths == paths
    expected_file_sizes = _get_file_sizes_bytes(paths, fs)
    assert file_sizes == expected_file_sizes

    # Many directories should not trigger error.
    if isinstance(fs, LocalFileSystem):
        dir_paths = [dir1, dir2] * num_dfs
        with caplog.at_level(logging.WARNING), patcher as mock_get:
            file_paths, file_sizes = map(
                list, zip(*meta_provider.expand_paths(dir_paths, fs))
            )
        assert len(file_paths) == len(paths) * num_dfs


if __name__ == "__main__":
    import sys

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