from __future__ import annotations

from typing import (
    AsyncIterator,
    Callable,
    Dict,
    Iterator,
    List,
    Optional,
    Protocol,
    Tuple,
    runtime_checkable,
)
import typing as t

import pandas as pd
import pyarrow as pa

try:
    import tensorflow as tf
except ModuleNotFoundError:
    pass  # Warning is displayed by typing.py

import warnings

from sarus_data_spec.storage.typing import HasStorage
import sarus_data_spec.dataspec_rewriter.typing as sdrt
import sarus_data_spec.dataspec_validator.typing as sdvt
import sarus_data_spec.protobuf as sp
import sarus_data_spec.typing as st

try:
    import sqlalchemy as sa

    sa_engine = sa.engine.Engine
except ModuleNotFoundError:
    warnings.warn("Sqlalchemy not installed, cannot send sql queries")
    sa_engine = t.Any  # type: ignore


@runtime_checkable
class Manager(st.Referrable[sp.Manager], HasStorage, Protocol):
    """Provide the dataset functionalities"""

    def to_arrow(
        self, dataset: st.Dataset, batch_size: int
    ) -> t.Iterator[pa.RecordBatch]:
        """Synchronous method based on async_to_arrow
        that returns an iterator of arrow batches
        for the input dataset"""
        ...

    async def async_to_arrow(
        self, dataset: st.Dataset, batch_size: int
    ) -> AsyncIterator[pa.RecordBatch]:
        """Asynchronous method. It orchestrates how
        the iterator is obtained: it can either be delegated
        via arrow_task and the result polled, or computed directly
        it via the op"""
        ...

    def schema(self, dataset: st.Dataset) -> st.Schema:
        """Synchronous method that returns the schema of a
        dataspec. Based on the asynchronous version"""
        ...

    async def async_schema(self, dataset: st.Dataset) -> st.Schema:
        """Asynchronous method that returns the schema of a
        dataspec. The computation can be either delegated to
        another manager via schema_task and the result polled
        or executed directly via async_schema_ops"""
        ...

    def value(self, scalar: st.Scalar) -> st.DataSpecValue:
        """Synchronous method that returns the value of a
        scalar. Based on the asynchronous version"""
        ...

    async def async_value(self, scalar: st.Scalar) -> st.DataSpecValue:
        """Asynchronous method that returns the value of a
        scalar. The computation can be either delegated to
        another manager via value_task and the result polled
        or executed directly via async_value_ops"""
        ...

    def prepare(self, dataspec: st.DataSpec) -> None:
        """Make sure a Dataspec is ready."""
        ...

    async def async_prepare(self, dataspec: st.DataSpec) -> None:
        """Make sure a Dataspec is ready asynchronously."""
        ...

    async def async_prepare_parents(self, dataspec: st.DataSpec) -> None:
        """Prepare all the parents of a Dataspec."""
        ...

    def sql_prepare(self, dataset: st.Dataset) -> None:
        """Make sure a dataset is sql ready"""
        ...

    async def async_sql_prepare(self, dataset: st.Dataset) -> None:
        """Make sure a dataset is sql ready asynchronously."""
        ...

    async def async_sql_prepare_parents(self, dataset: st.Dataset) -> None:
        """SQL prepare all the parents of a dataset. It should sql_prepare
        dataset parents and prepare Scalars parents.
        """
        ...

    def cache_scalar(self, scalar: st.Scalar) -> None:
        """Synchronous scalar caching"""
        ...

    async def async_cache_scalar(self, scalar: st.Scalar) -> None:
        """Asynchronous scalar caching"""
        ...

    def to_parquet(self, dataset: st.Dataset) -> None:
        """Synchronous parquet caching"""
        ...

    async def async_to_parquet(self, dataset: st.Dataset) -> None:
        """Asynchronous parquet caching"""
        ...

    def parquet_dir(self) -> str:
        ...

    def marginals(self, dataset: st.Dataset) -> st.Marginals:
        ...

    async def async_marginals(self, dataset: st.Dataset) -> st.Marginals:
        ...

    def bounds(self, dataset: st.Dataset) -> st.Bounds:
        ...

    async def async_bounds(self, dataset: st.Dataset) -> st.Bounds:
        ...

    def size(self, dataset: st.Dataset) -> st.Size:
        ...

    async def async_size(self, dataset: st.Dataset) -> st.Size:
        ...

    def to_pandas(self, dataset: st.Dataset) -> pd.DataFrame:
        ...

    async def async_to_pandas(self, dataset: st.Dataset) -> pd.DataFrame:
        ...

    def to_tensorflow(self, dataset: st.Dataset) -> tf.data.Dataset:
        ...

    async def async_to_tensorflow(
        self, dataset: st.Dataset
    ) -> tf.data.Dataset:
        ...

    def to_sql(self, dataset: st.Dataset) -> None:
        ...

    async def async_to_sql(self, dataset: st.Dataset) -> None:
        ...

    def status(
        self, dataspec: st.DataSpec, task_name: str
    ) -> t.Optional[st.Status]:
        ...

    def dataspec_rewriter(self) -> sdrt.DataspecRewriter:
        ...

    def dataspec_validator(self) -> sdvt.DataspecValidator:
        ...

    def is_remote(self, dataspec: st.DataSpec) -> bool:
        """Is the dataspec a remotely defined dataset."""
        ...

    def infer_output_type(
        self,
        transform: st.Transform,
        *arguments: st.DataSpec,
        **named_arguments: st.DataSpec,
    ) -> Tuple[str, Callable[[st.DataSpec], None]]:
        ...

    def foreign_keys(self, dataset: st.Dataset) -> Dict[st.Path, st.Path]:
        ...

    async def async_foreign_keys(
        self, dataset: st.Dataset
    ) -> Dict[st.Path, st.Path]:
        ...

    async def async_primary_keys(self, dataset: st.Dataset) -> List[st.Path]:
        ...

    def primary_keys(self, dataset: st.Dataset) -> List[st.Path]:
        ...

    def sql(
        self,
        dataset: st.Dataset,
        query: t.Union[str, t.Mapping[t.Union[str, t.Tuple[str, ...]], str]],
        dialect: Optional[st.SQLDialect] = None,
        batch_size: int = 10000,
    ) -> Iterator[pa.RecordBatch]:
        ...

    async def async_sql(
        self,
        dataset: st.Dataset,
        query: t.Union[str, t.Mapping[t.Union[str, t.Tuple[str, ...]], str]],
        dialect: Optional[st.SQLDialect] = None,
        batch_size: int = 10000,
    ) -> AsyncIterator[pa.RecordBatch]:
        ...

    async def execute_sql_query(
        self,
        dataset: st.Dataset,
        caching_properties: t.Mapping[str, str],
        query: t.Union[str, t.Mapping[t.Union[str, t.Tuple[str, ...]], str]],
        dialect: t.Optional[st.SQLDialect] = None,
        batch_size: int = 10000,
    ) -> t.AsyncIterator[pa.RecordBatch]:
        ...

    async def async_sql_op(
        self,
        dataset: st.Dataset,
        query: t.Union[str, t.Mapping[t.Union[str, t.Tuple[str, ...]], str]],
        dialect: t.Optional[st.SQLDialect] = None,
        batch_size: int = 10000,
    ) -> t.AsyncIterator[pa.RecordBatch]:
        ...

    def is_big_data(self, dataset: st.Dataset) -> bool:
        ...

    def is_cached(self, dataspec: st.DataSpec) -> bool:
        """Returns whether a dataspec should be cached
        or not"""
        ...

    def is_pushed_to_sql(self, dataspec: st.DataSpec) -> bool:
        """Returns whether a dataspec should be pushed to sql
        or not"""
        ...

    def attribute(
        self, name: str, dataspec: st.DataSpec
    ) -> t.Optional[st.Attribute]:
        ...

    def attributes(
        self, name: str, dataspec: st.DataSpec
    ) -> t.List[st.Attribute]:
        ...

    def links(self, dataset: st.Dataset) -> st.Links:
        ...

    async def async_links(self, dataset: st.Dataset) -> st.Links:
        ...

    def sql_pushing_schema_prefix(self, dataset: st.Dataset) -> str:
        ...

    def engine(self, uri: str) -> sa_engine:
        ...

    def mock_value(
        self,
        transform: st.Transform,
        *arguments: st.DataSpec,
        **named_arguments: st.DataSpec,
    ) -> t.Any:
        """Compute the mock value of an external transform applied on
        Dataspecs.
        """


