# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

"""
Transformation decorator classes for Featureform.

This module contains decorator classes for SQL and DataFrame transformations.
"""

from __future__ import annotations

import ast
import inspect
import re
from dataclasses import dataclass, field
from datetime import timedelta
from typing import TYPE_CHECKING, Callable, List, Optional, Union

import dill
from typeguard import typechecked

from ..config.compute import EmptySparkFlags, K8sArgs, SparkFlags
from ..config.offline_stores import ResourceSnowflakeConfig
from ..core.exceptions import InvalidSQLQuery
from ..enums import ResourceType
from ..parse import add_variant_to_name, canonicalize_function_definition
from ..resources import (
    DFTransformation,
    PartitionType,
    SourceVariant,
    SQLTransformation,
)
from ..utils.helpers import get_name_variant

if TYPE_CHECKING:
    from ..providers import FileStoreProvider, OnlineProvider
    from . import EntityRegistrar, Registrar, UserRegistrar
    from .column_mapping import ColumnMapping


__all__ = [
    "SubscriptableTransformation",
    "SQLTransformationDecorator",
    "DFTransformationDecorator",
]


class SubscriptableTransformation:
    """
    SubscriptableTransformation creates a wrapped decorator that's callable and subscriptable,
    which allows for the following syntax:

    ``` py
    @local.transformation(variant="quickstart")
    def average_user_transaction():
        return "SELECT CustomerID as user_id, avg(TransactionAmount) as avg_transaction_amt from {{transactions.v1}} GROUP BY user_id"

    feature = ff.Feature(average_user_transaction[["user_id", "avg_transaction_amt"]])
    ```

    Given the function type does not implement __getitem__ we need to wrap it in a class that
    enables this behavior while still maintaining the original function signature and behavior.
    """

    def __init__(
        self,
        fn,
        registrar,
        provider,
        decorator_register_resources_method,
        decorator_name_variant_method,
        transformation,
    ):
        # if not self.__has_return_statement(fn):
        #     raise Exception(
        #         "Transformation function seems to be missing a return statement"
        #     )

        self.fn = fn
        self.registrar = registrar
        self.provider = provider
        # Previously, the descriptor protocol was used to apply methods from the decorator classes
        # to instances of SubscriptableTransformation such that a user could call `fn.name_variant()`
        # and receive a tuple of (name, variant) where name was the name of the wrapped function and
        # variant was either the value passed to the decorator or the default value. This was achieved
        # via the following syntax: `self.name_variant = decorator_name_variant_method.__get__(self)`
        # For as-of-yet unknown reasons, this behavior was not working as expected in Python 3.11.2,
        # so the code has been reverted to the original syntax, which simply passes a reference to
        # the decorator methods to the SubscriptableTransformation class.
        self.register_resources = decorator_register_resources_method
        self.name_variant = decorator_name_variant_method
        self.transformation = transformation

    def name_variant(self):
        return self.transformation.name_variant()

    def get_resource_type(self):
        return self.transformation.get_resource_type()

    def to_key(self):
        n, v = self.name_variant()
        return self.get_resource_type(), n, v

    def __getitem__(self, columns: List[str]):
        col_len = len(columns)
        if col_len < 2:
            raise Exception(
                f"Expected 2 columns, but found {col_len}. Missing entity and/or source columns"
            )
        elif col_len > 3:
            raise Exception(
                f"Found unrecognized columns {', '.join(columns[3:])}. Expected 2 required columns and an optional 3rd timestamp column"
            )
        return (self.registrar, self.transformation, columns)

    def __call__(self, *args, **kwargs):
        return self.fn(*args, **kwargs)

    @staticmethod
    def __has_return_statement(fn):
        """
        Parses the function’s source code into an abstract syntax tree
        and then walks through the tree to check for any Return nodes.
        Not full-proof but will at least catch cases on the client.
        """
        tree = ast.parse(inspect.getsource(fn))
        for node in ast.walk(tree):
            if isinstance(node, ast.Return):
                return True
        return False


