from __future__ import annotations

from functools import cached_property
from typing import Any
from typing import Literal
from typing import NoReturn
from typing import overload
from typing import Protocol
from typing import TYPE_CHECKING

import numpy as np  # noqa: F401
import pandas as pd
import pyarrow as pa
import pyarrow.flight as flight

from tempora.utils.dataclass import IsDataclass
from tempora.utils.dataset import AsofDirection

if TYPE_CHECKING:
    from .base import JoinData
    from .base import Pivot


class IsDataset(IsDataclass, Protocol):
    # Since Datasets are "frozen" we require the attributes are read-only. If we define
    # them in the usual way within the Protocol we get the following mypy error:
    # "Protocol member IsDataset.name expected settable variable got readonly attribute"
    # ... and similarly for the other Dataset attributes.
    # To fix this we can define the Proto attributes as properties (without a setter):
    # https://github.com/python/typing/discussions/903

    # ----------- Start of attributes -----------
    @property
    def time_column(self) -> str | None: ...

    @property
    def entity_keys(self) -> list[str] | None: ...

    @property
    def pivot(self) -> Pivot | None: ...

    @property
    def materialize(self) -> bool: ...

    @property
    def targets(self) -> bool: ...

    @property
    def _join_data(self) -> list[JoinData] | None: ...

    # ------------ End of attributes ------------

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

    def _finalize(self) -> None: ...

    def __hash__(self) -> int: ...

    def __eq__(self, other: object) -> bool: ...

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

    @cached_property
    def _as_json(self) -> str: ...

    @property
    def _flight_reader(self) -> flight.FlightStreamReader: ...

    @property
    def schema(self) -> pa.Schema: ...

    @property
    def num_rows(self) -> int: ...

    def drop(self) -> None | NoReturn: ...

    @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: ...

    def to_pandas(self, use_nullable_dtypes: bool = False) -> pd.DataFrame: ...

    def df(self, use_nullable_dtypes: bool = False) -> pd.DataFrame: ...

    def to_numpy(self) -> np.ndarray: ...

    def np(self) -> np.ndarray: ...  # noqa: F811

    @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: ...

    def write_dataset(self, *args: Any, **kwargs: Any) -> None: ...

    @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:
        ...

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

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


class IsFilteredDataset(IsDataset, Protocol):
    # Define all attributes as read-only since mixin classes are "frozen" (see above)
    @property
    def ts_filter(self) -> str | None: ...

    @property
    def columns(self) -> list[str] | None: ...

    @classmethod
    def from_dataset(
        cls, dataset: IsDataset,
        ts_filter: str | None = None,
        columns: list[str] | None = None,
        materialize: bool = False,
    ) -> IsFilteredDataset:
        ...
