from __future__ import annotations

import dataclasses as dc
from typing import Any
from typing import ClassVar
from typing import Generic
from typing import Literal
from typing import TypeVar
from typing import Union

import pyarrow as pa
import pyarrow.csv as csv
import pyarrow.dataset as ds
import pyarrow.json as json
import pydantic.dataclasses as pydc

from .dataclass import pydc_config


ArrowReadOptions = Union[ds.ParquetReadOptions, csv.ReadOptions, json.ReadOptions]
ArrowParseOptions = Union[csv.ParseOptions, json.ParseOptions]
ArrowConvertOptions = csv.ConvertOptions
ArrowFSOptions = TypeVar(
    'ArrowFSOptions',
    ds.ParquetReadOptions, csv.ReadOptions, json.ReadOptions,
    csv.ParseOptions, json.ParseOptions, csv.ConvertOptions,
    ds.DirectoryPartitioning, ds.FilenamePartitioning, ds.HivePartitioning,
)


@dc.dataclass(frozen=True, repr=False)
class _ArrowFSOptionsMixin(Generic[ArrowFSOptions]):
    """Mixin class for serializable versions of Pyarrow's FS option classes.

    NB: we cannot add _pyarrow_fs_cls as a ClassVar attribute to the mixin since
    ClassVar's cannot contain type variables, see:
    https://peps.python.org/pep-0526/#class-and-instance-variable-annotations

    """  # noqa: E501

    def __repr__(self) -> str:
        # Custom __repr__ to display a Pyarrow schema in the same style as a dataclass
        field_reprs: list[str] = []
        for field in dc.fields(self):
            value = getattr(self, field.name)
            if isinstance(value, pa.Schema):
                value = str(value).replace('\n', ', ')
            field_reprs.append(f'{field.name}={value}')
        return f'{self.__class__.__name__}({", ".join(field_reprs)})'

    @property
    def as_pyarrow(self) -> ArrowFSOptions:
        """Return the equivalent Pyarrow object."""
        args = {}
        for field in dc.fields(self):
            if field.name != '_pyarrow_fs_cls':
                args[field.name] = getattr(self, field.name)
        return self._pyarrow_fs_cls(**args)  # type: ignore[attr-defined]


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class CSVReadOptions(_ArrowFSOptionsMixin[csv.ReadOptions]):
    """Read options for CSV files.
    Serializable version of [pa.csv.ReadOptions](https://arrow.apache.org/docs/python/generated/pyarrow.csv.ReadOptions).

    Args:
        use_threads: Whether to use multiple threads to accelerate reading.
        block_size: How much bytes to process at a time from the input stream.
        skip_rows: The number of rows to skip before the column names (if any) and the
            CSV data.
        skip_rows_after_names: The number of rows to skip after the column names.
        column_names: The column names of the target table. If empty, fall back on
            autogenerate_column_names.
        autogenerate_column_names: Whether to autogenerate column names if column_names
            is empty.
        encoding: The character encoding of the CSV data.

    """  # noqa: E501
    use_threads: bool = True
    block_size: int | None = None
    skip_rows: int = 0
    skip_rows_after_names: int = 0
    column_names: list[str] | None = None
    autogenerate_column_names: bool = False
    encoding: str = 'utf8'
    _pyarrow_fs_cls: ClassVar[type[csv.ReadOptions]] = csv.ReadOptions


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class CSVParseOptions(_ArrowFSOptionsMixin[csv.ParseOptions]):
    """Parse options for CSV files.
        Serializable version of [pa.csv.ParseOptions](https://arrow.apache.org/docs/python/generated/pyarrow.csv.ParseOptions).

    Args:
        delimiter: The character delimiting individual cells in the CSV data.
        quote_char: The character used optionally for quoting CSV values.
        double_quote: Whether two quotes in a quoted CSV value denote a single quote in
            the data.
        escape_char: The character used optionally for escaping special characters.
        newlines_in_values: Whether newline characters are allowed in CSV values.
            Setting this to True reduces the performance of multi-threaded CSV reading.
        ignore_empty_lines: Whether empty lines are ignored in CSV input.

    """  # noqa: E501
    delimiter: str = ','
    quote_char: str | bool = '"'
    double_quote: bool = True
    escape_char: str | bool = False
    newlines_in_values: bool = False
    ignore_empty_lines: bool = True
    _pyarrow_fs_cls: ClassVar[type[csv.ParseOptions]] = csv.ParseOptions


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class CSVConvertOptions(_ArrowFSOptionsMixin[csv.ConvertOptions]):
    """Conversion options for CSV files.
        Serializable version of [pa.csv.ConvertOptions](https://arrow.apache.org/docs/python/generated/pyarrow.csv.ConvertOptions).

    Args:
        check_utf8: Whether to check UTF8 validity of string columns.
        column_types: Explicitly map column names to column types.
        null_values: A sequence of strings that denote nulls in the data.
        decimal_point: The character used as decimal point in floating-point and decimal
            data.
        strings_can_be_null: Whether string / binary columns can have null values.
        quoted_strings_can_be_null: Whether quoted values can be null.
        auto_dict_encode: Whether to try to automatically dict-encode string / binary
            data.
        auto_dict_max_cardinality: The maximum dictionary cardinality for
            auto_dict_encode. This value is per chunk.
        timestamp_parsers: A sequence of strptime()-compatible format strings, tried in
            order when attempting to infer or convert timestamp values.

    """  # noqa: E501
    check_utf8: bool = True
    column_types: pa.Schema | None = None
    null_values: list[str] | None = None
    decimal_point: str = '.'
    strings_can_be_null: bool = False
    quoted_strings_can_be_null: bool = True
    auto_dict_encode: bool = False
    auto_dict_max_cardinality: int | None = None
    timestamp_parsers: list[str] | None = None
    _pyarrow_fs_cls: ClassVar[type[csv.ConvertOptions]] = csv.ConvertOptions


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class JSONReadOptions(_ArrowFSOptionsMixin[json.ReadOptions]):
    """Read options for JSON files.
        Serializable version of [pa.json.ReadOptions](https://arrow.apache.org/docs/python/generated/pyarrow.json.ReadOptions).

    Args:
        use_threads: Whether to use multiple threads to accelerate reading.
        block_size: How much bytes to process at a time from the input stream.

    """  # noqa: E501
    use_threads: bool = True
    block_size: int | None = None
    _pyarrow_fs_cls: ClassVar[type[json.ReadOptions]] = json.ReadOptions


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class JSONParseOptions(_ArrowFSOptionsMixin[json.ParseOptions]):
    """Parse options for JSON files.
        Serializable version of [pa.json.ParseOptions](https://arrow.apache.org/docs/python/generated/pyarrow.json.ParseOptions).

    Args:
        explicit_schema: Explicit schema (no type inference, ignores other fields).
        newlines_in_values: Whether objects may be printed across multiple lines
            (for example pretty printed). If false, input must end with an empty line.
        unexpected_field_behavior: How JSON fields outside of explicit_schema (if given)
            are treated.

    """  # noqa: E501
    explicit_schema: pa.Schema | None = None
    newlines_in_values: bool = False
    unexpected_field_behavior: Literal['ignore', 'error', 'infer'] = 'infer'
    _pyarrow_fs_cls: ClassVar[type[json.ParseOptions]] = json.ParseOptions


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class ParquetReadOptions(_ArrowFSOptionsMixin[ds.ParquetReadOptions]):
    """Read options for parquet files.
        Serializable version of [pa.ds.ParquetReadOptions](https://arrow.apache.org/docs/python/generated/pyarrow.dataset.ParquetReadOptions).

    Args:
        dictionary_columns: Names of columns which should be dictionary encoded as they
            are read in.
        coerce_int96_timestamp_unit: Cast timestamps that are stored in INT96 format to
            a particular resolution (e.g. ‘ms’). Setting to None is equivalent to ‘ns’
            and therefore INT96 timestamps will be inferred as timestamps in nanosecs.

    """  # noqa: E501
    dictionary_columns: list[str] | None = None
    coerce_int96_timestamp_unit: str | None = None
    _pyarrow_fs_cls: ClassVar[type[ds.ParquetReadOptions]] = ds.ParquetReadOptions


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class DirectoryPartitioning(_ArrowFSOptionsMixin[ds.DirectoryPartitioning]):
    """Directory partitioning options.
        Serializable version of [pa.ds.DirectoryPartitioning](https://arrow.apache.org/docs/python/generated/pyarrow.ds.DirectoryPartitioning).

    Args:
        schema: The schema that describes the partitions present in the file path.
        dictionaries: If the type of any field of schema is a dictionary type, the
            corresponding entry of dictionaries must be an array containing every value
            which may be taken by the corresponding column.
        segment_encoding: After splitting paths into segments, decode the segments.

    """  # noqa: E501
    schema: pa.Schema
    dictionaries: dict[str, list[Any]] | None = None
    segment_encoding: str = 'uri'
    _pyarrow_fs_cls: ClassVar[type[ds.DirectoryPartitioning]] = ds.DirectoryPartitioning


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class FilenamePartitioning(_ArrowFSOptionsMixin[ds.FilenamePartitioning]):
    """Filename partitioning options.
        Serializable version of [pa.ds.FilenamePartitioning](https://arrow.apache.org/docs/python/generated/pyarrow.ds.FilenamePartitioning).

    Args:
        schema: The schema that describes the partitions present in the file path.
        dictionaries: If the type of any field of schema is a dictionary type, the
            corresponding entry of dictionaries must be an array containing every value
            which may be taken by the corresponding column.
        segment_encoding: After splitting paths into segments, decode the segments.

    """  # noqa: E501
    schema: pa.Schema
    dictionaries: dict[str, list[Any]] | None = None
    segment_encoding: str = 'uri'
    _pyarrow_fs_cls: ClassVar[type[ds.FilenamePartitioning]] = ds.FilenamePartitioning


