"""Abstract base class for datasets, together with related utility classes.

- DatasetBase: Abstract base class for datasets.
- Pivot: Class defining the pivot operation for datasets in EAV / wide format
- JoinData: Class defining the various options for a dataset join operation.
"""
from __future__ import annotations

import dataclasses as dc
import json
import logging
import os
from abc import ABC
from functools import cached_property
from typing import Any
from typing import ClassVar
from typing import Literal
from typing import NoReturn
from typing import overload
from weakref import finalize

import numpy as np  # noqa: F401
import pandas as pd
import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.flight as flight
import pydantic.dataclasses as pydc

from .protocols import IsDataset
from .protocols import IsFilteredDataset
from tempora.utils.dataclass import get_args_dict
from tempora.utils.dataclass import get_args_tuple
from tempora.utils.dataclass import pydc_config
from tempora.utils.dataset import arrow_json_encoder
from tempora.utils.dataset import AsofDirection
from tempora.utils.dataset import cleanup
from tempora.utils.exceptions import parse_arrow_exception
from tempora.utils.exceptions import TemporaServerError
from tempora.utils.server import get_connection
from tempora.utils.server import SessionParams


log = logging.getLogger(__name__)


@pydc.dataclass(config=pydc_config)
class Pivot:
    """Pivot operation for datasets in EAV / long / melted format.

    Args:
        on: The name of the column whose values will form the names of the new columns
            that will be created during the dataset pivot operation.
        using: The name of the column that will provide the value for each of the newly
            created columns.
        agg_function: The SQL aggregate function to use for combining the grouped values
            from the pivot "on" column. To specify different aggregate functions for
            different columns, pass a dictionary mapping the column name to the desired
            function. NB: the dictionary must contain a key for every newly created
            column by the pivot operation.
        in_values: Filter the values of the on column such that only the specified
            values will be used to form new columns during the pivot operation.
        dtype: The dtype of the newly created columns, e.g. np.dtype('float64'). If set
            to None then the columns will all have the same dtype as that of the "using"
            column. To specify different dtypes for different columns, pass a dictionary
            mapping the column name to the desired dtype. Any columns without a key in
            the dictionary will default to the dtype of the "using" column.
        errors: How to handle errors during dtype conversion. If 'raise' any values
            that cannot be cast to the desired dtype will raise an exception. If
            'coerce' then any values that cannot be cast will be set to NaN.
    """
    on: str
    using: str
    agg_function: str | dict[str, str]
    in_values: list[str] | None = None
    dtype: str | np.dtype | dict[str, str | np.dtype] | None = None
    errors: Literal['raise', 'coerce'] = 'coerce'

    def __post_init__(self) -> None:
        if self.in_values:
            if isinstance(self.agg_function, dict):
                # Every "in_values" column should also be specified in the agg func dict
                for column in self.in_values:
                    if column not in self.agg_function:
                        raise ValueError(
                            f'column "{column}" not specified in agg_function dict',
                        )
                # Every agg function column should be present in the "in_values" list
                for column in self.agg_function:
                    if column not in self.in_values:
                        raise ValueError(
                            f'agg_function key "{column}" not present in "in_values"',
                        )

            if isinstance(self.dtype, dict):
                # dtype columns should be present in the "in_values" list (if passed)
                for column in self.dtype:
                    if column not in self.in_values:
                        raise ValueError(
                            f'dtype column "{column}" is not in the "in_values" list',
                        )


@dc.dataclass
class JoinData:
    """Class representing a dataset join operation.

    Args:
        dataset: The dataset to join to the parent dataset.
        join_condition: To join the two datasets when the names of the join columns
            are the same in both datasets, pass a list of the column names to join on.
            To join using an "ON" style SQL join condition, pass a string containing
            the specific join condition (omitting the ON keyword).
            NB: for ON style joins, any column names in the join condition that are
            the same in both datasets must be fully qualified using the dataset names.
        asof_join: Use an ASOF style JOIN on the time columns from each dataset. Both
            datasets must be temporal to use this option (i.e. specify a time column).
        direction: Direction of the ASOF JOIN. If 'forward', each row from the parent
            dataset is joined to the closest row from the join dataset whose time
            value is >= to the parent time value. If 'backward', the comparison is
            reversed, i.e. each parent row is joined to the closest row from the
            join dataset whose time value is <= to the parent time value.
        allow_exact_matches: If False, disallow exact matches of the time values during
            the ASOF JOIN, i.e. use < and > in place of <= and >= in the comparison op.
    """
    dataset: DatasetBase
    join_condition: str | list[str]
    asof_join: bool = False
    direction: AsofDirection = 'forward'
    allow_exact_matches: bool = True


