"""
Module to read an iceberg table into a Ray Dataset, by using the Ray Datasource API.
"""

import heapq
import itertools
import logging
from functools import partial
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import pyarrow as pa
from packaging import version

from ray.data._internal.planner.plan_expression.expression_visitors import _ExprVisitor
from ray.data._internal.util import _check_import
from ray.data.block import Block, BlockMetadata
from ray.data.datasource.datasource import Datasource, ReadTask
from ray.data.expressions import (
    AliasExpr,
    BinaryExpr,
    ColumnExpr,
    DownloadExpr,
    LiteralExpr,
    MonotonicallyIncreasingIdExpr,
    Operation,
    RandomExpr,
    StarExpr,
    UDFExpr,
    UnaryExpr,
    UUIDExpr,
)
from ray.util import log_once
from ray.util.annotations import DeveloperAPI

try:
    from pyiceberg.expressions import (
        And,
        EqualTo,
        GreaterThan,
        GreaterThanOrEqual,
        In,
        IsNull,
        LessThan,
        LessThanOrEqual,
        Literal,
        Not,
        NotEqualTo,
        NotIn,
        NotNull,
        Or,
        Reference,
        UnboundTerm,
        literal,
    )

    RAY_DATA_OPERATION_TO_ICEBERG = {
        Operation.EQ: EqualTo,
        Operation.NE: NotEqualTo,
        Operation.GT: GreaterThan,
        Operation.GE: GreaterThanOrEqual,
        Operation.LT: LessThan,
        Operation.LE: LessThanOrEqual,
        Operation.AND: And,
        Operation.OR: Or,
        Operation.IN: In,
        Operation.NOT_IN: NotIn,
        Operation.IS_NULL: IsNull,
        Operation.IS_NOT_NULL: NotNull,
        Operation.NOT: Not,
    }
except ImportError:
    log_once("pyiceberg.expressions not found. Please install pyiceberg >= 0.9.0")

if TYPE_CHECKING:
    from pyiceberg.catalog import Catalog
    from pyiceberg.expressions import BooleanExpression
    from pyiceberg.io import FileIO
    from pyiceberg.manifest import DataFile
    from pyiceberg.schema import Schema
    from pyiceberg.table import DataScan, FileScanTask, Table
    from pyiceberg.table.metadata import TableMetadata

    from ray.data.context import DataContext

logger = logging.getLogger(__name__)


