"""File system datasets.

- FSDataset: Main class representing a file system dataset.
- FilteredFSDataset: Filtered version of an FSDataset instance.
"""
from __future__ import annotations

import dataclasses as dc
import logging
from pathlib import Path
from typing import Any
from typing import ClassVar
from typing import Literal

import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.flight as flight
import pydantic.dataclasses as pydc

from .base import DatasetBase
from .filter import FilterMixin
from tempora.utils.dataclass import pydc_config
from tempora.utils.dataset import cleanup
from tempora.utils.exceptions import parse_arrow_exception
from tempora.utils.exceptions import TemporaServerError
from tempora.utils.fs import get_fs_type
from tempora.utils.fs import load_fs_dataset
from tempora.utils.fs_options import ConvertOptions
from tempora.utils.fs_options import FileFormat
from tempora.utils.fs_options import FSType
from tempora.utils.fs_options import ParseOptions
from tempora.utils.fs_options import Partitioning
from tempora.utils.fs_options import ReadOptions
from tempora.utils.server import get_connection
from tempora.utils.server import SessionParams


log = logging.getLogger(__name__)


@dc.dataclass(frozen=True)
class _FSDatasetArgs:
    source: str | Path | list[str | Path]
    file_format: FileFormat | None = None
    file_schema: pa.Schema | None = None
    columns: list[str] | None = None
    read_options: ReadOptions | None = None
    parse_options: ParseOptions | None = None
    convert_options: ConvertOptions | None = None
    schema_inf_depth: float | None = None  # in MiB
    partitioning: Literal['hive'] | Partitioning | None = None
    fs_config: dict[str, Any] | None = None
    dataset_name: str | None = None


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class FSDataset(DatasetBase, _FSDatasetArgs):
    """Class representing a file system backed dataset.

    Args:
        source: The filepath(s) that make up the dataset. For non-local filesystems
            prefix each path with the appropriate URI (i.e. s3://, gcs://, hdfs://).
            For a dataset located on the tempora server use the path prefix "server:".
        file_format: Format of the dataset files. If not specified the format will be
            inferred from the file extensions during dataset loading.
        file_schema: Schema for the dataset as a Pyarrow Schema object.
        columns: The specific columns to load from the dataset.
        read_options: Options to customize file reading
        parse_options: Options to customize file parsing
        convert_options: Options to specify data conversions.
        schema_inf_depth: The amount of data (in MiBs) which is used to infer the schema
            of the dataset. Defaults to 8 MiB.
        partitioning: The dataset partitioning structure.
        fs_config: Custom arguments to pass to Pyarrow's
            [LocalFileSystem](https://arrow.apache.org/docs/python/generated/pyarrow.fs.LocalFileSystem.html),
            [GcsFileSystem](https://arrow.apache.org/docs/python/generated/pyarrow.fs.GcsFileSystem.html)
            or [HadoopFileSystem](https://arrow.apache.org/docs/python/generated/pyarrow.fs.HadoopFileSystem.html).
            For S3 specify the parameters in the form expected by
            [boto3.session.Session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html).
        dataset_name: Optional name for the dataset.
    """  # noqa: E501
    filtered_dataset_cls: ClassVar[type[FilteredFSDataset]]

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)

    @property
    def fs_type(self) -> FSType:
        """The file system type for the dataset - one of: server, s3, gcs, or hdfs."""
        path = str(self.source[0] if isinstance(self.source, list) else self.source)
        return get_fs_type(path)

    def _init_dataset(self, _op_dict: dict[str, Any] | None = None) -> None:
        if _op_dict or self.fs_type != 'local':
            # Don't upload source data for filter() / join() ops (_op_dict provided)
            super()._init_dataset(_op_dict=_op_dict)
        else:
            self._init_local_dataset()

    def _init_local_dataset(self) -> None:
        """Upload the files for a local dataset to the server."""
        try:
            # Create the flight descriptor
            desc = flight.FlightDescriptor.for_command(self._as_json.encode())

            arrow_data = load_fs_dataset(
                source=self.source,
                file_format=self.file_format,
                file_schema=self.file_schema,
                read_options=self.read_options,
                parse_options=self.parse_options,
                convert_options=self.convert_options,
                schema_inf_depth=self.schema_inf_depth,
                partitioning=self.partitioning,
                fs_config=None,
                columns=self.columns,
                time_column=self.time_column,
            )
            if isinstance(arrow_data, ds.Dataset):
                reader = pa.RecordBatchReader.from_batches(
                    schema=arrow_data.schema, batches=arrow_data.to_batches(),
                )
            elif isinstance(arrow_data, (ds.Scanner, pa.Table)):
                reader = arrow_data.to_reader()
            else:
                reader = arrow_data

            # Send the data to the server
            client, auth_token = get_connection()
            call_options = flight.FlightCallOptions(headers=[auth_token])
            writer, _ = client.do_put(desc, reader.schema, call_options)
            for batch in reader:
                try:
                    writer.write_batch(batch)
                except KeyboardInterrupt:
                    # Currently it's not possible with FlightSreamWriter to signal to
                    # the server that the DoPut operation should be cancelled, so we
                    # need to handle this manually by calling the cleanup function, see:
                    # https://github.com/apache/arrow/issues/34431
                    writer.close()
                    session_params = SessionParams(client=client, auth_token=auth_token)
                    cleanup(json_key=self._as_json, session_params=session_params)
                    break
            else:
                writer.close()

        except flight.FlightError as e:
            msg = parse_arrow_exception(e)
            log.error(msg)
            raise TemporaServerError(msg) from None

        return None


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class FilteredFSDataset(FilterMixin, FSDataset):
    """Class representing a filtered file system backed dataset."""

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)


FSDataset.filtered_dataset_cls = FilteredFSDataset