@dataclass
class SQLTransformationDecorator:
    registrar: "Registrar"
    owner: str
    provider: str
    tags: List[str]
    properties: dict
    run: str = ""
    variant: str = ""
    name: str = ""
    schedule: str = ""
    incremental: bool = False
    incremental_inputs: list = field(default_factory=list)
    inputs: list = field(default_factory=list)
    description: str = ""
    args: Union[K8sArgs, None] = None
    query: str = field(default_factory=str, init=False)
    partition_options: Optional[PartitionType] = None
    func_params_to_inputs: dict = field(default_factory=dict, init=False)
    max_job_duration: timedelta = timedelta(hours=48)
    spark_flags: SparkFlags = field(default_factory=lambda: EmptySparkFlags)
    resource_snowflake_config: Optional[ResourceSnowflakeConfig] = None

    def __post_init__(self):
        if self.inputs is None:
            self.inputs = []
        if self.incremental_inputs is None:
            self.incremental_inputs = []

    def __call__(self, fn: Callable[[], str]):
        if self.description == "" and fn.__doc__ is not None:
            self.description = fn.__doc__
        if self.name == "":
            self.name = fn.__name__

        func_params = inspect.signature(fn).parameters

        if len(func_params) > 0:
            if len(func_params) > len(self.inputs):
                raise ValueError(
                    f"Transformation function has more parameters than inputs. \n"
                    f"Make sure each function parameter has a corresponding input in the decorator."
                )

            if len(func_params) < len(self.inputs):
                raise ValueError(
                    f"Too many inputs for transformation function. Expected {len(func_params)} inputs, but found {len(self.inputs)}.\n"
                )

            if not isinstance(self.inputs, list):
                raise ValueError("Dataframe transformation inputs must be a list")

            self.func_params_to_inputs = dict(zip(func_params, self.inputs))

            self.__set_query(fn(*self.inputs))
        else:
            self.__set_query(fn())

        self.registrar.map_client_object_to_resource(self, self.to_source())
        self.registrar.add_resource(self.to_source())
        return SubscriptableTransformation(
            fn,
            self.registrar,
            self.provider,
            self.register_resources,
            self.name_variant,
            self,
        )

    @typechecked
    def __set_query(self, query: str):
        if query == "":
            raise ValueError("Query cannot be an empty string")

        self._assert_query_contains_at_least_one_source(query)
        if len(self.inputs) > 0:
            # if inputs are specified, then the query will be resolved at the time of creation (when #kwargs is called)
            self.query = query
        else:
            self.query = add_variant_to_name(query, self.run)

    def to_source(self) -> SourceVariant:
        return SourceVariant(
            created=None,
            name=self.name,
            variant=self.variant,
            definition=SQLTransformation(
                query=self.query,
                args=self.args,
                func_params_to_inputs=self.func_params_to_inputs,
                is_incremental=self.incremental,
                incremental_inputs=self.incremental_inputs,
                partition_options=self.partition_options,
                spark_flags=self.spark_flags,
                resource_snowflake_config=self.resource_snowflake_config,
            ),
            owner=self.owner,
            schedule=self.schedule,
            provider=self.provider,
            description=self.description,
            tags=self.tags,
            properties=self.properties,
            max_job_duration=self.max_job_duration,
        )

    def name_variant(self):
        return (self.name, self.variant)

    def get_resource_type(self):
        return ResourceType.SOURCE_VARIANT

    def register_resources(
        self,
        entity: Union[str, "EntityRegistrar"],
        entity_column: str,
        owner: Union[str, "UserRegistrar"] = "",
        inference_store: Union[str, "OnlineProvider", "FileStoreProvider"] = "",
        features: List["ColumnMapping"] = None,
        labels: List["ColumnMapping"] = None,
        timestamp_column: str = "",
        description: str = "",
        schedule: str = "",
    ):
        return self.registrar.register_column_resources(
            source=self,
            entity=entity,
            entity_column=entity_column,
            owner=owner,
            inference_store=inference_store,
            features=features,
            labels=labels,
            timestamp_column=timestamp_column,
            description=description,
            schedule=schedule,
            client_object=self,
        )

    @staticmethod
    def _assert_query_contains_at_least_one_source(query):
        # Checks to verify that the query contains a FROM {{ name.variant }}

        # the pattern pulls the string within the double curly braces
        pattern = r"\{\{\s*(.*?)\s*\}\}"
        matches = re.findall(pattern, query)
        if len(matches) == 0:
            raise InvalidSQLQuery(query, "No source specified.")

        for m in matches:
            name, variant = get_name_variant(query, m)
            if name == "":
                raise InvalidSQLQuery(query, "Source name is empty.")

            # Check for invalid characters in the source name and variant
            if name.startswith(" ") or name.endswith(" "):
                raise InvalidSQLQuery(
                    query, "Source name cannot start or end with a space."
                )
            if variant.startswith(" ") or variant.endswith(" "):
                raise InvalidSQLQuery(
                    query, "Source variant cannot start or end with a space."
                )

            if name.startswith("_") or name.endswith("_"):
                raise InvalidSQLQuery(
                    query, "Source name cannot start or end with an underscore."
                )
            if variant.startswith("_") or variant.endswith("_"):
                raise InvalidSQLQuery(
                    query, "Source variant cannot start or end with an underscore."
                )

            if "__" in name or "__" in variant:
                raise InvalidSQLQuery(
                    query,
                    "Source name and variant cannot contain consecutive underscores.",
                )


