import abc
import datetime
from dataclasses import dataclass
from types import MappingProxyType
from typing import Dict
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._internals.type_utils import sdk_type_from_tecton_type
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.framework.expression_parser import expression_to_proto
from tecton.types import Field
from tecton.types import SdkDataType
from tecton_core.specs.feature_spec import FeatureMetadataSpec
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):
    description: Optional[str]
    tags: Optional[Dict[str, str]]

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

    ```python
    from tecton import Aggregate, batch_feature_view, TimeWindow
    from tecton.types import Int64
    from datetime import timedelta

    @batch_feature_view(
        # ...
        features=[
            Aggregate(
                input_column=Field("my_column", Int64),
                function="mean",
                time_window=TimeWindow(window_size=timedelta(days=7)),
            ),
            Aggregate(
                input_column=Field("another_column", Int64),
                function="mean",
                time_window=TimeWindow(window_size=timedelta(days=1)),
                name="1d_average",
                description="my aggregate feature description",
                tags={"tag": "value"}
            ),
        ],
    )
    def my_fv(data_source):
        pass
    ```

    :param function: One of the built-in aggregation functions, such as "sum", "count", `last(2)` etc.
    :param time_window: The window_size and optional offset over which to aggregate over.
    :param name: The name of this feature. Defaults to an autogenerated name, e.g. transaction_count_7d_1d.
    :param input_column: Describes name and type of the column that will be used in the aggregation.
    :param description: A human-readable description of the feature
    :param tags: Tags associated with the feature (key-value pairs of user-defined metadata).
    """

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

    def __init__(
        self,
        function: Union[str, AggregationFunction],
        time_window: Union[TimeWindow, TimeWindowSeries, LifetimeWindow, datetime.timedelta],
        input_column: Field,
        name: Optional[str] = None,
        description: Optional[str] = None,
        tags: Optional[Dict[str, str]] = None,
    ):
        super().__init__(
            function=function,
            time_window=time_window,
            name=name,
            input_column=input_column,
            description=description,
            tags=tags,
        )

    @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.input_column.name,
            self.function,
            self.time_window,
            aggregation_interval,
            is_continuous,
            compaction_enabled,
            is_streaming_fv,
            self.description,
            self.tags,
        )
        proto.column_dtype.CopyFrom(self.input_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.

    ```python
    from tecton import Attribute, batch_feature_view
    from tecton.types import String

    @batch_feature_view(
        # ...
        features=[
            Attribute(
                name="my_column",
                dtype=String
            )
            Attribute(
                name="my_other_column",
                dtype=String,
                description="my attribute feature description",
                tags={"tag": "value"}
            )
        ],
    )
    def my_fv(data_source):
        pass
    ```

    :param name: Name of a column from the transformation output.
    :param dtype: Datatype of a column from the transformation output
    :param description: A human-readable description of the feature
    :param tags: Tags associated with the feature (key-value pairs of user-defined metadata).
    """

    name: str
    dtype: SdkDataType

    def __init__(
        self, name: str, dtype: SdkDataType, description: Optional[str] = None, tags: Optional[Dict[str, str]] = None
    ):
        super().__init__(name=name, dtype=dtype, description=description, tags=tags)

    def _to_proto(self):
        return feature_view_pb2.Attribute(
            name=self.name, column_dtype=self.dtype.tecton_type.proto, description=self.description, tags=self.tags
        )