@dc.dataclass(frozen=True)
class _DBDatasetArgs:
    database: str
    table: str
    db_schema: str | None = None


@dc.dataclass(frozen=True)
class DatasetBase(ABC):
    """Abstract base class for datasets.

    Args:
        time_column: The name of the column that contains the time (or sequence) values
            that are used to order the rows in the dataset. The time values may be
            timestamps, times or dates, or more generally any numeric value that can be
            used to order the rows. If set to None, the dataset is assumed to be a
            non-temporal (i.e. static) dataset.
        entity_keys: The column name(s) that serve as the primary key(s) for identifying
            the rows associated with distinct entities in the dataset. For example, for
            a medical dataset the entity_keys might be a patient id, for an industrial
            dataset the entity_keys column might be a machine or equipment id.
        pivot: Optional pivot settings for a dataset in EAV / wide format.
        materialize: Whether to materialize the dataset on the server.
        targets: Whether the dataset contains target data for training batches.
    """
    filtered_dataset_cls: ClassVar[type[IsFilteredDataset]]
    time_column: str | None = None
    entity_keys: list[str] | None = None
    pivot: Pivot | None = None
    materialize: bool = False
    targets: bool = False
    _join_data: list[JoinData] | None = dc.field(default=None, repr=False)
    _op_dict: dc.InitVar[dict[str, Any] | None] = None

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

    def _finalize(self) -> None:
        """Set up the finalizer to initiate cleanup on the server.

        We implement this in an explicit method (vs __post_init__) so that we can mock
        it during unit testing.
        """
        # NB: we must pass an actual json dataset string and not a reference to
        # "self._as_json" so that self can be garbage collected, see:
        # https://docs.python.org/3/library/weakref.html#weakref.finalize
        json_dataset = self._as_json
        client, auth_token = get_connection()
        session_params = SessionParams(client=client, auth_token=auth_token)
        finalize(self, cleanup, json_key=json_dataset, session_params=session_params)
        return None

    def __hash__(self) -> int:
        # Subclasses must pass "eq=False" in the dc.dataclass decorator so that they
        # inherit this __hash__ method. From the dataclass docs:
        # "If eq is false, __hash__() will be left untouched meaning the __hash__()
        # method of the superclass will be used."
        return hash(get_args_tuple(self))

    def __eq__(self, other: object) -> bool:
        return (
            isinstance(other, type(self))
            and get_args_tuple(self) == get_args_tuple(other)
        )

    def _init_dataset(self, _op_dict: dict[str, Any] | None = None) -> None:
        """Initialize an instance of the dataset on the server.

        Args:
            _op_dict: Optional dictionary containing the details of an operation to
                be performed on the dataset (filter, join etc.)
        """
        if _op_dict:
            args_dict = get_args_dict(self)
            args_dict['op_dict'] = _op_dict
            json_dataset = json.dumps(args_dict, default=arrow_json_encoder)
        else:
            json_dataset = self._as_json

        action = flight.Action('init_dataset', json_dataset.encode())
        client, auth_token = get_connection()
        call_options = flight.FlightCallOptions(headers=[auth_token])
        result_iter = client.do_action(action, call_options)
        try:
            # Server responds with an empty result iter upon success of "init_dataset"
            next(result_iter)
        except flight.FlightError as e:
            msg = parse_arrow_exception(e)
            log.error(msg)
            raise TemporaServerError(msg) from None
        except StopIteration:
            # Empty response from server - all good
            ...
        return None

    @cached_property
    def _as_json(self) -> str:
        return json.dumps(get_args_dict(self), default=arrow_json_encoder)

    @property
    def _flight_reader(self) -> flight.FlightStreamReader:
        try:
            client, auth_token = get_connection()
            call_options = flight.FlightCallOptions(headers=[auth_token])
            ticket = flight.Ticket(self._as_json.encode())
            return client.do_get(ticket, call_options)
        except flight.FlightError as e:
            msg = f'Unable to retrieve data from the server: {parse_arrow_exception(e)}'
            log.error(msg)
            raise TemporaServerError(msg) from None

    @property
    def schema(self) -> pa.Schema:
        """The dataset schema."""
        action = flight.Action('get_dataset_schema', self._as_json.encode())
        try:
            client, auth_token = get_connection()
            call_options = flight.FlightCallOptions(headers=[auth_token])
            schema_bytes = next(client.do_action(action, call_options))
            return pa.ipc.read_schema(pa.py_buffer(schema_bytes.body))
        except flight.FlightError as e:
            msg = f'Problem encountered getting schema: {parse_arrow_exception(e)}'
            raise TemporaServerError(msg) from None
        except StopIteration:
            raise TemporaServerError('No response from server...') from None

    @property
    def num_rows(self) -> int:
        """The number of rows in the dataset."""
        action = flight.Action('count_rows', self._as_json.encode())
        try:
            client, auth_token = get_connection()
            call_options = flight.FlightCallOptions(headers=[auth_token])
            result = next(client.do_action(action, call_options))
            return int.from_bytes(result.body.to_pybytes(), 'little')
        except flight.FlightError as e:
            msg = f'Problem encountered getting row count: {parse_arrow_exception(e)}'
            raise TemporaServerError(msg) from None
        except StopIteration:
            raise TemporaServerError('No response from server...') from None

    def drop(self) -> None | NoReturn:
        """Drop the dataset from the server."""
        action = flight.Action('drop_dataset', self._as_json.encode())
        try:
            client, auth_token = get_connection()
            call_options = flight.FlightCallOptions(headers=[auth_token])
            next(client.do_action(action, call_options))
        except flight.FlightError as e:
            msg = parse_arrow_exception(e)
            log.error(msg)
            raise TemporaServerError(msg) from None
        except StopIteration:
            # Empty response from server - all good
            ...
        return None

    @overload
    def head(self, n: int = 10, *, as_arrow: Literal[False] = ...) -> pd.DataFrame: ...

    @overload
    def head(self, n: int = 10, *, as_arrow: Literal[True]) -> pa.Table: ...

    def head(self, n: int = 10, *, as_arrow: bool = False) -> pd.DataFrame | pa.Table:
        """Return the first n rows of the dataset as a dataframe.

        Args:
            n: Number of rows to return
            as_arrow: Return a Pyarrow Table if True, otherwise a dataframe.

        Returns:
            Pandas dataframe (or Pyarrow Table).
        """
        batches: list[pa.RecordBatch] = []
        num_rows_left = n

        flight_reader = self._flight_reader
        for chunk in flight_reader:
            batch = chunk.data
            if batch.num_rows > num_rows_left:
                indices = list(range(num_rows_left))
                batch = batch.take(indices)
            batches.append(batch)
            num_rows_left -= batch.num_rows
            if num_rows_left <= 0:
                break

        flight_reader.cancel()
        table = pa.Table.from_batches(batches)
        if as_arrow:
            return table
        else:
            return table.to_pandas()

    def to_pandas(self, use_nullable_dtypes: bool = False) -> pd.DataFrame:
        """Return the dataset as a Pandas dataframe.

        Args:
            use_nullable_dtypes: If True, use dtypes that use pd.NA as a missing value
                indicator for the resulting DataFrame.

        Returns:
            Dataset as a dataframe.
        """
        dtype_mapping = {
            pa.int8(): pd.Int8Dtype(),
            pa.int16(): pd.Int16Dtype(),
            pa.int32(): pd.Int32Dtype(),
            pa.int64(): pd.Int64Dtype(),
            pa.uint8(): pd.UInt8Dtype(),
            pa.uint16(): pd.UInt16Dtype(),
            pa.uint32(): pd.UInt32Dtype(),
            pa.uint64(): pd.UInt64Dtype(),
            pa.bool_(): pd.BooleanDtype(),
            pa.float32(): pd.Float32Dtype(),
            pa.float64(): pd.Float64Dtype(),
            pa.string(): pd.StringDtype(),
        }
        if use_nullable_dtypes:
            return self._flight_reader.read_pandas(types_mapper=dtype_mapping.get)
        else:
            return self._flight_reader.read_pandas()

    def df(self, use_nullable_dtypes: bool = False) -> pd.DataFrame:
        """Alias for [to_pandas][tempora.datasets.base.DatasetBase.to_pandas]."""
        return self.to_pandas(use_nullable_dtypes=use_nullable_dtypes)

    def to_numpy(self) -> np.ndarray:
        """Return the dataset as a numpy array (if possible)."""
        return self.to_pandas().to_numpy()

    def np(self) -> np.ndarray:
        """Alias for [to_numpy][tempora.datasets.base.DatasetBase.to_numpy]."""
        return self.to_numpy()

    @overload
    def to_arrow(self, *, materialize: Literal[True] = ...) -> pa.Table: ...

    @overload
    def to_arrow(self, *, materialize: Literal[False]) -> pa.RecordBatchReader: ...

    def to_arrow(self, materialize: bool = True) -> pa.Table | pa.RecordBatchReader:
        """Return the dataset as a Pyarrow Table or RecordBatchReader.

        Args:
            materialize: If True return a Table, otherwise a RecordBatchReader.

        Returns:
            Pyarrow Table or RecordBatchReader.
        """
        if materialize:
            return self._flight_reader.read_all()
        else:
            return self._flight_reader.to_reader()

    def write_dataset(self, *args: Any, **kwargs: Any) -> None:
        """Write a dataset to a given format and partitioning.
        See [pa.ds.write_dataset](https://arrow.apache.org/docs/python/generated/pyarrow.dataset.write_dataset.html).

        """  # noqa: E501
        # Disable unnecessary Arrow array alignment checking that can lead to warnings:
        # https://github.com/apache/arrow/issues/36301
        os.environ['ACERO_ALIGNMENT_HANDLING'] = 'ignore'
        ds.write_dataset(self.to_arrow(materialize=False), *args, **kwargs)

    @classmethod
    def _add_join_data(
        cls,
        left_dataset: IsDataset,
        right_dataset: IsDataset,
        join_condition: str | list[str],
        asof_join: bool = False,
        direction: AsofDirection = 'forward',
        allow_exact_matches: bool = True,
    ) -> IsDataset:
        args: dict[str, Any] = {}
        for field in dc.fields(left_dataset):  # noqa
            args[field.name] = getattr(left_dataset, field.name)

        join_data = JoinData(
            dataset=right_dataset,  # type: ignore[arg-type]
            join_condition=join_condition,
            asof_join=asof_join,
            direction=direction,
            allow_exact_matches=allow_exact_matches,
        )
        if left_dataset._join_data:
            args['_join_data'] = left_dataset._join_data + [join_data]
        else:
            args['_join_data'] = [join_data]
        args['_op_dict'] = {'op': 'join', 'parent': get_args_dict(left_dataset)}
        return cls(**args)  # noqa

    def join(
        self, dataset: IsDataset,
        join_condition: str | list[str], *,
        asof_join: bool = False,
        direction: AsofDirection = 'forward',
        allow_exact_matches: bool = True,
    ) -> IsDataset:
        """Join with another dataset.

        Args:
            dataset: The dataset to join to the parent dataset.
            join_condition: To join the two datasets when the names of the join columns
                are the same in both datasets, pass a list of the column names to join
                on. To join using an "ON" style SQL join condition, pass a string
                containing the specific join condition (omitting the ON keyword).
                NB: for ON style joins, any column names in the join condition that are
                the same in both datasets must be fully qualified using the dataset
                names.
            asof_join: Use an ASOF style JOIN on the time columns from each dataset.
                Both datasets must be temporal to use this option (i.e. specify a time
                column).
            direction: Direction of the ASOF JOIN. If 'forward', each row from the parent
                dataset is joined to the closest row from the join dataset whose time
                value is >= to the parent time value. If 'backward', the comparison is
                reversed, i.e. each parent row is joined to the closest row from the
                join dataset whose time value is <= to the parent time value.
            allow_exact_matches: If False, disallow exact matches of the time values
                during the ASOF JOIN, i.e. use < and > in place of <= and >= in the
                comparison op.

        Returns:
            The joined dataset object.

        """  # noqa: E501

        return self._add_join_data(
            left_dataset=self,  # noqa
            right_dataset=dataset,  # noqa
            join_condition=join_condition,
            asof_join=asof_join,
            direction=direction,
            allow_exact_matches=allow_exact_matches,
        )

    def filter(
        self, ts_filter: str | None = None, /, *,
        columns: list[str] | None = None,
        materialize: bool = False,
    ) -> IsFilteredDataset:
        """Filter the dataset.

        Args:
            ts_filter: Row filter condition in the form of a SQL WHERE clause
                (e.g."id = '12345' AND model = 'xyz'")
            columns: The specific columns to return.
            materialize: If True materialize the result of the filter operation as a
                temporary table within Snowflake.

        Returns:
            The filtered dataset object.
        """
        return type(self).filtered_dataset_cls.from_dataset(
            dataset=self,  # noqa
            ts_filter=ts_filter,
            columns=columns,
            materialize=materialize,
        )
