# 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 FeatureColumnSchema:
    """Schema for a feature column with optional timestamp.

    Represents a single feature column with its metadata, matching the proto structure.
    """

    feature_column: str  # e.g., "FEATURE_VARIANT__total_amount__v1"
    timestamp_column: (
        str  # e.g., "FEATURE_VARIANT__total_amount__v1__ts" (empty if no timestamp)
    )
    name: str  # e.g., "total_amount"
    variant: str  # e.g., "v1"


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

    entity_column: str
    feature_columns: List[FeatureColumnSchema]


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

    Attributes:
        feature_schema: The schema describing the feature view structure as a Python dataclass.
            Order: precomputed features first, then realtime features.
        rows: All values including timestamps (interleaved) and realtime feature results.
            Order: [entity_id, precomputed1_val, precomputed1_ts, ..., realtime1_val, ...]
        _raw_proto: The raw proto response from the server (for backwards compatibility).

    Backwards Compatibility:
        - result.schema: Returns the raw proto schema (same as pre-FeatureViewResult API)
        - result.value_lists: Returns the raw proto value_lists (same as pre-FeatureViewResult API)

    Important Notes:
        - Dependencies are filtered out from feature_schema and rows (internal use only)
        - Realtime feature values are computed client-side and included in rows
        - The backwards-compat properties (schema, value_lists) return raw proto data only,
          which does NOT include realtime feature values. Use rows or get_feature_values()
          when working with realtime features.

    Usage:
        # For most use cases (handles both precomputed and realtime features):
        result.get_feature_values()  # [entity, feat1, feat2, ..., rtf1, rtf2, ...]

        # For full data including timestamps:
        result.rows  # [entity, feat1, ts1, feat2, ts2, ..., rtf1, rtf2, ...]

        # For backwards compatibility (precomputed features only, no realtime values):
        result.value_lists[row_idx].values[col_idx]
    """

    feature_schema: FeatureViewSchema
    rows: List[List[Any]]
    _raw_proto: Any = None

    @property
    def schema(self):
        """Backwards-compatible access to raw proto schema.

        Note: This returns the raw proto schema from the server. If you're using
        realtime features, use feature_schema instead which has the correct order
        matching rows.

        Returns:
            The schema from the raw proto response.

        Raises:
            AttributeError: If the raw proto is not available.
        """
        if self._raw_proto is not None:
            return self._raw_proto.schema
        raise AttributeError(
            "Raw proto not available. This FeatureViewResult was not created from a server response."
        )

    @property
    def value_lists(self):
        """Backwards-compatible access to raw proto value_lists.

        This property provides access to the original proto response format
        for code that was written against the 0.16.x client API.

        Important: This returns ONLY server-side data. Realtime feature values
        are NOT included because they are computed client-side. If you're using
        realtime features, use rows or get_feature_values() instead.

        Returns:
            The value_lists from the raw proto response (precomputed features only).

        Raises:
            AttributeError: If the raw proto is not available.
        """
        if self._raw_proto is not None:
            return self._raw_proto.value_lists
        raise AttributeError(
            "Raw proto not available. This FeatureViewResult was not created from a server response."
        )

    def get_feature_values(self) -> List[List[Any]]:
        """Get rows with only feature values, excluding timestamps.

        This is the recommended method for most use cases. It returns feature values
        for both precomputed and realtime features, with timestamps filtered out.

        Order: [entity_id, precomputed1, precomputed2, ..., realtime1, realtime2, ...]

        Note: This method works correctly with realtime features, unlike the
        backwards-compat value_lists property which only contains server-side data.

        Returns:
            List of rows where each row is [entity_id, feat1_val, feat2_val, ...]
            without timestamp values. Includes realtime feature values.

        Example:
            # rows = [[entity, feat1, ts1, feat2, ts2, rtf1]]
            # get_feature_values() returns [[entity, feat1, feat2, rtf1]]
        """
        result = []
        for row in self.rows:
            filtered_row = [row[0]]  # Start with entity_id
            value_idx = 1  # Skip entity_id
            for feature_col in self.feature_schema.feature_columns:
                if value_idx < len(row):
                    filtered_row.append(row[value_idx])
                    value_idx += 1  # Move past feature value
                    if feature_col.timestamp_column:
                        value_idx += 1  # Skip timestamp
            # Include any remaining values (e.g., realtime features at the end)
            # that don't have timestamps
            while value_idx < len(row):
                filtered_row.append(row[value_idx])
                value_idx += 1
            result.append(filtered_row)
        return result


@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 (<)
    )
    has_timestamp: bool = False  # True if this feature has an associated timestamp


@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),
            has_timestamp=bool(col.timestamp_column),
        )

        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, val1, ts1, val2, ts2, ...]
        Non-dependency features with their timestamps interleaved, plus realtime results.
    """
    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] = {}

        # First pass: build precomputed value map for realtime feature execution
        # and collect non-dependency values for output
        # This uses the feature-based indexing, accounting for timestamps
        value_idx = 0
        for feature in all_precomputed:
            value = get_value_fn(row_idx, value_idx)
            key = (feature.name, feature.use_inclusive_join)
            precomputed_value_map[key] = value

            # Only include non-dependency values in output
            if not feature.is_dependency:
                processed_values.append(value)
                # Include timestamp if this feature has one
                if feature.has_timestamp:
                    ts_value = get_value_fn(row_idx, value_idx + 1)
                    processed_values.append(ts_value)

            value_idx += 1  # Move past the feature value
            # Skip timestamp in index
            if feature.has_timestamp:
                value_idx += 1

        # Add realtime feature results at the end
        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_feature_schema(
    entity_column: str,
    proto_feature_columns,
) -> FeatureViewSchema:
    """Build FeatureViewSchema from proto feature columns.

    The schema order matches how rows are built in process_rows:
    1. Precomputed features (non-dependency, non-realtime) in proto order
    2. Realtime features at the end

    Args:
        entity_column: Name of the entity column
        proto_feature_columns: Feature columns from the proto response

    Returns:
        FeatureViewSchema with full feature column metadata
    """
    precomputed_columns = []
    realtime_columns = []

    for col in proto_feature_columns:
        # Skip dependency columns - they are internal
        if getattr(col, "is_dependency", False):
            continue

        schema_col = FeatureColumnSchema(
            feature_column=col.feature_column,
            timestamp_column=col.timestamp_column if col.timestamp_column else "",
            name=col.name,
            variant=col.variant,
        )

        # Separate precomputed and realtime to match rows order
        if getattr(col, "serialized_realtime_definition", None):
            realtime_columns.append(schema_col)
        else:
            precomputed_columns.append(schema_col)

    # Order: precomputed first, then realtime (matches process_rows)
    return FeatureViewSchema(
        entity_column=entity_column,
        feature_columns=precomputed_columns + realtime_columns,
    )


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,
    raw_proto: Any = None,
    proto_feature_columns=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
        raw_proto: The raw proto response from the server (for backwards compatibility)
        proto_feature_columns: Feature columns from the proto response (for building schema)

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

    # Build schema from proto feature columns if available
    if proto_feature_columns is not None:
        feature_schema = build_feature_schema(entity_column, proto_feature_columns)
    else:
        # Fallback for cases without proto (e.g., local testing)
        feature_schema = FeatureViewSchema(
            entity_column=entity_column,
            feature_columns=[],
        )

    return FeatureViewResult(
        feature_schema=feature_schema,
        rows=processed_rows,
        _raw_proto=raw_proto,
    )
