"""File reading generators for common file types."""

import abc
from typing import Any, Dict, Union

import pandas as pd
import pyarrow
import pyarrow.dataset
import pyarrow.parquet as pq
from astropy.io import ascii as ascii_reader
from astropy.table import Table
from hipscat.io import FilePointer, file_io

# pylint: disable=too-few-public-methods,too-many-arguments


def get_file_reader(
    file_format,
    chunksize=500_000,
    schema_file=None,
    column_names=None,
    skip_column_names=None,
    type_map=None,
    **kwargs,
):
    """Get a generator file reader for common file types

    Args:
        file_format (str): specifier for the file type and extension.
            Currently supported formats include:

            - `csv`, comma separated values. may also be tab- or pipe-delimited
              includes `.csv.gz` and other compressed csv files
            - `fits`, flexible image transport system. often used for astropy tables.
            - `parquet`, compressed columnar data format
            - `ecsv`, astropy's enhanced CSV
            - `indexed_csv`, "index" style reader, that accepts a file with a list
              of csv files that are appended in-memory
            - `indexed_parquet`, "index" style reader, that accepts a file with a list
              of parquet files that are appended in-memory

        chunksize (int): number of rows to read in a single iteration.
            for single-file readers, large files are split into batches based on this value.
            for index-style readers, we read files until we reach this chunksize and
            create a single batch in-memory.
        schema_file (str): path to a parquet schema file. if provided, header names
            and column types will be pulled from the parquet schema metadata.
        column_names (list[str]): for CSV files, the names of columns if no header
            is available. for fits files, a list of columns to *keep*.
        skip_column_names (list[str]): for fits files, a list of columns to remove.
        type_map (dict): for CSV files, the data types to use for columns
        kwargs: additional keyword arguments to pass to the underlying file reader.
    """
    if file_format == "csv":
        return CsvReader(
            chunksize=chunksize,
            schema_file=schema_file,
            column_names=column_names,
            type_map=type_map,
            **kwargs,
        )
    if file_format == "ecsv":
        return AstropyEcsvReader(**kwargs)
    if file_format == "fits":
        return FitsReader(
            chunksize=chunksize,
            column_names=column_names,
            skip_column_names=skip_column_names,
            **kwargs,
        )
    if file_format == "parquet":
        return ParquetReader(chunksize=chunksize, **kwargs)
    if file_format == "indexed_csv":
        return IndexedCsvReader(
            chunksize=chunksize,
            schema_file=schema_file,
            column_names=column_names,
            type_map=type_map,
            **kwargs,
        )
    if file_format == "indexed_parquet":
        return IndexedParquetReader(chunksize=chunksize, **kwargs)
    raise NotImplementedError(f"File Format: {file_format} not supported")


class InputReader(abc.ABC):
    """Base class for chunking file readers."""

    @abc.abstractmethod
    def read(self, input_file, read_columns=None):
        """Read the input file, or chunk of the input file.

        Args:
            input_file(str): path to the input file.
            read_columns(List[str]): subset of columns to read.
                if None, all columns are read
        Yields:
            DataFrame containing chunk of file info.
        """

    def provenance_info(self) -> dict:
        """Create dictionary of parameters for provenance tracking.

        If any `storage_options` have been provided as kwargs, we will replace the
        value with ``REDACTED`` for the purpose of writing to provenance info, as it
        may contain user names or API keys.

        Returns:
            dictionary with all argument_name -> argument_value as key -> value pairs.
        """
        all_args = vars(self)
        if "kwargs" in all_args and "storage_options" in all_args["kwargs"]:
            all_args["kwargs"]["storage_options"] = "REDACTED"
        return {"input_reader_type": type(self).__name__, **vars(self)}

    def regular_file_exists(self, input_file, storage_options: Union[Dict[Any, Any], None] = None, **_kwargs):
        """Check that the `input_file` points to a single regular file

        Raises:
            FileNotFoundError: if nothing exists at path, or directory found.
        """
        if not file_io.does_file_or_directory_exist(input_file, storage_options=storage_options):
            raise FileNotFoundError(f"File not found at path: {input_file}")
        if not file_io.is_regular_file(input_file, storage_options=storage_options):
            raise FileNotFoundError(f"Directory found at path - requires regular file: {input_file}")

    def read_index_file(self, input_file, storage_options: Union[Dict[Any, Any], None] = None, **kwargs):
        """Read an "indexed" file.

        This should contain a list of paths to files to be read and batched.

        Raises:
            FileNotFoundError: if nothing exists at path, or directory found.
        """
        self.regular_file_exists(input_file, **kwargs)
        file_names = file_io.load_text_file(input_file, storage_options=storage_options)
        file_names = [f.strip() for f in file_names]
        file_names = [f for f in file_names if f]
        return file_names