@runtime_checkable
class HasManager(Protocol):
    """Has a manager."""

    def manager(self) -> Manager:
        """Return a manager (usually a singleton)."""
        ...


T = t.TypeVar("T", covariant=True)


class Computation(t.Protocol[T]):
    """Protocol for classes that perform tasks computations.
    It sets how computations are scheduled and launched
    depending on statuses. A computation is mainly defined by two methods:
     - launch : a method that does not return a value but
     that only has side effects, changing either the storage or the cache
    and that updates statuses during the process
    - result: a method that allows to get the value of the computation
    either by reading the cache/storage or via the ops.

    Furthermore, a computation has a method to monitor task completion.
    """

    task_name: str = ''

    def launch_task(self, dataspec: st.DataSpec) -> t.Optional[t.Awaitable]:
        """This methods launches a task in the background
        but returns immediately without waiting for the
        result. It updates the statuses during its process."""
        ...

    async def task_result(self, dataspec: st.DataSpec, **kwargs: t.Any) -> T:
        """Returns the value for the given computed task. It either
        retrieves it from the cache or computes it via the ops."""
        ...

    async def complete_task(self, dataspec: st.DataSpec) -> st.Status:
        """Monitors a task: it launches it if there is no status
        and then polls until it is ready/error"""
        ...


