import json
import logging
from datetime import datetime, timedelta
from typing import List, Optional, Set, Union

import pyarrow as pa

from feast import BatchFeatureView, StreamFeatureView
from feast.data_source import DataSource
from feast.infra.compute_engines.backends.base import DataFrameBackend
from feast.infra.compute_engines.dag.context import ColumnInfo, ExecutionContext
from feast.infra.compute_engines.dag.model import DAGFormat
from feast.infra.compute_engines.dag.node import DAGNode
from feast.infra.compute_engines.local.arrow_table_value import ArrowTableValue
from feast.infra.compute_engines.local.local_node import LocalNode
from feast.infra.compute_engines.utils import (
    create_offline_store_retrieval_job,
)
from feast.infra.offline_stores.offline_utils import (
    infer_event_timestamp_from_entity_df,
)
from feast.utils import _convert_arrow_to_proto

logger = logging.getLogger(__name__)

ENTITY_TS_ALIAS = "__entity_event_timestamp"


class LocalSourceReadNode(LocalNode):
    def __init__(
        self,
        name: str,
        source: DataSource,
        column_info: ColumnInfo,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
    ):
        super().__init__(name)
        self.source = source
        self.column_info = column_info
        self.start_time = start_time
        self.end_time = end_time

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        retrieval_job = create_offline_store_retrieval_job(
            data_source=self.source,
            context=context,
            start_time=self.start_time,
            end_time=self.end_time,
            column_info=self.column_info,
        )
        arrow_table = retrieval_job.to_arrow()
        if self.column_info.field_mapping:
            arrow_table = arrow_table.rename_columns(
                [
                    self.column_info.field_mapping.get(col, col)
                    for col in arrow_table.column_names
                ]
            )
        return ArrowTableValue(data=arrow_table)


class LocalJoinNode(LocalNode):
    def __init__(
        self,
        name: str,
        column_info: ColumnInfo,
        backend: DataFrameBackend,
        inputs: Optional[List["DAGNode"]] = None,
        how: str = "inner",
    ):
        super().__init__(name, inputs or [])
        self.column_info = column_info
        self.backend = backend
        self.how = how

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        input_values = self.get_input_values(context)
        for val in input_values:
            val.assert_format(DAGFormat.ARROW)

        # Convert all upstream ArrowTables to backend DataFrames
        joined_df = self.backend.from_arrow(input_values[0].data)
        for val in input_values[1:]:
            next_df = self.backend.from_arrow(val.data)
            joined_df = self.backend.join(
                joined_df,
                next_df,
                on=self.column_info.join_keys,
                how=self.how,
            )

        # If entity_df is provided, join it in last
        if context.entity_df is not None:
            entity_df = self.backend.from_arrow(pa.Table.from_pandas(context.entity_df))

            entity_schema = dict(zip(entity_df.columns, entity_df.dtypes))
            entity_ts_col = infer_event_timestamp_from_entity_df(entity_schema)

            if entity_ts_col != ENTITY_TS_ALIAS:
                entity_df = self.backend.rename_columns(
                    entity_df, {entity_ts_col: ENTITY_TS_ALIAS}
                )

            joined_df = self.backend.join(
                entity_df,
                joined_df,
                on=self.column_info.join_keys,
                how="left",
            )

        result = self.backend.to_arrow(joined_df)
        return ArrowTableValue(result)


class LocalFilterNode(LocalNode):
    def __init__(
        self,
        name: str,
        column_info: ColumnInfo,
        backend: DataFrameBackend,
        filter_expr: Optional[str] = None,
        ttl: Optional[timedelta] = None,
        inputs=None,
    ):
        super().__init__(name, inputs=inputs)
        self.column_info = column_info
        self.backend = backend
        self.filter_expr = filter_expr
        self.ttl = ttl

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        input_table = self.get_single_table(context).data
        df = self.backend.from_arrow(input_table)

        timestamp_column = self.column_info.timestamp_column

        if ENTITY_TS_ALIAS in self.backend.columns(df):
            # filter where feature.ts <= entity.event_timestamp
            df = df[df[timestamp_column] <= df[ENTITY_TS_ALIAS]]

            # TTL: feature.ts >= entity.event_timestamp - ttl
            if self.ttl:
                lower_bound = df[ENTITY_TS_ALIAS] - self.backend.to_timedelta_value(
                    self.ttl
                )
                df = df[df[timestamp_column] >= lower_bound]

        # Optional user-defined filter expression (e.g., "value > 0")
        if self.filter_expr:
            df = self.backend.filter(df, self.filter_expr)

        result = self.backend.to_arrow(df)
        output = ArrowTableValue(result)
        context.node_outputs[self.name] = output
        return output


class LocalAggregationNode(LocalNode):
    def __init__(
        self,
        name: str,
        backend: DataFrameBackend,
        group_keys: list[str],
        agg_ops: dict,
        inputs=None,
    ):
        super().__init__(name, inputs=inputs)
        self.backend = backend
        self.group_keys = group_keys
        self.agg_ops = agg_ops

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        input_table = self.get_single_table(context).data
        df = self.backend.from_arrow(input_table)
        grouped_df = self.backend.groupby_agg(df, self.group_keys, self.agg_ops)
        result = self.backend.to_arrow(grouped_df)
        output = ArrowTableValue(result)
        return output