# get_name_variant moved to utils/helpers.py


@dataclass
class DFTransformationDecorator:
    registrar: "Registrar"
    owner: str
    provider: str
    tags: List[str]
    properties: dict
    variant: str = ""
    name: str = ""
    description: str = ""
    incremental: bool = False
    incremental_inputs: list = field(default_factory=list)
    partition_options: Optional[PartitionType] = None
    inputs: list = field(default_factory=list)
    args: Union[K8sArgs, None] = None
    source_text: str = ""
    canonical_func_text: str = ""
    query: bytes = field(default_factory=bytes, init=False)
    max_job_duration: timedelta = timedelta(hours=48)
    spark_flags: SparkFlags = field(default_factory=lambda: EmptySparkFlags)

    def __call__(self, fn):
        if self.description == "" and fn.__doc__ is not None:
            self.description = fn.__doc__
        if self.name == "":
            self.name = fn.__name__

        func_params = inspect.signature(fn).parameters
        if len(func_params) > len(self.inputs):
            raise ValueError(
                f"Transformation function has more parameters than inputs. \n"
                f"Make sure each function parameter has a corresponding input in the decorator."
            )

        if not isinstance(self.inputs, list):
            raise ValueError("Dataframe transformation inputs must be a list")

        # Check that the function doesn't have free variables.
        # A free variable is one that is captured from the environment. This is an issue
        # because it's not serialized with the function, so later when it's deserialized
        # and attempted to run, it will complain about unbound variables.
        if fn.__code__.co_freevars:
            raise ValueError(
                f"a dataframe transformation must be top-level and self-contained; free vars: {fn.__code__.co_freevars}"
            )

        # check that input isn't self referencing
        for nv in self.inputs:
            if isinstance(
                nv, tuple
            ):  # TODO all that should be called here is name_variant()
                n, v = nv
            else:
                n, v = nv.name_variant()
            if self.name is n and self.variant is v:
                raise ValueError(
                    f"Transformation cannot be input for itself: {self.name} {self.variant}"
                )
        self.query = dill.dumps(fn.__code__)
        self.source_text = dill.source.getsource(fn)
        self.canonical_func_text = canonicalize_function_definition(fn)
        self.registrar.map_client_object_to_resource(self, self.to_source())
        self.registrar.add_resource(self.to_source())
        return SubscriptableTransformation(
            fn,
            self.registrar,
            self.provider,
            self.register_resources,
            self.name_variant,
            self,
        )

    def to_source(self) -> SourceVariant:
        return SourceVariant(
            created=None,
            name=self.name,
            variant=self.variant,
            definition=DFTransformation(
                query=self.query,
                inputs=self.inputs,
                args=self.args,
                source_text=self.source_text,
                canonical_func_text=self.canonical_func_text,
                spark_flags=self.spark_flags,
                is_incremental=self.incremental,
                incremental_inputs=self.incremental_inputs,
                partition_options=self.partition_options,
            ),
            owner=self.owner,
            provider=self.provider,
            description=self.description,
            tags=self.tags,
            properties=self.properties,
            max_job_duration=self.max_job_duration,
        )

    def name_variant(self):
        return (self.name, self.variant)

    def register_resources(
        self,
        entity: Union[str, "EntityRegistrar"],
        entity_column: str,
        owner: Union[str, "UserRegistrar"] = "",
        inference_store: Union[str, "OnlineProvider", "FileStoreProvider"] = "",
        features: List["ColumnMapping"] = None,
        labels: List["ColumnMapping"] = None,
        timestamp_column: str = "",
        description: str = "",
    ):
        return self.registrar.register_column_resources(
            source=self,
            entity=entity,
            entity_column=entity_column,
            owner=owner,
            inference_store=inference_store,
            features=features,
            labels=labels,
            timestamp_column=timestamp_column,
            description=description,
            client_object=self,
        )


# ColumnSourceRegistrar, ResourceRegistrar, and ModelRegistrar are imported from .registrar
