"""
Window function handler for MockDataFrame.

This module handles window function evaluation (row_number, rank, lag, lead, etc.)
following the Single Responsibility Principle.
"""

from typing import List, Dict, Any, Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    pass


class WindowFunctionHandler:
    """Handles window function evaluation (row_number, rank, lag, lead, etc.)."""

    def __init__(self, dataframe: Any):
        """Initialize window function handler.

        Args:
            dataframe: The MockDataFrame instance this handler belongs to
        """
        self.dataframe = dataframe

    def evaluate_window_functions(
        self, data: List[Dict[str, Any]], window_functions: List[Tuple[Any, ...]]
    ) -> List[Dict[str, Any]]:
        """Evaluate window functions across all rows."""
        result_data = data.copy()

        for col_index, window_func in window_functions:
            col_name = window_func.name

            if window_func.function_name == "row_number":
                # For row_number(), we need to handle partitionBy and orderBy
                if hasattr(window_func, "window_spec") and window_func.window_spec:
                    window_spec = window_func.window_spec

                    # Get partition by columns from window spec
                    partition_by_cols = getattr(window_spec, "_partition_by", [])
                    # Get order by columns from window spec
                    order_by_cols = getattr(window_spec, "_order_by", [])

                    if partition_by_cols:
                        # Handle partitioning - group by partition columns
                        partition_groups: Dict[Any, List[int]] = {}
                        for i, row in enumerate(result_data):
                            # Create partition key
                            partition_key = tuple(
                                (
                                    row.get(col.name)
                                    if hasattr(col, "name")
                                    else row.get(str(col))
                                )
                                for col in partition_by_cols
                            )
                            if partition_key not in partition_groups:
                                partition_groups[partition_key] = []
                            partition_groups[partition_key].append(i)

                        # Assign row numbers within each partition
                        for partition_indices in partition_groups.values():
                            if order_by_cols:
                                # Sort within partition by order by columns using corrected ordering logic
                                sorted_partition_indices = (
                                    self._apply_ordering_to_indices(
                                        result_data,
                                        partition_indices,
                                        order_by_cols,
                                    )
                                )
                            else:
                                # No order by - use original order within partition
                                sorted_partition_indices = partition_indices

                            # Assign row numbers starting from 1 within each partition
                            for i, original_index in enumerate(
                                sorted_partition_indices
                            ):
                                result_data[original_index][col_name] = i + 1
                    elif order_by_cols:
                        # No partitioning, just sort by order by columns using corrected ordering logic
                        sorted_indices = self._apply_ordering_to_indices(
                            result_data, list(range(len(result_data))), order_by_cols
                        )

                        # Assign row numbers based on sorted order
                        for i, original_index in enumerate(sorted_indices):
                            result_data[original_index][col_name] = i + 1
                    else:
                        # No partition or order by - just assign sequential row numbers
                        for i in range(len(result_data)):
                            result_data[i][col_name] = i + 1
                else:
                    # No window spec - assign sequential row numbers
                    for i in range(len(result_data)):
                        result_data[i][col_name] = i + 1
            elif window_func.function_name == "lag":
                # Handle lag function - get previous row value
                self._evaluate_lag_lead(
                    result_data, window_func, col_name, is_lead=False
                )
            elif window_func.function_name == "lead":
                # Handle lead function - get next row value
                self._evaluate_lag_lead(
                    result_data, window_func, col_name, is_lead=True
                )
            elif window_func.function_name in ["rank", "dense_rank"]:
                # Handle rank and dense_rank functions
                self._evaluate_rank_functions(result_data, window_func, col_name)
            elif window_func.function_name in [
                "avg",
                "sum",
                "count",
                "countDistinct",
                "max",
                "min",
            ]:
                # Handle aggregate window functions
                self._evaluate_aggregate_window_functions(
                    result_data, window_func, col_name
                )
            else:
                # For other window functions, assign None for now
                for row in result_data:
                    row[col_name] = None

        return result_data

    def _evaluate_lag_lead(
        self, data: List[Dict[str, Any]], window_func: Any, col_name: str, is_lead: bool
    ) -> None:
        """Evaluate lag or lead window function."""
        if not window_func.column_name:
            # No column specified, set to None
            for row in data:
                row[col_name] = None
            return

        # Get offset and default value
        offset = getattr(window_func, "offset", 1)
        default_value = getattr(window_func, "default_value", None)

        # Handle window specification if present
        if hasattr(window_func, "window_spec") and window_func.window_spec:
            window_spec = window_func.window_spec
            partition_by_cols = getattr(window_spec, "_partition_by", [])
            order_by_cols = getattr(window_spec, "_order_by", [])

            if partition_by_cols:
                # Handle partitioning
                partition_groups: Dict[Any, List[int]] = {}
                for i, row in enumerate(data):
                    partition_key = tuple(
                        row.get(col.name) if hasattr(col, "name") else row.get(str(col))
                        for col in partition_by_cols
                    )
                    if partition_key not in partition_groups:
                        partition_groups[partition_key] = []
                    partition_groups[partition_key].append(i)

                # Process each partition
                for partition_indices in partition_groups.values():
                    # Apply ordering to partition indices
                    ordered_indices = self._apply_ordering_to_indices(
                        data, partition_indices, order_by_cols
                    )
                    self._apply_lag_lead_to_partition(
                        data,
                        ordered_indices,
                        window_func.column_name,
                        col_name,
                        offset,
                        default_value,
                        is_lead,
                    )
            else:
                # No partitioning, apply to entire dataset with ordering
                all_indices = list(range(len(data)))
                ordered_indices = self._apply_ordering_to_indices(
                    data, all_indices, order_by_cols
                )
                self._apply_lag_lead_to_partition(
                    data,
                    ordered_indices,
                    window_func.column_name,
                    col_name,
                    offset,
                    default_value,
                    is_lead,
                )
        else:
            # No window spec, apply to entire dataset
            self._apply_lag_lead_to_partition(
                data,
                list(range(len(data))),
                window_func.column_name,
                col_name,
                offset,
                default_value,
                is_lead,
            )

    def _apply_ordering_to_indices(
        self, data: List[Dict[str, Any]], indices: List[int], order_by_cols: List[Any]
    ) -> List[int]:
        """Apply ordering to a list of indices based on order by columns."""
        if not order_by_cols:
            return indices

        def sort_key(idx: int) -> Tuple[Any, ...]:
            row = data[idx]
            key_values = []
            for col in order_by_cols:
                # Handle MockColumnOperation objects (like col("salary").desc())
                if hasattr(col, "column") and hasattr(col.column, "name"):
                    col_name = col.column.name
                elif hasattr(col, "name"):
                    col_name = col.name
                else:
                    col_name = str(col)
                value = row.get(col_name)
                # Handle None values for sorting - put them at the end
                if value is None:
                    key_values.append(float("inf"))  # Sort None values last
                else:
                    key_values.append(value)
            return tuple(key_values)

        # Check if any column has desc operation
        has_desc = any(
            hasattr(col, "operation") and col.operation == "desc"
            for col in order_by_cols
        )

        # Sort indices based on the ordering
        return sorted(indices, key=sort_key, reverse=has_desc)

    def _apply_lag_lead_to_partition(
        self,
        data: List[Dict[str, Any]],
        indices: List[int],
        source_col: str,
        target_col: str,
        offset: int,
        default_value: Any,
        is_lead: bool,
    ) -> None:
        """Apply lag or lead to a specific partition."""
        if is_lead:
            # Lead: get next row value
            for i, idx in enumerate(indices):
                source_idx = i + offset
                if source_idx < len(indices):
                    actual_idx = indices[source_idx]
                    data[idx][target_col] = data[actual_idx].get(source_col)
                else:
                    data[idx][target_col] = default_value
        else:
            # Lag: get previous row value
            for i, idx in enumerate(indices):
                source_idx = i - offset
                if source_idx >= 0:
                    actual_idx = indices[source_idx]
                    data[idx][target_col] = data[actual_idx].get(source_col)
                else:
                    data[idx][target_col] = default_value

    def _evaluate_rank_functions(
        self, data: List[Dict[str, Any]], window_func: Any, col_name: str
    ) -> None:
        """Evaluate rank or dense_rank window function."""
        is_dense = window_func.function_name == "dense_rank"

        # Handle window specification if present
        if hasattr(window_func, "window_spec") and window_func.window_spec:
            window_spec = window_func.window_spec
            partition_by_cols = getattr(window_spec, "_partition_by", [])
            order_by_cols = getattr(window_spec, "_order_by", [])

            if partition_by_cols:
                # Handle partitioning
                partition_groups: Dict[Any, List[int]] = {}
                for i, row in enumerate(data):
                    partition_key = tuple(
                        row.get(col.name) if hasattr(col, "name") else row.get(str(col))
                        for col in partition_by_cols
                    )
                    if partition_key not in partition_groups:
                        partition_groups[partition_key] = []
                    partition_groups[partition_key].append(i)

                # Process each partition
                for partition_indices in partition_groups.values():
                    self._apply_rank_to_partition(
                        data, partition_indices, order_by_cols, col_name, is_dense
                    )
            else:
                # No partitioning, apply to entire dataset
                self._apply_rank_to_partition(
                    data, list(range(len(data))), order_by_cols, col_name, is_dense
                )
        else:
            # No window spec, assign ranks based on original order
            for i in range(len(data)):
                data[i][col_name] = i + 1

    def _apply_rank_to_partition(
        self,
        data: List[Dict[str, Any]],
        indices: List[int],
        order_by_cols: List[Any],
        col_name: str,
        is_dense: bool,
    ) -> None:
        """Apply rank or dense_rank to a specific partition."""
        if not order_by_cols:
            # No order by, assign ranks based on original order
            for i, idx in enumerate(indices):
                data[idx][col_name] = i + 1
            return

        # Sort partition by order by columns using the corrected ordering logic
        sorted_indices = self._apply_ordering_to_indices(data, indices, order_by_cols)

        # Assign ranks in sorted order
        if is_dense:
            # Dense rank: consecutive ranks without gaps
            current_rank = 1
            previous_values = None

            for i, idx in enumerate(sorted_indices):
                row = data[idx]
                current_values = []
                for col in order_by_cols:
                    # Handle MockColumnOperation objects (like col("salary").desc())
                    if hasattr(col, "column") and hasattr(col.column, "name"):
                        order_col_name = col.column.name
                    elif hasattr(col, "name"):
                        order_col_name = col.name
                    else:
                        order_col_name = str(col)
                    value = row.get(order_col_name)
                    current_values.append(value)

                if previous_values is not None and current_values != previous_values:
                    current_rank += 1

                data[idx][col_name] = current_rank
                previous_values = current_values
        else:
            # Regular rank: ranks with gaps for ties
            current_rank = 1

            for i, idx in enumerate(sorted_indices):
                if i > 0:
                    prev_idx = sorted_indices[i - 1]
                    # Check if current and previous rows have different values
                    row = data[idx]
                    prev_row = data[prev_idx]

                    current_values = []
                    prev_values = []
                    for col in order_by_cols:
                        # Handle MockColumnOperation objects (like col("salary").desc())
                        if hasattr(col, "column") and hasattr(col.column, "name"):
                            order_col_name = col.column.name
                        elif hasattr(col, "name"):
                            order_col_name = col.name
                        else:
                            order_col_name = str(col)
                        current_values.append(row.get(order_col_name))
                        prev_values.append(prev_row.get(order_col_name))

                    if current_values != prev_values:
                        current_rank = i + 1
                else:
                    current_rank = 1

                data[idx][col_name] = current_rank

    def _evaluate_aggregate_window_functions(
        self, data: List[Dict[str, Any]], window_func: Any, col_name: str
    ) -> None:
        """Evaluate aggregate window functions like avg, sum, count, etc."""
        if not window_func.column_name and window_func.function_name not in ["count"]:
            # No column specified for functions that need it
            for row in data:
                row[col_name] = None
            return

        # Handle window specification if present
        if hasattr(window_func, "window_spec") and window_func.window_spec:
            window_spec = window_func.window_spec
            partition_by_cols = getattr(window_spec, "_partition_by", [])
            order_by_cols = getattr(window_spec, "_order_by", [])

            if partition_by_cols:
                # Handle partitioning
                partition_groups: Dict[Any, List[int]] = {}
                for i, row in enumerate(data):
                    partition_key = tuple(
                        row.get(col.name) if hasattr(col, "name") else row.get(str(col))
                        for col in partition_by_cols
                    )
                    if partition_key not in partition_groups:
                        partition_groups[partition_key] = []
                    partition_groups[partition_key].append(i)

                # Process each partition
                for partition_indices in partition_groups.values():
                    # Apply ordering to partition indices
                    ordered_indices = self._apply_ordering_to_indices(
                        data, partition_indices, order_by_cols
                    )
                    self._apply_aggregate_to_partition(
                        data, ordered_indices, window_func, col_name
                    )
            else:
                # No partitioning, apply to entire dataset with ordering
                all_indices = list(range(len(data)))
                ordered_indices = self._apply_ordering_to_indices(
                    data, all_indices, order_by_cols
                )
                self._apply_aggregate_to_partition(
                    data, ordered_indices, window_func, col_name
                )
        else:
            # No window spec, apply to entire dataset
            all_indices = list(range(len(data)))
            self._apply_aggregate_to_partition(data, all_indices, window_func, col_name)

    def _apply_aggregate_to_partition(
        self,
        data: List[Dict[str, Any]],
        indices: List[int],
        window_func: Any,
        col_name: str,
    ) -> None:
        """Apply aggregate function to a specific partition."""
        if not indices:
            return

        source_col = window_func.column_name
        func_name = window_func.function_name

        # Get window boundaries if specified
        rows_between = (
            getattr(window_func.window_spec, "_rows_between", None)
            if hasattr(window_func, "window_spec") and window_func.window_spec
            else None
        )

        for i, idx in enumerate(indices):
            # Determine the window for this row
            if rows_between:
                start_offset, end_offset = rows_between
                window_start = max(0, i + start_offset)
                window_end = min(len(indices), i + end_offset + 1)
            else:
                # Default: all rows up to current row
                window_start = 0
                window_end = i + 1

            # Get values in the window
            window_values = []
            for j in range(window_start, window_end):
                if j < len(indices):
                    row_idx = indices[j]
                    if source_col:
                        value = data[row_idx].get(source_col)
                        if value is not None:
                            window_values.append(value)
                    else:
                        # For count(*) - count all rows
                        window_values.append(1)

            # Apply aggregate function
            if func_name == "avg":
                data[idx][col_name] = (
                    sum(window_values) / len(window_values) if window_values else None
                )
            elif func_name == "sum":
                data[idx][col_name] = sum(window_values) if window_values else None
            elif func_name == "count":
                data[idx][col_name] = len(window_values)
            elif func_name == "countDistinct":
                data[idx][col_name] = len(set(window_values))
            elif func_name == "max":
                data[idx][col_name] = max(window_values) if window_values else None
            elif func_name == "min":
                data[idx][col_name] = min(window_values) if window_values else None