class CsvReader(InputReader):
    """CSV reader for the most common CSV reading arguments.

    This uses `pandas.read_csv`, and you can find more information on
    additional arguments in the pandas documentation:
    https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html

    Attributes:
        chunksize (int): number of rows to read in a single iteration.
        header (int, list of int, None, default 'infer'): rows to
            use as the header with column names
        schema_file (str): path to a parquet schema file. if provided, header names
            and column types will be pulled from the parquet schema metadata.
        column_names (list[str]): the names of columns if no header is available
        type_map (dict): the data types to use for columns
        parquet_kwargs (dict): additional keyword arguments to use when
            reading the parquet schema metadata, passed to pandas.read_parquet.
            See https://pandas.pydata.org/docs/reference/api/pandas.read_parquet.html
        kwargs (dict): additional keyword arguments to use when reading
            the CSV files with pandas.read_csv.
            See https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html
    """

    def __init__(
        self,
        chunksize=500_000,
        header="infer",
        schema_file=None,
        column_names=None,
        type_map=None,
        parquet_kwargs=None,
        **kwargs,
    ):
        self.chunksize = chunksize
        self.header = header
        self.schema_file = schema_file
        self.column_names = column_names
        self.type_map = type_map
        self.parquet_kwargs = parquet_kwargs
        self.kwargs = kwargs

        schema_parquet = None
        if self.schema_file:
            if self.parquet_kwargs is None:
                self.parquet_kwargs = {}
            schema_parquet = file_io.load_parquet_to_pandas(
                FilePointer(self.schema_file),
                **self.parquet_kwargs,
            )

        if self.column_names:
            self.kwargs["names"] = self.column_names
        elif not self.header and schema_parquet is not None:
            self.kwargs["names"] = schema_parquet.columns

        if self.type_map:
            self.kwargs["dtype"] = self.type_map
        elif schema_parquet is not None:
            self.kwargs["dtype"] = schema_parquet.dtypes.to_dict()

    def read(self, input_file, read_columns=None):
        self.regular_file_exists(input_file, **self.kwargs)

        if read_columns:
            self.kwargs["usecols"] = read_columns

        with file_io.load_csv_to_pandas(
            FilePointer(input_file),
            chunksize=self.chunksize,
            header=self.header,
            **self.kwargs,
        ) as reader:
            yield from reader


class IndexedCsvReader(CsvReader):
    """Reads an index file, containing paths to CSV files to be read and batched

    See CsvReader for additional configuration for reading CSV files.
    """

    def read(self, input_file, read_columns=None):
        file_names = self.read_index_file(input_file=input_file, **self.kwargs)

        batch_size = 0
        batch_frames = []
        for file in file_names:
            for single_frame in super().read(file, read_columns=read_columns):
                if batch_size + len(single_frame) >= self.chunksize:
                    # We've hit our chunksize, send the batch off to the task.
                    if len(batch_frames) == 0:
                        yield single_frame
                        batch_size = 0
                    else:
                        yield pd.concat(batch_frames, ignore_index=True)
                        batch_frames = []
                        batch_frames.append(single_frame)
                        batch_size = len(single_frame)
                else:
                    batch_frames.append(single_frame)
                    batch_size += len(single_frame)

        if len(batch_frames) > 0:
            yield pd.concat(batch_frames, ignore_index=True)


class AstropyEcsvReader(InputReader):
    """Reads astropy ascii .ecsv files.

    Note that this is NOT a chunked reader. Use caution when reading
    large ECSV files with this reader.

    Attributes:
        kwargs: keyword arguments passed to astropy ascii reader.
            See https://docs.astropy.org/en/stable/api/astropy.io.ascii.read.html#astropy.io.ascii.read
    """

    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def read(self, input_file, read_columns=None):
        self.regular_file_exists(input_file, **self.kwargs)
        if read_columns:
            self.kwargs["include_names"] = read_columns

        astropy_table = ascii_reader.read(input_file, format="ecsv", **self.kwargs)
        yield astropy_table.to_pandas()