class _IcebergExpressionVisitor(
    _ExprVisitor["BooleanExpression | UnboundTerm | Literal"]
):
    """
    Visitor that converts Ray Data expressions to PyIceberg expressions.

    This enables Ray Data users to write filters using the familiar col() syntax
    while leveraging Iceberg's native filtering capabilities.

    Example:
        >>> from ray.data.expressions import col
        >>> ray_expr = (col("date") >= "2024-01-01") & (col("status") == "active")
        >>> iceberg_expr = _IcebergExpressionVisitor().visit(ray_expr)
        >>> # iceberg_expr can now be used with PyIceberg's filter APIs
    """

    def visit_column(self, expr: "ColumnExpr") -> "UnboundTerm":
        """Convert a column reference to an Iceberg reference."""
        return Reference(expr.name)

    def visit_literal(self, expr: "LiteralExpr") -> "Literal":
        """Convert a literal value to an Iceberg literal."""
        return literal(expr.value)

    def visit_binary(self, expr: "BinaryExpr") -> "BooleanExpression":
        """Convert a binary operation to an Iceberg expression."""
        # Handle IN/NOT_IN specially since they don't visit the right operand
        # (the right operand is a list literal that can't be converted)
        if expr.op in (Operation.IN, Operation.NOT_IN):
            left = self.visit(expr.left)
            if not isinstance(expr.right, LiteralExpr):
                raise ValueError(
                    f"{expr.op.name} operation requires right operand to be a literal list, "
                    f"got {type(expr.right).__name__}"
                )
            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](left, expr.right.value)

        # For all other operations, visit both operands
        left = self.visit(expr.left)
        right = self.visit(expr.right)

        if expr.op in RAY_DATA_OPERATION_TO_ICEBERG:
            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](left, right)
        else:
            # Arithmetic operations are not supported in filter expressions
            raise ValueError(
                f"Unsupported binary operation for Iceberg filters: {expr.op}. "
                f"Iceberg filters support: {RAY_DATA_OPERATION_TO_ICEBERG.keys()}. "
                f"Arithmetic operations (ADD, SUB, MUL, DIV) cannot be used in filters."
            )

    def visit_unary(self, expr: "UnaryExpr") -> "BooleanExpression":
        """Convert a unary operation to an Iceberg expression."""
        operand = self.visit(expr.operand)

        if expr.op in RAY_DATA_OPERATION_TO_ICEBERG:
            return RAY_DATA_OPERATION_TO_ICEBERG[expr.op](operand)
        else:
            raise ValueError(
                f"Unsupported unary operation for Iceberg: {expr.op}. "
                f"Supported operations: {RAY_DATA_OPERATION_TO_ICEBERG.keys()}"
            )

    def visit_alias(
        self, expr: "AliasExpr"
    ) -> "BooleanExpression | UnboundTerm | Literal":
        """Convert an aliased expression (just unwrap the alias)."""
        return self.visit(expr.expr)

    def visit_udf(self, expr: "UDFExpr") -> "BooleanExpression | UnboundTerm | Literal":
        """UDF expressions cannot be converted to Iceberg expressions."""
        raise TypeError(
            "UDF expressions cannot be converted to Iceberg expressions. "
            "Iceberg filters must use simple column comparisons and boolean operations."
        )

    def visit_download(
        self, expr: "DownloadExpr"
    ) -> "BooleanExpression | UnboundTerm | Literal":
        """Download expressions cannot be converted to Iceberg expressions."""
        raise TypeError(
            "Download expressions cannot be converted to Iceberg expressions."
        )

    def visit_star(
        self, expr: "StarExpr"
    ) -> "BooleanExpression | UnboundTerm | Literal":
        """Star expressions cannot be converted to Iceberg expressions."""
        raise TypeError(
            "Star expressions cannot be converted to Iceberg filter expressions."
        )

    def visit_monotonically_increasing_id(
        self, expr: "MonotonicallyIncreasingIdExpr"
    ) -> "BooleanExpression | UnboundTerm | Literal":
        """Monotonically increasing ID expressions cannot be converted to Iceberg expressions."""
        raise TypeError(
            "monotonically_increasing_id expressions cannot be converted to Iceberg filter expressions."
        )

    def visit_random(
        self, expr: "RandomExpr"
    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
        """Random expressions cannot be converted to Iceberg expressions."""
        raise TypeError(
            "Random expressions cannot be converted to Iceberg filter expressions."
        )

    def visit_uuid(
        self, expr: "UUIDExpr"
    ) -> "BooleanExpression | UnboundTerm[Any] | Literal[Any]":
        """UUID expressions cannot be converted to Iceberg expressions."""
        raise TypeError(
            "UUID expressions cannot be converted to Iceberg filter expressions."
        )


def _get_read_task(
    tasks: Iterable["FileScanTask"],
    table_io: "FileIO",
    table_metadata: "TableMetadata",
    row_filter: "BooleanExpression",
    case_sensitive: bool,
    limit: Optional[int],
    schema: "Schema",
) -> Iterable[Block]:
    # Determine the PyIceberg version to handle backward compatibility
    import pyiceberg

    def _generate_tables() -> Iterable[pa.Table]:
        if version.parse(pyiceberg.__version__) >= version.parse("0.9.0"):
            # Modern implementation using ArrowScan (PyIceberg 0.9.0+)
            from pyiceberg.io.pyarrow import ArrowScan

            # Initialize scanner with Iceberg metadata and query parameters
            scanner = ArrowScan(
                table_metadata=table_metadata,
                io=table_io,
                row_filter=row_filter,
                projected_schema=schema,
                case_sensitive=case_sensitive,
                limit=limit,
            )

            # Convert scanned data to Arrow Table format
            result_table = scanner.to_table(tasks=tasks)

            # Stream results as RecordBatches for memory efficiency
            for batch in result_table.to_batches():
                yield pa.Table.from_batches([batch])

        else:
            # Legacy implementation using project_table (PyIceberg <0.9.0)
            from pyiceberg.io import pyarrow as pyi_pa_io

            # Use the PyIceberg API to read only a single task (specifically, a
            # FileScanTask) - note that this is not as simple as reading a single
            # parquet file, as there might be delete files, etc. associated, so we
            # must use the PyIceberg API for the projection.
            table = pyi_pa_io.project_table(
                tasks=tasks,
                table_metadata=table_metadata,
                io=table_io,
                row_filter=row_filter,
                projected_schema=schema,
                case_sensitive=case_sensitive,
                limit=limit,
            )
            yield table

    yield from _generate_tables()


@DeveloperAPI
class IcebergDatasource(Datasource):
    """
    Iceberg datasource to read Iceberg tables into a Ray Dataset. This module heavily
    uses PyIceberg to read iceberg tables. All the routines in this class override
    `ray.data.Datasource`.
    """

    def __init__(
        self,
        table_identifier: str,
        row_filter: Union[str, "BooleanExpression"] = None,
        selected_fields: Tuple[str, ...] = ("*",),
        snapshot_id: Optional[int] = None,
        scan_kwargs: Optional[Dict[str, Any]] = None,
        catalog_kwargs: Optional[Dict[str, Any]] = None,
    ):
        """
        Initialize an IcebergDatasource.

        Args:
            table_identifier: Fully qualified table identifier (i.e.,
                "db_name.table_name")
            row_filter: A PyIceberg BooleanExpression to use to filter the data *prior*
                 to reading
            selected_fields: Which columns from the data to read, passed directly to
                PyIceberg's load functions
            snapshot_id: Optional snapshot ID for the Iceberg table
            scan_kwargs: Optional arguments to pass to PyIceberg's Table.scan()
                function
            catalog_kwargs: Optional arguments to use when setting up the Iceberg
                catalog
        """
        # Initialize parent class to set up predicate pushdown mixin
        super().__init__()

        _check_import(self, module="pyiceberg", package="pyiceberg")
        from pyiceberg.expressions import AlwaysTrue

        self._scan_kwargs = scan_kwargs if scan_kwargs is not None else {}
        self._catalog_kwargs = catalog_kwargs if catalog_kwargs is not None else {}

        if "name" in self._catalog_kwargs:
            self._catalog_name = self._catalog_kwargs.pop("name")
        else:
            self._catalog_name = "default"

        self.table_identifier = table_identifier

        self._row_filter = row_filter if row_filter is not None else AlwaysTrue()
        # Convert selected_fields to projection_map (identity mapping if specified)
        # Note: Empty tuple () means no columns, None/"*" means all columns
        if selected_fields is None or selected_fields == ("*",):
            self._projection_map = None
        else:
            self._projection_map = {col: col for col in selected_fields}

        if snapshot_id:
            self._scan_kwargs["snapshot_id"] = snapshot_id

        self._plan_files = None
        self._table = None

    def _get_catalog(self) -> "Catalog":
        from pyiceberg import catalog

        return catalog.load_catalog(self._catalog_name, **self._catalog_kwargs)

    @property
    def table(self) -> "Table":
        """
        Return the table reference from the catalog
        """
        if self._table is None:
            catalog = self._get_catalog()
            self._table = catalog.load_table(self.table_identifier)
        return self._table

    @property
    def plan_files(self) -> List["FileScanTask"]:
        """
        Return the plan files specified by this query
        """
        # Calculate and cache the plan_files if they don't already exist
        if self._plan_files is None:
            data_scan = self._get_data_scan()
            self._plan_files = data_scan.plan_files()

        return self._plan_files

    def _get_combined_filter(self) -> "BooleanExpression":
        """Get the combined filter including both row_filter and pushed-down predicates."""
        combined_filter = self._row_filter

        if self._predicate_expr is not None:
            # Convert Ray Data expression to PyIceberg expression using internal visitor
            visitor = _IcebergExpressionVisitor()
            iceberg_filter = visitor.visit(self._predicate_expr)

            # Combine with existing row_filter using AND
            from pyiceberg.expressions import AlwaysTrue, And

            if not isinstance(combined_filter, AlwaysTrue):
                combined_filter = And(combined_filter, iceberg_filter)
            else:
                combined_filter = iceberg_filter

        return combined_filter

    def _get_data_scan(self) -> "DataScan":
        # Get the combined filter
        combined_filter = self._get_combined_filter()

        # Convert back to tuple for PyIceberg API (None -> ("*",))
        data_columns = self._get_data_columns()
        selected_fields = ("*",) if data_columns is None else tuple(data_columns)

        data_scan = self.table.scan(
            row_filter=combined_filter,
            selected_fields=selected_fields,
            **self._scan_kwargs,
        )

        return data_scan

    def estimate_inmemory_data_size(self) -> Optional[int]:
        # Approximate the size by using the plan files - this will not
        # incorporate the deletes, but that's a reasonable approximation
        # task
        return sum(task.file.file_size_in_bytes for task in self.plan_files)

    def supports_predicate_pushdown(self) -> bool:
        """Returns True to indicate this datasource supports predicate pushdown."""
        return True

    def supports_projection_pushdown(self) -> bool:
        """Returns True to indicate this datasource supports projection pushdown."""
        return True

    @staticmethod
    def _distribute_tasks_into_equal_chunks(
        plan_files: Iterable["FileScanTask"], n_chunks: int
    ) -> List[List["FileScanTask"]]:
        """
        Implement a greedy knapsack algorithm to distribute the files in the scan
        across tasks, based on their file size, as evenly as possible
        """
        chunks = [list() for _ in range(n_chunks)]

        chunk_sizes = [(0, chunk_id) for chunk_id in range(n_chunks)]
        heapq.heapify(chunk_sizes)

        # From largest to smallest, add the plan files to the smallest chunk one at a
        # time
        for plan_file in sorted(
            plan_files, key=lambda f: f.file.file_size_in_bytes, reverse=True
        ):
            smallest_chunk = heapq.heappop(chunk_sizes)
            chunks[smallest_chunk[1]].append(plan_file)
            heapq.heappush(
                chunk_sizes,
                (
                    smallest_chunk[0] + plan_file.file.file_size_in_bytes,
                    smallest_chunk[1],
                ),
            )

        return chunks

    def get_read_tasks(
        self,
        parallelism: int,
        per_task_row_limit: Optional[int] = None,
        data_context: Optional["DataContext"] = None,
    ) -> List[ReadTask]:
        from pyiceberg.io import pyarrow as pyi_pa_io
        from pyiceberg.manifest import DataFileContent

        # Get the PyIceberg scan
        data_scan = self._get_data_scan()
        # Get the plan files in this query
        plan_files = self.plan_files

        # Get the projected schema for this scan, given all the row filters,
        # snapshot ID, etc.
        projected_schema = data_scan.projection()
        # Get the arrow schema, to set in the metadata
        pya_schema = pyi_pa_io.schema_to_pyarrow(projected_schema)

        # Set the n_chunks to the min of the number of plan files and the actual
        # requested n_chunks, so that there are no empty tasks
        if parallelism > len(list(plan_files)):
            parallelism = len(list(plan_files))
            logger.warning(
                f"Reducing the parallelism to {parallelism}, as that is the number of files"
            )

        # Get required properties for reading tasks - table IO, table metadata,
        # row filter, case sensitivity,limit and projected schema to pass
        # them directly to `_get_read_task` to avoid capture of `self` reference
        # within the closure carrying substantial overhead invoking these tasks
        #
        # See https://github.com/ray-project/ray/issues/49107 for more context
        table_io = self.table.io
        table_metadata = self.table.metadata
        row_filter = self._get_combined_filter()
        case_sensitive = self._scan_kwargs.get("case_sensitive", True)
        limit = self._scan_kwargs.get("limit")

        get_read_task = partial(
            _get_read_task,
            table_io=table_io,
            table_metadata=table_metadata,
            row_filter=row_filter,
            case_sensitive=case_sensitive,
            limit=limit,
            schema=projected_schema,
        )

        read_tasks = []
        # Chunk the plan files based on the requested parallelism
        for chunk_tasks in IcebergDatasource._distribute_tasks_into_equal_chunks(
            plan_files, parallelism
        ):
            unique_deletes: Set[DataFile] = set(
                itertools.chain.from_iterable(
                    [task.delete_files for task in chunk_tasks]
                )
            )
            # Get a rough estimate of the number of deletes by just looking at
            # position deletes. Equality deletes are harder to estimate, as they
            # can delete multiple rows.
            position_delete_count = sum(
                delete.record_count
                for delete in unique_deletes
                if delete.content == DataFileContent.POSITION_DELETES
            )
            metadata = BlockMetadata(
                num_rows=sum(task.file.record_count for task in chunk_tasks)
                - position_delete_count,
                size_bytes=sum(task.file.file_size_in_bytes for task in chunk_tasks),
                input_files=[task.file.file_path for task in chunk_tasks],
                exec_stats=None,
            )
            read_tasks.append(
                ReadTask(
                    read_fn=lambda tasks=chunk_tasks: get_read_task(tasks),
                    metadata=metadata,
                    schema=pya_schema,
                    per_task_row_limit=per_task_row_limit,
                )
            )

        return read_tasks