import abc
import datetime
from types import MappingProxyType
from typing import List
from typing import Optional
from typing import Union

from tecton._internals.tecton_pydantic import StrictModel
from tecton._internals.tecton_pydantic import pydantic_v1
from tecton.aggregation_functions import AggregationFunction
from tecton.framework.configs import LifetimeWindow
from tecton.framework.configs import TimeWindow
from tecton.framework.configs import TimeWindowSeries
from tecton.framework.configs import build_aggregation_proto
from tecton.types import Field
from tecton.types import SdkDataType
from tecton.types import String
from tecton_proto.args import feature_view__client_pb2 as feature_view_pb2
from tecton_proto.common import schema__client_pb2 as schema_pb2


class Feature(StrictModel, abc.ABC):
    column: str
    column_dtype: SdkDataType

    @abc.abstractmethod
    def _to_proto(self, **kwargs):
        pass


class Aggregate(Feature):
    """
    The `Aggregate` class describes an aggregation feature that is applied to a Batch or Stream Feature View via `features` param.

    :param column: Column name of the feature.
    :type column: str
    :param column_dtype: The data type of the column.
    :type column_dtype: SdkDataType
    :param function: One of the built-in aggregation functions, such as "sum", "count", `last(2)` etc.
    :type function: Union[str, AggregationFunction]
    :param time_window: The window_size and optional offset over which to aggregate over.
    :type time_window: Union[TimeWindow, TimeWindowSeries, LifetimeWindow, datetime.timedelta]
    :param name: The name of this feature. Defaults to an autogenerated name, e.g. transaction_count_7d_1d.
    :type name: str

    .. code-block:: python
        from tecton import Aggregate, batch_feature_view, TimeWindow
        from tecton.types import Int64
        from datetime import timedelta


        @batch_feature_view(
            # ...
            features=[
                Aggregate(
                    column="my_column",
                    column_dtype=Int64,
                    function="mean",
                    time_window=TimeWindow(window_size=timedelta(days=7)),
                ),
                Aggregate(
                    column="another_column",
                    column_dtype=Int64,
                    function="mean",
                    time_window=TimeWindow(window_size=timedelta(days=1)),
                    name="1d_average",
                ),
            ],
        )
        def my_fv(data_source):
            pass

    """

    function: AggregationFunction
    time_window: Union[TimeWindow, TimeWindowSeries, LifetimeWindow]
    name: Optional[str] = None

    def __init__(
        self,
        column: str,
        column_dtype: SdkDataType,
        function: Union[str, AggregationFunction],
        time_window: Union[TimeWindow, TimeWindowSeries, LifetimeWindow, datetime.timedelta],
        name: Optional[str] = None,
    ):
        super().__init__(
            column=column, column_dtype=column_dtype, function=function, time_window=time_window, name=name
        )

    @pydantic_v1.validator("time_window", pre=True)
    def timedelta_to_time_window(cls, v):
        if isinstance(v, datetime.timedelta):
            return TimeWindow(window_size=v)
        return v

    @pydantic_v1.validator("function", pre=True)
    def str_to_aggregation_function(cls, v):
        if isinstance(v, str):
            return AggregationFunction(base_name=v, resolved_name=v, params=MappingProxyType({}))
        return v

    def _to_proto(
        self,
        aggregation_interval: datetime.timedelta,
        is_continuous: bool,
        compaction_enabled: bool = False,
        is_streaming_fv: bool = False,
    ) -> feature_view_pb2.FeatureAggregation:
        proto = build_aggregation_proto(
            self.name,
            self.column,
            self.function,
            self.time_window,
            aggregation_interval,
            is_continuous,
            compaction_enabled,
            is_streaming_fv,
        )
        proto.column_dtype.CopyFrom(self.column_dtype.tecton_type.proto)
        return proto


class Attribute(Feature):
    """
    The `Attribute` class describes an attribute feature that is applied to a Feature View or Feature Table via `features` param.

    :param column: Column name of the feature.
    :type column: str
    :param column_dtype: The data type of the column.
    :type column_dtype: SdkDataType


    .. code-block:: python
    from tecton import Attribute, batch_feature_view
    from tecton.types import String

    @batch_feature_view(
        # ...
        features=[Attribute(column="my_column", column_dtype=String)],
    )
    def my_fv(data_source):
        pass
    """

    def __init__(self, column: str, column_dtype: SdkDataType):
        super().__init__(column=column, column_dtype=column_dtype)

    def _to_proto(self):
        return feature_view_pb2.Attribute(
            column=self.column,
            column_dtype=self.column_dtype.tecton_type.proto,
        )


class Embedding(Feature):
    """
    The `Embedding` class describes an embedding feature that is applied to a Batch Feature View via `features` param.

    :param column: The name of the column to be embedded.
    :type column: str
    :param column_dtype: The data type of the column to be embedded.
    :type column_dtype: SdkDataType
    :param model: The model name that is used to compute the embedding feature. Check https://docs.tecton.ai/docs/beta/defining-features/feature-views/embeddings for all supported models.
    :type model: str
    :param name: The name of this feature. Defaults to an autogenerated name, e.g. "my_embedding_feature".
    :type name: str


    .. code-block:: python
    from tecton import Embedding, batch_feature_view
    from tecton.types import String

    @batch_feature_view(
        # ...
        features=[
            Embedding(
                column="my_column",
                column_dtype=String,
                model="sentence-transformers/all-MiniLM-L6-v2",
                name="my_embedding_feature")
        ],
    )
    def my_fv(data_source):
        return data_source
    """

    model: str
    name: Optional[str] = None

    def __init__(self, column: str, column_dtype: SdkDataType, model: str, name: Optional[str] = None):
        super().__init__(column=column, column_dtype=column_dtype, model=model, name=name)

    def _to_proto(self) -> feature_view_pb2.Embedding:
        name = self.name if self.name else f"{self.column}_embedding"
        return feature_view_pb2.Embedding(
            name=name, column=self.column, column_dtype=self.column_dtype.tecton_type.proto, model=self.model
        )


class Inference(Feature):
    """
    The `Inference` class describes a model inference feature that is applied to a Batch Feature View via `features` param.

    :param input_columns: A list of columns that are used as inputs to the model.
    :type input_columns: List[Field]
    :param model: The model name that is used to compute the inference feature. The model needs to be registered in advance using `tecton model` CLI.
    :type model: str
    :param name: The name of this feature. Defaults to an autogenerated name, e.g. "my_inference_feature".
    :type name: Optional[str]


    .. code-block:: python
    from tecton import Inference, batch_feature_view
    from tecton.types import String, Int64

    @batch_feature_view(
        # ...
        features=[
            Inference(
                input_columns=[Field("my_field1", String), Field("my_field2", Int64)],
                model="my_custom_model",
                name="my_inference_feature")
        ],
    )
    def my_fv(data_source):
        return data_source
    """

    model: str
    input_columns: List[Field]
    name: Optional[str]

    def __init__(self, model: str, input_columns: List[Field], name: str):
        # TODO(jiadong): `column` and `column_dtype` are required right now. Need to refactor `Features` to accommodate new `Inferece` class.
        super().__init__(model=model, input_columns=input_columns, name=name, column="", column_dtype=String)

    def _to_proto(self) -> feature_view_pb2.Inference:
        # TODO(jiadong): Figure out a default name for `Inference` feature.
        return feature_view_pb2.Inference(
            model=self.model,
            input_columns=[
                schema_pb2.Column(name=field.name, offline_data_type=field.dtype.tecton_type.proto)
                for field in self.input_columns
            ],
            name=self.name,
        )