@pydc.dataclass(frozen=True, repr=False, config=pydc_config)
class HivePartitioning(_ArrowFSOptionsMixin[ds.HivePartitioning]):
    """Hive partitioning options.
        Serializable version of [pa.ds.HivePartitioning](https://arrow.apache.org/docs/python/generated/pyarrow.ds.HivePartitioning).

    Args:
        schema: The schema that describes the partitions present in the file path.
        dictionaries: If the type of any field of schema is a dictionary type, the
            corresponding entry of dictionaries must be an array containing every value
            which may be taken by the corresponding column.
        null_fallback: If any field is None then this fallback will be used as a label.
        segment_encoding: After splitting paths into segments, decode the segments.

    """  # noqa: E501
    schema: pa.Schema
    dictionaries: dict[str, list[Any]] | None = None
    null_fallback: str = '__HIVE_DEFAULT_PARTITION__'
    segment_encoding: str = 'uri'
    _pyarrow_fs_cls: ClassVar[type[ds.HivePartitioning]] = ds.HivePartitioning


ReadOptions = Union[ParquetReadOptions, CSVReadOptions, JSONReadOptions]
"""File read options for Parquet, CSV and JSON datasets."""

ParseOptions = Union[CSVParseOptions, JSONParseOptions]
"""File parse options for CSV and JSON datasets."""

ConvertOptions = CSVConvertOptions
"""Data conversion options for CSV datasets."""

Partitioning = Union[DirectoryPartitioning, FilenamePartitioning, HivePartitioning]
"""Dataset partitioning options for directory, filename and hive partitioning."""

ArrowPartitioning = Union[
    ds.DirectoryPartitioning, ds.FilenamePartitioning, ds.HivePartitioning,
]

FileFormat = Literal['parquet', 'csv', 'json', 'feather', 'orc']
"""Format of the files making up a dataset."""

ArrowFileFormat = Union[ds.ParquetFileFormat, ds.CsvFileFormat, ds.JsonFileFormat]

FSType = Literal['local', 'server', 's3', 'gcs', 'hdfs']
"""File system type.

- local: the user's local filesystem
- server: the server's filesystem
- s3: Amazon S3
- gcs: Google Cloud Storage
- hdfs: HDFS file system
"""