class DelegatedComputation(Computation, t.Protocol[T]):
    def delegate_manager_status(
        self, dataspec: st.DataSpec
    ) -> Optional[st.Status]:
        """While some managers are eager and always execute tasks,
        others are lazy and delegate them. This is an interface
        for a manager to get the status of the eager manager
        that is executing the task"""
        ...


@t.runtime_checkable
class ExternalOpImplementation(t.Protocol):
    """External Op implementation class.

    The `allowed_pep_args` is a list of combinations of arguments' names which
    are managed by the Op. The result of the Op will be PEP only if the set of
    PEP arguments passed to the Op are in this list.

    For instance, if we have an op that takes 3 arguments `a`, `b` and `c` and
    the `allowed_pep_args` are [{'a'}, {'b'}, {'a','b'}] then the following
    combinations will yield a PEP output:
        - `a` is a PEP dataspec, `b` and `c` are either not dataspecs or public
          dataspecs
        - `b` is a PEP dataspec, `a` and `c` are either not dataspecs or public
          dataspecs
        - `a` and `b` are PEP dataspecs, `c` is either not a dataspec or a
          public dataspec
    """

    transform_id: str
    allowed_pep_args: t.List[t.Set[str]]

    def dataspec(self) -> st.DataSpec:
        ...

    def dp_equivalent(self) -> t.Optional[ExternalOpImplementation]:
        ...

    @staticmethod
    async def call(*args: t.Any, **kwargs: t.Any) -> t.Any:
        """Compute the external op output value.

        DP implementation take additional arguments:
            - `seed` an integer used to parametrize rangom number generators
            - `budget` the privacy budget that can be spend in the op
            - `pe` the protected entity information of each row
        """

    def pep_kind(self, *args: t.Any, **kwargs: t.Any) -> sdvt.PEPKind:
        """Return the PEP properties of the transform.

        It takes the transform arguments as input because it can depend on some
        transform parameters. For instance, it is not PEP if we are aggregating
        the rows (axis=0) and it is PEP if we are aggregating the columns
        (axis=1).
        """

    def is_dp(self, *args: t.Any, **kwargs: t.Any) -> bool:
        """Return True if the DP transform is compatible with the arguments.

        It takes the transform arguments as input because it can depend on some
        transform parameters. For instance, if we are aggregating the rows
        (axis=0), then there might be an equivalent DP transform but if we are
        aggregating the columns there might not (axis=1).

        NB: This function should have the same signature as the call function.
        """

    async def private_queries(
        self, *args: t.Any, **kwargs: t.Any
    ) -> t.List[st.PrivateQuery]:
        """Return the PrivateQueries summarizing DP characteristics."""
