"""Delta Lake Maintenance Operations Implementation.

This module implements DataFrame maintenance operations for Delta Lake,
providing full ACID compliance for TPC-H RF1/RF2 and TPC-DS maintenance testing.

Delta Lake supports:
- INSERT: Append new data with transaction guarantees
- DELETE: Row-level deletes with predicate pushdown
- UPDATE: Row-level updates
- MERGE: Upsert operations (insert or update)

Delta Lake provides:
- ACID transactions with optimistic concurrency
- Time travel (query historical versions)
- Schema enforcement and evolution
- Partition pruning for efficient queries

Copyright 2026 Joe Harris / BenchBox Project

Licensed under the MIT License. See LICENSE file in the project root for details.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any

try:
    import deltalake
    from deltalake import DeltaTable, write_deltalake

    DELTA_AVAILABLE = True
except ImportError:
    deltalake = None  # type: ignore[assignment]
    DeltaTable = None  # type: ignore[assignment, misc]
    write_deltalake = None  # type: ignore[assignment]
    DELTA_AVAILABLE = False

try:
    import pyarrow as pa

    PYARROW_AVAILABLE = True
except ImportError:
    pa = None  # type: ignore[assignment]
    PYARROW_AVAILABLE = False

from benchbox.core.dataframe.maintenance_interface import (
    DELTA_LAKE_CAPABILITIES,
    BaseDataFrameMaintenanceOperations,
    DataFrameMaintenanceCapabilities,
)

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


class DeltaLakeMaintenanceOperations(BaseDataFrameMaintenanceOperations):
    """Delta Lake maintenance operations implementation.

    Implements full ACID maintenance operations using delta-rs:
    - INSERT: Append with transaction guarantees
    - DELETE: Row-level deletes using predicates
    - UPDATE: Row-level updates
    - MERGE: Upsert operations

    Delta Lake provides snapshot isolation and optimistic concurrency,
    making it suitable for TPC-H and TPC-DS maintenance tests.

    Example:
        ops = DeltaLakeMaintenanceOperations()

        # Insert new rows (transactional)
        result = ops.insert_rows(
            table_path="/data/orders",
            dataframe=new_orders_df,
            mode="append"
        )

        # Delete rows (row-level)
        result = ops.delete_rows(
            table_path="/data/orders",
            condition="order_date < '2020-01-01'"
        )

        # Update rows
        result = ops.update_rows(
            table_path="/data/orders",
            condition="status = 'pending'",
            updates={"status": "'cancelled'"}
        )

    Note:
        Requires delta-rs: pip install deltalake
    """

    def __init__(self, working_dir: str | Path | None = None) -> None:
        """Initialize Delta Lake maintenance operations.

        Args:
            working_dir: Optional working directory for temporary files

        Raises:
            ImportError: If delta-rs is not installed
        """
        super().__init__()

        if not DELTA_AVAILABLE:
            raise ImportError(
                "delta-rs is not installed. Install with: pip install deltalake\n"
                "For TPC-H/TPC-DS maintenance tests with Delta Lake, install:\n"
                "pip install 'benchbox[delta]'"
            )

        if not PYARROW_AVAILABLE:
            raise ImportError(
                "PyArrow is not installed. Install with: pip install pyarrow\n"
                "PyArrow is required for Delta Lake operations."
            )

        self.working_dir = Path(working_dir) if working_dir else None
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

    def _get_capabilities(self) -> DataFrameMaintenanceCapabilities:
        """Return Delta Lake maintenance capabilities.

        Returns:
            DELTA_LAKE_CAPABILITIES (full ACID support)
        """
        return DELTA_LAKE_CAPABILITIES

    # Note: _convert_to_arrow() is inherited from BaseDataFrameMaintenanceOperations

    def _do_insert(
        self,
        table_path: Path | str,
        dataframe: Any,
        partition_columns: list[str] | None,
        mode: str,
    ) -> int:
        """Insert rows using Delta Lake write.

        Args:
            table_path: Path to the Delta table
            dataframe: DataFrame containing rows to insert
            partition_columns: Columns to partition by
            mode: Write mode ("append" or "overwrite")

        Returns:
            Number of rows inserted
        """
        table_path = str(table_path)

        # Convert to PyArrow
        arrow_table = self._convert_to_arrow(dataframe)
        row_count = arrow_table.num_rows

        if row_count == 0:
            self.logger.info("No rows to insert")
            return 0

        # Map mode to Delta Lake mode
        delta_mode = "append" if mode == "append" else "overwrite"

        # Write to Delta table
        write_deltalake(
            table_path,
            arrow_table,
            mode=delta_mode,
            partition_by=partition_columns,
        )

        self.logger.info(f"Inserted {row_count} rows to Delta table at {table_path}")
        return row_count

    def _do_delete(
        self,
        table_path: Path | str,
        condition: str | Any,
    ) -> int:
        """Delete rows using Delta Lake predicate.

        Args:
            table_path: Path to the Delta table
            condition: SQL-like delete predicate (e.g., "id > 100")

        Returns:
            Number of rows deleted
        """
        table_path = str(table_path)

        # Open existing Delta table
        try:
            dt = DeltaTable(table_path)
        except Exception as e:
            self.logger.warning(f"Could not open Delta table at {table_path}: {e}")
            return 0

        # Get row count before delete
        rows_before = dt.to_pyarrow_table().num_rows

        # Execute delete with predicate
        dt.delete(predicate=str(condition))

        # Get row count after delete
        rows_after = dt.to_pyarrow_table().num_rows
        rows_deleted = rows_before - rows_after

        self.logger.info(f"Deleted {rows_deleted} rows from Delta table at {table_path}")
        return rows_deleted

    def _do_update(
        self,
        table_path: Path | str,
        condition: str | Any,
        updates: dict[str, Any],
    ) -> int:
        """Update rows using Delta Lake.

        Args:
            table_path: Path to the Delta table
            condition: SQL-like update predicate
            updates: Column name to new value expression mapping

        Returns:
            Number of rows updated
        """
        table_path = str(table_path)

        # Open existing Delta table
        try:
            dt = DeltaTable(table_path)
        except Exception as e:
            raise RuntimeError(f"Could not open Delta table at {table_path}: {e}") from e

        # Count matching rows before update
        # Note: This is approximate since concurrent changes could occur
        full_table = dt.to_pyarrow_table()
        rows_before = full_table.num_rows

        # Execute update
        dt.update(
            predicate=str(condition),
            updates=updates,
        )

        # Estimate rows updated (Delta Lake doesn't return this directly)
        # We count rows matching the predicate before update
        # This is a heuristic - exact count would require additional query
        self.logger.info(f"Updated rows in Delta table at {table_path} where {condition}")

        # Return estimated count (rows that matched predicate)
        # Note: delta-rs update() doesn't return affected row count
        # We estimate based on table size change or return a placeholder
        return rows_before  # Placeholder - actual count not available from delta-rs

    def _do_merge(
        self,
        table_path: Path | str,
        source_dataframe: Any,
        merge_condition: str | Any,
        when_matched: dict[str, Any] | None,
        when_not_matched: dict[str, Any] | None,
    ) -> int:
        """Merge (upsert) rows using Delta Lake.

        Args:
            table_path: Path to the Delta table
            source_dataframe: Source DataFrame
            merge_condition: Join condition for matching
            when_matched: Updates to apply when matched
            when_not_matched: Inserts when not matched

        Returns:
            Number of rows affected
        """
        table_path = str(table_path)

        # Convert source to PyArrow
        source_arrow = self._convert_to_arrow(source_dataframe)
        source_rows = source_arrow.num_rows

        # Open existing Delta table
        try:
            dt = DeltaTable(table_path)
        except Exception as e:
            raise RuntimeError(f"Could not open Delta table at {table_path}: {e}") from e

        # Build merge operation
        merge_builder = dt.merge(
            source=source_arrow,
            predicate=str(merge_condition),
            source_alias="source",
            target_alias="target",
        )

        # Add when_matched clause
        if when_matched:
            merge_builder = merge_builder.when_matched_update(updates=when_matched)

        # Add when_not_matched clause
        if when_not_matched:
            merge_builder = merge_builder.when_not_matched_insert(updates=when_not_matched)

        # Execute merge
        merge_result = merge_builder.execute()

        # Extract metrics from merge result
        rows_affected = source_rows  # Estimate based on source size
        if merge_result and isinstance(merge_result, dict):
            # delta-rs returns metrics dict
            rows_affected = merge_result.get("num_target_rows_updated", 0) + merge_result.get(
                "num_target_rows_inserted", 0
            )

        self.logger.info(f"Merged {rows_affected} rows into Delta table at {table_path}")
        return rows_affected


def get_delta_lake_maintenance_operations(
    working_dir: str | Path | None = None,
) -> DeltaLakeMaintenanceOperations | None:
    """Get Delta Lake maintenance operations if delta-rs is available.

    Args:
        working_dir: Optional working directory

    Returns:
        DeltaLakeMaintenanceOperations if delta-rs is available, None otherwise
    """
    if not DELTA_AVAILABLE:
        logger.debug("Delta Lake maintenance not available (deltalake not installed)")
        return None

    if not PYARROW_AVAILABLE:
        logger.debug("Delta Lake maintenance not available (pyarrow not installed)")
        return None

    return DeltaLakeMaintenanceOperations(working_dir=working_dir)
