# 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/.

"""
Shared Feature View Processing Logic

This module contains common functionality for processing feature view responses.

Key Components:
- RealtimeFeatureExecutor: Class that executes realtime feature functions with caching
- Module-level functions for stateless feature view processing operations
"""

import inspect
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional

import cloudpickle

from featureform.proto import metadata_pb2


@dataclass
class FeatureViewSchema:
    """Schema for a feature view result."""

    entity_column: str
    feature_columns: List[str]


@dataclass
class FeatureViewResult:
    """Result from a feature view query."""

    schema: FeatureViewSchema
    rows: List[List[Any]]


@dataclass
class ClassifiedFeature:
    """Represents a classified feature with its metadata."""

    name: str
    variant: str
    column_name: str  # e.g., "FEATURE_VARIANT__name__variant"
    is_realtime: bool = False
    is_dependency: bool = False
    serialized_realtime_definition: Optional[bytes] = None
    use_inclusive_join: bool = (
        False  # True for inclusive PIT join (<=), False for exclusive (<)
    )


@dataclass
class ClassifiedFeatures:
    """Container for classified features."""

    precomputed: List[ClassifiedFeature]
    dependency: List[ClassifiedFeature]
    realtime: List[ClassifiedFeature]

    def all_precomputed(self) -> List[ClassifiedFeature]:
        """Return all precomputed features including dependencies."""
        return self.precomputed + self.dependency

    def user_visible_columns(self) -> List[str]:
        """Return column names visible to the user (excludes dependencies)."""
        columns = [f.column_name for f in self.precomputed]
        columns.extend([f.name for f in self.realtime])
        return columns


class RealtimeFeatureExecutor:
    """Executes realtime feature functions with caching.

    This class maintains a cache of deserialized functions to avoid
    repeated cloudpickle.loads() calls for the same function.
    """

    def __init__(self):
        self._fn_cache: Dict[int, Callable] = {}

    def execute(
        self,
        serialized_definition: bytes,
        precomputed_values: Dict[str, Any],
        params: Optional[Dict[str, Any]] = None,
    ) -> Any:
        """Execute a realtime feature function.

        Args:
            serialized_definition: RealtimeFeatureDefinition proto bytes
            precomputed_values: Dict of precomputed feature values keyed by feature name
            params: Optional dict of runtime parameters for RealtimeInput

        Returns:
            The result of executing the realtime feature function
        """
        params = params or {}

        realtime_def = metadata_pb2.RealtimeFeatureDefinition()
        realtime_def.ParseFromString(serialized_definition)

        func = self._get_cached_function(realtime_def.serialized_function)
        kwargs = self._build_kwargs(
            func, realtime_def.inputs, precomputed_values, params
        )
        return func(**kwargs)

    def _get_cached_function(self, serialized_function: bytes) -> Callable:
        """Get a deserialized function, using cache if available."""
        cache_key = hash(serialized_function)
        if cache_key not in self._fn_cache:
            self._fn_cache[cache_key] = cloudpickle.loads(serialized_function)
        return self._fn_cache[cache_key]

    def _build_kwargs(
        self,
        func: Callable,
        inputs,
        precomputed_values: Dict[Any, Any],
        params: Dict[str, Any],
    ) -> Dict[str, Any]:
        """Build function keyword arguments from input specifications.

        Maps inputs to function parameters by position (not by name). This allows
        function parameters to have any name - they don't need to match feature names.

        For FeatureInput: looks up the feature value from precomputed_values using
        (name, False) as the key (exclusive PIT join).

        For RealtimeInput:
        - At serving time: looks up the value from params using parameter_name as key
        - At training time (when training_feature is set): looks up the backing
          feature value from precomputed_values using (name, True) as the key
          (inclusive PIT join)
        """
        kwargs = {}

        # Get function parameter names in order
        sig = inspect.signature(func)
        param_names = list(sig.parameters.keys())

        for i, input_spec in enumerate(inputs):
            if i >= len(param_names):
                raise ValueError(
                    f"More inputs ({len(inputs)}) than function parameters ({len(param_names)})"
                )
            param_name = param_names[i]

            if input_spec.HasField("feature_input"):
                feature = input_spec.feature_input.feature
                # FeatureInput uses EXCLUSIVE PIT join semantics (excludes events at label ts)
                key = (feature.name, False)
                # Also try string key for backward compatibility with older code
                if (
                    key not in precomputed_values
                    and feature.name not in precomputed_values
                ):
                    raise ValueError(
                        f"Missing required FeatureInput dependency '{feature.name}' "
                        f"for realtime feature."
                    )
                if key in precomputed_values:
                    kwargs[param_name] = precomputed_values[key]
                else:
                    # Backward compatibility: string key
                    kwargs[param_name] = precomputed_values[feature.name]
            elif input_spec.HasField("realtime_input"):
                realtime_input = input_spec.realtime_input

                # Check if this input has a training_feature backing it
                # (used when executing during training set iteration)
                if (
                    realtime_input.HasField("training_feature")
                    and realtime_input.training_feature.name
                ):
                    backing_name = realtime_input.training_feature.name
                    # RealtimeInput.training_feature uses INCLUSIVE PIT join semantics
                    # (includes events at exact label timestamp)
                    inclusive_key = (backing_name, True)
                    if inclusive_key in precomputed_values:
                        kwargs[param_name] = precomputed_values[inclusive_key]
                        continue
                    # Backward compatibility: try string key
                    if backing_name in precomputed_values:
                        kwargs[param_name] = precomputed_values[backing_name]
                        continue

                # Fall back to params (for serving time)
                # Use the realtime_input's parameter_name to look up in params
                realtime_param_name = realtime_input.parameter_name
                if realtime_param_name not in params:
                    raise ValueError(
                        f"Missing required parameter '{realtime_param_name}' for realtime feature. "
                        f"Please provide it in the params argument."
                    )
                kwargs[param_name] = params[realtime_param_name]

        return kwargs


