import io
import json
import warnings

from .core import url_to_fs
from .utils import merge_offset_ranges

# Parquet-Specific Utilities for fsspec
#
# Most of the functions defined in this module are NOT
# intended for public consumption. The only exception
# to this is `open_parquet_file`, which should be used
# place of `fs.open()` to open parquet-formatted files
# on remote file systems.


def open_parquet_file(
    path,
    fs=None,
    columns=None,
    row_groups=None,
    storage_options=None,
    engine="auto",
    max_gap=64_000,
    max_block=256_000_000,
    footer_sample_size=1_000_000,
    **kwargs,
):
    """
    Return a file-like object for a single Parquet file.

    The specified parquet `engine` will be used to parse the
    footer metadata, and determine the required byte ranges
    from the file. The target path will then be opened with
    the "parts" (`KnownPartsOfAFile`) caching strategy.

    Note that this method is intended for usage with remote
    file systems, and is unlikely to improve parquet-read
    performance on local file systems.

    Parameters
    ----------
    path: str
        Target file path.
    fs: AbstractFileSystem, optional
        Filesystem object to use for opening the file. If nothing is
        specified, an `AbstractFileSystem` object will be inferred.
    engine : str, default "auto"
        Parquet engine to use for metadata parsing. Allowed options
        include "fastparquet", "pyarrow", and "auto". The specified
        engine must be installed in the current environment. If
        "auto" is specified, and both engines are installed,
        "fastparquet" will take precedence over "pyarrow".
    columns: list, optional
        List of all column names that may be read from the file.
    row_groups : list, optional
        List of all row-group indices that may be read from the file.
    storage_options : dict, optional
        Used to generate an `AbstractFileSystem` object if `fs` was
        not specified.
    max_gap : int, optional
        Neighboring byte ranges will only be merged when their
        inter-range gap is <= `max_gap`. Default is 64KB.
    max_block : int, optional
        Neighboring byte ranges will only be merged when the size of
        the aggregated range is <= `max_block`. Default is 256MB.
    footer_sample_size : int, optional
        Number of bytes to read from the end of the path to look
        for the footer metadata. If the sampled bytes do not contain
        the footer, a second read request will be required, and
        performance will suffer. Default is 1MB.
    **kwargs :
        Optional key-word arguments to pass to `fs.open`
    """

    # Make sure we have an `AbstractFileSystem` object
    # to work with
    if fs is None:
        fs = url_to_fs(path, storage_options=(storage_options or {}))[0]

    # Fetch the known byte ranges needed to read
    # `columns` and/or `row_groups`
    data = _get_parquet_byte_ranges(
        [path],
        fs,
        columns=columns,
        row_groups=row_groups,
        engine=engine,
        max_gap=max_gap,
        max_block=max_block,
        footer_sample_size=footer_sample_size,
    )

    # Call self.open with "parts" caching
    options = kwargs.pop("cache_options", {}).copy()
    return fs.open(
        path,
        mode="rb",
        cache_type="parts",
        cache_options={**options, **{"data": data[path]}},
        **kwargs,
    )


def _get_parquet_byte_ranges(
    paths,
    fs,
    columns=None,
    row_groups=None,
    max_gap=64_000,
    max_block=256_000_000,
    footer_sample_size=1_000_000,
    engine="auto",
):
    """Get a dictionary of the known byte ranges needed
    to read a specific column/row-group selection from a
    Parquet dataset. Each value in the output dictionary
    is intended for use as the `data` argument for the
    `KnownPartsOfAFile` caching strategy of a single path.
    """

    # Set the engine
    engine = _set_engine(engine)

    # Get file sizes asynchronously
    file_sizes = fs.sizes(paths)

    # Populate global paths, starts, & ends
    result = {}
    data_paths = []
    data_starts = []
    data_ends = []
    add_header_magic = True
    if columns is None and row_groups is None:
        # We are NOT selecting specific columns or row-groups.
        #
        # We can avoid sampling the footers, and just transfer
        # all file data with cat_ranges
        for i, path in enumerate(paths):
            result[path] = {}
            for b in range(0, file_sizes[i], max_block):
                data_paths.append(path)
                data_starts.append(b)
                data_ends.append(min(b + max_block, file_sizes[i]))
        add_header_magic = False  # "Magic" should already be included
    else:
        # We ARE selecting specific columns or row-groups.
        #
        # Gather file footers.
        # We just take the last `footer_sample_size` bytes of each
        # file (or the entire file if it is smaller than that)
        footer_starts = []
        footer_ends = []
        for i, path in enumerate(paths):
            footer_ends.append(file_sizes[i])
            sample_size = max(0, file_sizes[i] - footer_sample_size)
            footer_starts.append(sample_size)