class Calculation(Feature):
    """
    The `Calculation` class describes a Calculation feature that is applied to a Realtime Feature View via the `features` param.

    ```python
    from tecton import Calculation, RealtimeFeatureView

    transaction_counts_fv = RealtimeFeatureView(
        name="transaction_counts_features",
        sources=[request_ds, user_transaction_metrics],
        # mode=None implied
        ...
        features = [
            Calculation(
                name="transaction_count_fill_na",
                expr="COALESCE(user_transaction_metrics.transaction_count, request_ds.default_count)",
            ),
            Calculation(
                name="transaction_count_add_1",
                expr="user_transaction_metrics.transaction_count + 1",
                description="a test calculation utilizing addition",
                tags={"tag": "value"}
            ),
        ],
    )
    ```

    :param name: The name of this feature. Must be explicitly defined.
    :param expr: The calculation string expressing the operations and input operands.
    :param description: A human-readable description of the feature.
    :param tags: Tags associated with the feature (key-value pairs of user-defined metadata).
    """

    name: str
    expr: str
    description: Optional[str] = None
    tags: Optional[Dict[str, str]] = None

    def __init__(
        self,
        name: str,
        expr: str,
        description: Optional[str] = None,
        tags: Optional[Dict[str, str]] = None,
    ):
        super().__init__(name=name, expr=expr, description=description, tags=tags)

    def _to_proto(self, input_schema: Dict[str, SdkDataType]):
        root = expression_to_proto(self.expr, input_schema)
        return feature_view_pb2.Calculation(
            name=self.name, expr=self.expr, abstract_syntax_tree_root=root, column_dtype=root.dtype
        )


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


    ```python
    from tecton import Embedding, batch_feature_view
    from tecton.types import String

    @batch_feature_view(
        # ...
        features=[
            Embedding(
                input_column=Field("my_column", String),
                model="sentence-transformers/all-MiniLM-L6-v2",
                name="my_embedding_feature"
            ),
            Embedding(
                input_column=Field("my_other_column", String),
                model="sentence-transformers/all-MiniLM-L6-v2",
                name="my_other_embedding_feature",
                description="my embedding feature description",
                tags={"tag": "value"}
            )
        ],
    )
    def my_fv(data_source):
        return data_source
    ```
    :param input_column: Describes name and type of the column that will be used in the aggregation.
    :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.
    :param name: The name of this feature. Defaults to an autogenerated name, e.g. "my_embedding_feature".
    :param description: A human-readable description of the feature
    :param tags: Tags associated with the feature (key-value pairs of user-defined metadata).
    """

    input_column: Field
    model: str
    name: Optional[str] = None

    def __init__(
        self,
        input_column: Field,
        model: str,
        name: Optional[str] = None,
        description: Optional[str] = None,
        tags: Optional[Dict[str, str]] = None,
    ):
        super().__init__(model=model, name=name, input_column=input_column, description=description, tags=tags)

    def _to_proto(self) -> feature_view_pb2.Embedding:
        name = self.name or f"{self.input_column.name}_embedding"
        return feature_view_pb2.Embedding(
            name=name,
            column=self.input_column.name,
            column_dtype=self.input_column.dtype.tecton_type.proto,
            model=self.model,
            description=self.description,
            tags=self.tags,
        )


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

    ```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"
            ),
            Inference(
                input_columns=[Field("my_field3", String), Field("my_field4", Int64)],
                model="my_custom_model",
                name="my_other_inference_feature"m
                description="my inference feature description",
                tags={"tag": "value"}
            )
        ],
    )
    def my_fv(data_source):
        return data_source
    ```

    :param input_columns: A list of columns that are used as inputs to the model.
    :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.
    :param name: The name of this feature. Defaults to an autogenerated name, e.g. "my_inference_feature".
    :param description: A human-readable description of the feature
    :param tags: Tags associated with the feature (key-value pairs of user-defined metadata).
    """

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

    def __init__(
        self,
        model: str,
        input_columns: List[Field],
        name: str,
        description: Optional[str] = None,
        tags: Optional[Dict[str, str]] = None,
    ):
        super().__init__(model=model, input_columns=input_columns, name=name, description=description, tags=tags)

    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.Field(name=field.name, dtype=field.dtype.tecton_type.proto) for field in self.input_columns
            ],
            name=self.name,
            description=self.description,
            tags=self.tags,
        )


@dataclass
class FeatureMetadata:
    """
    `FeatureMetadata` contains metadata about your Feature.

    :param name: Name of the feature
    :param dtype: Data type of the feature
    :param description: A human-readable description of the feature
    :param tags: Tags associated with the feature (key-value pairs of user-defined metadata).
    """

    name: str
    dtype: SdkDataType
    description: Optional[str]
    tags: Optional[Dict[str, str]]

    @classmethod
    def _from_spec(cls, spec: FeatureMetadataSpec):
        return cls(
            name=spec.name, dtype=sdk_type_from_tecton_type(spec.dtype), description=spec.description, tags=spec.tags
        )