# Module-level functions for stateless feature view processing


def parse_proto_columns(feature_columns) -> ClassifiedFeatures:
    """Parse feature columns from server proto response.

    Args:
        feature_columns: List of feature column protos from server response

    Returns:
        ClassifiedFeatures containing precomputed, dependency, and realtime features
    """
    precomputed = []
    dependency = []
    realtime = []

    for col in feature_columns:
        # Use explicit name/variant fields from proto - no parsing needed
        feature = ClassifiedFeature(
            name=col.name,
            variant=col.variant,
            column_name=col.feature_column,
            is_realtime=bool(col.serialized_realtime_definition),
            is_dependency=col.is_dependency,
            serialized_realtime_definition=(
                col.serialized_realtime_definition
                if col.serialized_realtime_definition
                else None
            ),
            use_inclusive_join=getattr(col, "use_inclusive_join", False),
        )

        if feature.is_realtime:
            realtime.append(feature)
        elif feature.is_dependency:
            dependency.append(feature)
        else:
            precomputed.append(feature)

    return ClassifiedFeatures(
        precomputed=precomputed, dependency=dependency, realtime=realtime
    )


def process_rows(
    classified: ClassifiedFeatures,
    entity_ids: List[str],
    get_value_fn: Callable[[int, int], Any],
    executor: RealtimeFeatureExecutor,
    params: Optional[Dict[str, Any]] = None,
) -> List[List[Any]]:
    """Process rows using a generic value getter function.

    Args:
        classified: Classified features from parse_proto_columns
        entity_ids: List of entity IDs for each row
        get_value_fn: Function that takes (row_idx, col_idx) and returns the value
        executor: RealtimeFeatureExecutor for executing realtime features
        params: Optional parameters for realtime feature execution

    Returns:
        List of processed rows, each containing [entity_id, feature1, feature2, ...]
    """
    all_precomputed = classified.all_precomputed()
    processed_rows = []

    for row_idx, entity_id in enumerate(entity_ids):
        processed_values = [entity_id]
        # Key precomputed values by (name, use_inclusive_join) to distinguish
        # between inclusive and exclusive PIT join values for the same feature.
        precomputed_value_map: Dict[tuple, Any] = {}

        for col_idx, feature in enumerate(all_precomputed):
            value = get_value_fn(row_idx, col_idx)
            key = (feature.name, feature.use_inclusive_join)
            precomputed_value_map[key] = value

            if not feature.is_dependency:
                processed_values.append(value)

        for feature in classified.realtime:
            if feature.serialized_realtime_definition:
                result = executor.execute(
                    feature.serialized_realtime_definition,
                    precomputed_value_map,
                    params,
                )
                processed_values.append(result)

        processed_rows.append(processed_values)

    return processed_rows


def build_response(
    classified: ClassifiedFeatures,
    entity_ids: List[str],
    get_value_fn: Callable[[int, int], Any],
    entity_column: str,
    executor: RealtimeFeatureExecutor,
    params: Optional[Dict[str, Any]] = None,
) -> FeatureViewResult:
    """Build the complete feature view result.

    Args:
        classified: Classified features from parse_proto_columns
        entity_ids: List of entity IDs for each row
        get_value_fn: Function that takes (row_idx, col_idx) and returns the value
        entity_column: Name of the entity column
        executor: RealtimeFeatureExecutor for executing realtime features
        params: Optional parameters for realtime feature execution

    Returns:
        FeatureViewResult with schema and rows
    """
    processed_rows = process_rows(
        classified=classified,
        entity_ids=entity_ids,
        get_value_fn=get_value_fn,
        executor=executor,
        params=params,
    )

    return FeatureViewResult(
        schema=FeatureViewSchema(
            entity_column=entity_column,
            feature_columns=classified.user_visible_columns(),
        ),
        rows=processed_rows,
    )