class FitsReader(InputReader):
    """Chunked FITS file reader.

    There are two column-level arguments for reading fits files:
    `column_names` and `skip_column_names`.

        - If neither is provided, we will read and process all columns in the fits file.
        - If `column_names` is given, we will use *only* those names, and
          `skip_column_names` will be ignored.
        - If `skip_column_names` is provided, we will remove those columns from processing stages.

    NB: Uses astropy table memmap to avoid reading the entire file into memory.
    See: https://docs.astropy.org/en/stable/io/fits/index.html#working-with-large-files


    Attributes:
        chunksize (int): number of rows of the file to process at once.
            For large files, this can prevent loading the entire file
            into memory at once.
        column_names (list[str]): list of column names to keep. only use
            one of `column_names` or `skip_column_names`
        skip_column_names (list[str]): list of column names to skip. only use
            one of `column_names` or `skip_column_names`
        kwargs: keyword arguments passed along to astropy.Table.read.
            See https://docs.astropy.org/en/stable/api/astropy.table.Table.html#astropy.table.Table.read
    """

    def __init__(self, chunksize=500_000, column_names=None, skip_column_names=None, **kwargs):
        self.chunksize = chunksize
        self.column_names = column_names
        self.skip_column_names = skip_column_names
        self.kwargs = kwargs

    def read(self, input_file, read_columns=None):
        self.regular_file_exists(input_file, **self.kwargs)
        table = Table.read(input_file, memmap=True, **self.kwargs)
        if read_columns:
            table.keep_columns(read_columns)
        elif self.column_names:
            table.keep_columns(self.column_names)
        elif self.skip_column_names:
            table.remove_columns(self.skip_column_names)

        total_rows = len(table)
        read_rows = 0

        while read_rows < total_rows:
            df_chunk = table[read_rows : read_rows + self.chunksize].to_pandas()
            for column in df_chunk.columns:
                if (
                    df_chunk[column].dtype == object
                    and df_chunk[column].apply(lambda x: isinstance(x, bytes)).any()
                ):
                    df_chunk[column] = df_chunk[column].apply(
                        lambda x: x.decode("utf-8") if isinstance(x, bytes) else x
                    )

            yield df_chunk

            read_rows += self.chunksize


class ParquetReader(InputReader):
    """Parquet reader for the most common Parquet reading arguments.

    Attributes:
        chunksize (int): number of rows of the file to process at once.
            For large files, this can prevent loading the entire file
            into memory at once.
        column_names (list[str] or None): Names of columns to use from the input dataset.
            If None, use all columns.
        kwargs: arguments to pass along to pyarrow.parquet.ParquetFile.
            See https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetFile.html
    """

    def __init__(self, chunksize=500_000, column_names=None, **kwargs):
        self.chunksize = chunksize
        self.column_names = column_names
        self.kwargs = kwargs

    def read(self, input_file, read_columns=None):
        self.regular_file_exists(input_file, **self.kwargs)
        columns = read_columns or self.column_names
        parquet_file = pq.ParquetFile(input_file, **self.kwargs)
        for smaller_table in parquet_file.iter_batches(
            batch_size=self.chunksize, columns=columns, use_pandas_metadata=True
        ):
            yield smaller_table.to_pandas()


class IndexedParquetReader(InputReader):
    """Reads an index file, containing paths to parquet files to be read and batched

    Attributes:
        chunksize (int): maximum number of rows to process at once.
            Large files will be processed in chunks. Small files will be concatenated.
            Also passed to pyarrow.dataset.Dataset.to_batches as `batch_size`.
        batch_readahead (int): number of batches to read ahead.
            Passed to pyarrow.dataset.Dataset.to_batches.
        fragment_readahead (int): number of fragments to read ahead.
            Passed to pyarrow.dataset.Dataset.to_batches.
        use_threads (bool): whether to use multiple threads for reading.
            Passed to pyarrow.dataset.Dataset.to_batches.
        column_names (list[str] or None): Names of columns to use from the input dataset.
            If None, use all columns.
        kwargs: additional arguments to pass along to InputReader.read_index_file.
    """

    def __init__(
        self,
        chunksize=500_000,
        batch_readahead=16,
        fragment_readahead=4,
        use_threads=True,
        column_names=None,
        **kwargs,
    ):
        self.chunksize = chunksize
        self.batch_readahead = batch_readahead
        self.fragment_readahead = fragment_readahead
        self.use_threads = use_threads
        self.column_names = column_names
        self.kwargs = kwargs

    def read(self, input_file, read_columns=None):
        columns = read_columns or self.column_names
        file_names = self.read_index_file(input_file=input_file, **self.kwargs)
        (_, input_dataset) = file_io.read_parquet_dataset(file_names, **self.kwargs)

        batches, nrows = [], 0
        for batch in input_dataset.to_batches(
            batch_size=self.chunksize,
            batch_readahead=self.batch_readahead,
            fragment_readahead=self.fragment_readahead,
            use_threads=self.use_threads,
            columns=columns,
        ):
            if nrows + batch.num_rows > self.chunksize:
                # We've hit the chunksize so load to a DataFrame and yield.
                # There should always be at least one batch in here since batch_size == self.chunksize above.
                yield pyarrow.Table.from_batches(batches).to_pandas()
                batches, nrows = [], 0

            batches.append(batch)
            nrows += batch.num_rows

        if len(batches) > 0:
            yield pyarrow.Table.from_batches(batches).to_pandas()