class LocalDedupNode(LocalNode):
    def __init__(
        self, name: str, column_info: ColumnInfo, backend: DataFrameBackend, inputs=None
    ):
        super().__init__(name, inputs=inputs)
        self.column_info = column_info
        self.backend = backend

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        input_table = self.get_single_table(context).data
        df = self.backend.from_arrow(input_table)

        # Extract join_keys, timestamp, and created_ts from context

        # Dedup strategy: sort and drop_duplicates
        dedup_keys = self.column_info.join_keys
        if dedup_keys:
            sort_keys = [self.column_info.timestamp_column]
            if (
                self.column_info.created_timestamp_column
                and self.column_info.created_timestamp_column in df.columns
            ):
                sort_keys.append(self.column_info.created_timestamp_column)

            df = self.backend.drop_duplicates(
                df, keys=dedup_keys, sort_by=sort_keys, ascending=False
            )
        result = self.backend.to_arrow(df)
        output = ArrowTableValue(result)
        context.node_outputs[self.name] = output
        return output


class LocalTransformationNode(LocalNode):
    def __init__(
        self, name: str, transformation_fn, backend: DataFrameBackend, inputs=None
    ):
        super().__init__(name, inputs=inputs)
        self.transformation_fn = transformation_fn
        self.backend = backend

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        input_table = self.get_single_table(context).data
        df = self.backend.from_arrow(input_table)
        transformed_df = self.transformation_fn(df)
        result = self.backend.to_arrow(transformed_df)
        output = ArrowTableValue(result)
        context.node_outputs[self.name] = output
        return output


class LocalValidationNode(LocalNode):
    def __init__(
        self, name: str, validation_config, backend: DataFrameBackend, inputs=None
    ):
        super().__init__(name, inputs=inputs)
        self.validation_config = validation_config
        self.backend = backend

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        input_table = self.get_single_table(context).data

        if self.validation_config:
            self._validate_schema(input_table)

        output = ArrowTableValue(input_table)
        context.node_outputs[self.name] = output
        return output

    def _validate_schema(self, table: pa.Table):
        """Validate that the input table conforms to the expected schema.

        Checks that all expected columns are present, that their types
        are compatible with the declared Feast types, and that Json columns
        contain well-formed JSON. Logs warnings for type mismatches but
        raises on missing columns or invalid JSON content.
        """
        expected_columns = self.validation_config.get("columns", {})
        if not expected_columns:
            logger.debug(
                "[Validation: %s] No column schema to validate against.",
                self.name,
            )
            return

        actual_columns = set(table.column_names)
        expected_names = set(expected_columns.keys())

        missing = expected_names - actual_columns
        if missing:
            raise ValueError(
                f"[Validation: {self.name}] Missing expected columns: {missing}. "
                f"Actual columns: {sorted(actual_columns)}"
            )

        for col_name, expected_type in expected_columns.items():
            actual_type = table.schema.field(col_name).type
            if expected_type is not None and actual_type != expected_type:
                # PyArrow map columns and struct columns are compatible
                # with the Feast Map type — skip warning for these cases
                if pa.types.is_map(expected_type) and (
                    pa.types.is_map(actual_type)
                    or pa.types.is_struct(actual_type)
                    or pa.types.is_large_list(actual_type)
                    or pa.types.is_list(actual_type)
                ):
                    continue

                # JSON type (large_string) is compatible with string types
                if pa.types.is_large_string(expected_type) and (
                    pa.types.is_string(actual_type)
                    or pa.types.is_large_string(actual_type)
                ):
                    continue

                # Struct type — expected struct is compatible with actual
                # struct or map representations
                if pa.types.is_struct(expected_type) and (
                    pa.types.is_struct(actual_type)
                    or pa.types.is_map(actual_type)
                    or pa.types.is_list(actual_type)
                ):
                    continue

                logger.warning(
                    "[Validation: %s] Column '%s' type mismatch: expected %s, got %s",
                    self.name,
                    col_name,
                    expected_type,
                    actual_type,
                )

        # Validate JSON well-formedness for declared Json columns
        json_columns: Set[str] = self.validation_config.get("json_columns", set())
        for col_name in json_columns:
            if col_name not in actual_columns:
                continue

            column = table.column(col_name)
            invalid_count = 0
            first_error = None
            first_error_row = None

            for i in range(len(column)):
                value = column[i]
                if not value.is_valid:
                    continue

                str_value = value.as_py()
                if not isinstance(str_value, str):
                    continue

                try:
                    json.loads(str_value)
                except (json.JSONDecodeError, TypeError) as e:
                    invalid_count += 1
                    if first_error is None:
                        first_error = str(e)
                        first_error_row = i

            if invalid_count > 0:
                raise ValueError(
                    f"[Validation: {self.name}] Column '{col_name}' declared as Json "
                    f"contains {invalid_count} invalid JSON value(s). "
                    f"First error at row {first_error_row}: {first_error}"
                )

        logger.debug("[Validation: %s] Schema validation passed.", self.name)


class LocalOutputNode(LocalNode):
    def __init__(
        self,
        name: str,
        feature_view: Union[BatchFeatureView, StreamFeatureView],
        inputs=None,
    ):
        super().__init__(name, inputs=inputs)
        self.feature_view = feature_view

    def execute(self, context: ExecutionContext) -> ArrowTableValue:
        input_table = self.get_single_table(context).data
        context.node_outputs[self.name] = input_table

        if input_table.num_rows == 0:
            return input_table

        if self.feature_view.online:
            online_store = context.online_store

            join_key_to_value_type = {
                entity.name: entity.dtype.to_value_type()
                for entity in self.feature_view.entity_columns
            }

            rows_to_write = _convert_arrow_to_proto(
                input_table, self.feature_view, join_key_to_value_type
            )

            online_store.online_write_batch(
                config=context.repo_config,
                table=self.feature_view,
                data=rows_to_write,
                progress=lambda x: None,
            )

        if self.feature_view.offline:
            offline_store = context.offline_store
            offline_store.offline_write_batch(
                config=context.repo_config,
                feature_view=self.feature_view,
                table=input_table,
                progress=lambda x: None,
            )

        return input_table
