"""Apache Hudi Maintenance Operations Implementation.

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

Hudi 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)

Hudi provides:
- ACID transactions with snapshot isolation
- Time travel (query historical commits)
- Schema evolution
- Copy-on-Write (COW) and Merge-on-Read (MOR) table types
- Record-level change tracking via record key

IMPORTANT: Unlike Delta Lake (delta-rs) and Iceberg (pyiceberg), Hudi has NO pure
Python maintenance library. All maintenance operations MUST use PySpark SQL,
which requires an active SparkSession with the hudi-spark-bundle.

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 Any

# Check for PySpark availability
try:
    from pyspark.sql import SparkSession

    PYSPARK_AVAILABLE = True
except ImportError:
    SparkSession = None  # type: ignore[assignment, misc]
    PYSPARK_AVAILABLE = False

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

logger = logging.getLogger(__name__)


class HudiMaintenanceOperations(BaseDataFrameMaintenanceOperations):
    """Apache Hudi maintenance operations implementation.

    Implements full ACID maintenance operations using PySpark SQL with Hudi:
    - INSERT: INSERT INTO table SELECT ... or DataFrameWriter
    - DELETE: DELETE FROM table WHERE condition
    - UPDATE: UPDATE table SET ... WHERE condition
    - MERGE: MERGE INTO table USING source ON condition WHEN MATCHED/NOT MATCHED

    Hudi provides snapshot isolation and ACID guarantees with record-level
    change tracking, making it suitable for TPC-H and TPC-DS maintenance tests.

    Key Hudi Concepts:
    - Record Key: Unique identifier for each record (required)
    - Precombine Field: Ordering field to resolve duplicates (typically timestamp)
    - Table Type: COPY_ON_WRITE (faster reads) or MERGE_ON_READ (faster writes)

    Example:
        from pyspark.sql import SparkSession

        spark = SparkSession.builder \\
            .appName("benchbox") \\
            .config("spark.jars.packages", "org.apache.hudi:hudi-spark3.5-bundle_2.12:0.14.0") \\
            .getOrCreate()

        ops = HudiMaintenanceOperations(
            spark_session=spark,
            record_key="order_key",
            precombine_field="order_date",
        )

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

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

    Note:
        Requires PySpark with hudi-spark-bundle:
        pip install pyspark
        Add to Spark config: spark.jars.packages=org.apache.hudi:hudi-spark3.5-bundle_2.12:0.14.0
    """

    def __init__(
        self,
        spark_session: Any,
        record_key: str | None = None,
        precombine_field: str | None = None,
        table_type: str = "COPY_ON_WRITE",
        working_dir: str | Path | None = None,
    ) -> None:
        """Initialize Hudi maintenance operations.

        Args:
            spark_session: Active SparkSession with Hudi bundle loaded
            record_key: Field name for record key (required for writes).
                If not provided, must be set per-table via table properties.
            precombine_field: Field name for ordering during updates.
                If not provided, uses a default constant.
            table_type: Hudi table type - "COPY_ON_WRITE" (default) or "MERGE_ON_READ".
                COW is better for analytics (faster reads), MOR for streaming (faster writes).
            working_dir: Optional working directory for temporary files

        Raises:
            ImportError: If PySpark is not installed
            ValueError: If spark_session is None
        """
        super().__init__()

        if not PYSPARK_AVAILABLE:
            raise ImportError(
                "PySpark is not installed. Install with: pip install pyspark\n"
                "For Hudi support, also configure Spark with:\n"
                "  spark.jars.packages=org.apache.hudi:hudi-spark3.5-bundle_2.12:0.14.0"
            )

        if spark_session is None:
            raise ValueError(
                "spark_session is required. Create one with:\n"
                "  spark = SparkSession.builder \\\n"
                "      .appName('benchbox') \\\n"
                "      .config('spark.jars.packages', 'org.apache.hudi:hudi-spark3.5-bundle_2.12:0.14.0') \\\n"
                "      .getOrCreate()"
            )

        self.spark = spark_session
        self.record_key = record_key
        self.precombine_field = precombine_field
        self.table_type = table_type
        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 Hudi maintenance capabilities.

        Returns:
            HUDI_CAPABILITIES (full ACID support via PySpark SQL)
        """
        return HUDI_CAPABILITIES

    def _convert_to_spark_df(self, dataframe: Any) -> Any:
        """Convert various DataFrame types to Spark DataFrame.

        Args:
            dataframe: Input DataFrame (PySpark, Pandas, Polars, or PyArrow)

        Returns:
            Spark DataFrame

        Raises:
            TypeError: If dataframe type is not supported
        """
        from pyspark.sql import DataFrame as SparkDataFrame

        # Already a Spark DataFrame
        if isinstance(dataframe, SparkDataFrame):
            return dataframe

        # PyArrow Table - convert via Pandas
        try:
            import pyarrow as pa

            if isinstance(dataframe, pa.Table):
                return self.spark.createDataFrame(dataframe.to_pandas())
        except ImportError:
            pass

        # Pandas DataFrame - convert using spark.createDataFrame
        if hasattr(dataframe, "to_dict") and hasattr(dataframe, "columns"):
            return self.spark.createDataFrame(dataframe)

        # Polars DataFrame - convert via Pandas
        if hasattr(dataframe, "to_pandas"):
            return self.spark.createDataFrame(dataframe.to_pandas())

        # List of dicts or tuples
        if isinstance(dataframe, list):
            return self.spark.createDataFrame(dataframe)

        raise TypeError(
            f"Unsupported DataFrame type: {type(dataframe)}. "
            f"Expected Spark DataFrame, Pandas DataFrame, Polars DataFrame, PyArrow Table, or list."
        )

    def _normalize_table_identifier(self, table_path: str | Path) -> str:
        """Normalize table path to Hudi table identifier.

        Hudi tables in Spark can be referenced by:
        - Catalog path: database.table_name
        - File path: /path/to/hudi/table

        Args:
            table_path: Table path or identifier

        Returns:
            Normalized identifier for SQL operations
        """
        path_str = str(table_path)

        # If it looks like a catalog path (database.table), use as-is
        if "." in path_str and "/" not in path_str and "\\" not in path_str:
            return path_str

        # For file paths, use backtick quoting for Spark SQL
        if path_str.startswith("/") or path_str.startswith("s3://") or path_str.startswith("gs://"):
            return f"`hudi`.`{path_str}`"

        return path_str

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

        Uses Spark SQL INSERT INTO for compatibility across all Hudi table types.

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

        Returns:
            Number of rows inserted
        """
        table_id = self._normalize_table_identifier(str(table_path))

        # Convert to Spark DataFrame
        spark_df = self._convert_to_spark_df(dataframe)
        row_count = spark_df.count()

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

        # Register as temp view for SQL INSERT
        temp_view = f"_hudi_insert_source_{id(spark_df)}"
        spark_df.createOrReplaceTempView(temp_view)

        try:
            if mode == "overwrite":
                # Use INSERT OVERWRITE for full table replacement
                self.spark.sql(f"INSERT OVERWRITE TABLE {table_id} SELECT * FROM {temp_view}")
            else:
                # Use INSERT INTO for append
                self.spark.sql(f"INSERT INTO {table_id} SELECT * FROM {temp_view}")

            self.logger.info(f"Inserted {row_count} rows to Hudi table {table_id}")
            return row_count

        finally:
            # Clean up temp view
            self.spark.catalog.dropTempView(temp_view)

    def _do_delete(
        self,
        table_path: Path | str,
        condition: str | Any,
    ) -> int:
        """Delete rows using Hudi DELETE statement.

        Args:
            table_path: Path to the Hudi table
            condition: SQL-like delete predicate (e.g., "order_date < '2020-01-01'")

        Returns:
            Number of rows deleted
        """
        table_id = self._normalize_table_identifier(str(table_path))

        # Get row count before delete
        count_before = self.spark.sql(f"SELECT COUNT(*) as cnt FROM {table_id}").collect()[0]["cnt"]

        # Execute DELETE
        self.spark.sql(f"DELETE FROM {table_id} WHERE {condition}")

        # Get row count after delete
        count_after = self.spark.sql(f"SELECT COUNT(*) as cnt FROM {table_id}").collect()[0]["cnt"]

        rows_deleted = count_before - count_after
        self.logger.info(f"Deleted {rows_deleted} rows from Hudi table {table_id}")
        return rows_deleted

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

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

        Returns:
            Estimated number of rows updated
        """
        table_id = self._normalize_table_identifier(str(table_path))

        # Count matching rows before update
        matching_count = self.spark.sql(f"SELECT COUNT(*) as cnt FROM {table_id} WHERE {condition}").collect()[0]["cnt"]

        if matching_count == 0:
            self.logger.info("No rows match update condition")
            return 0

        # Build SET clause
        set_clauses = ", ".join([f"{col} = {val}" for col, val in updates.items()])

        # Execute UPDATE
        self.spark.sql(f"UPDATE {table_id} SET {set_clauses} WHERE {condition}")

        self.logger.info(f"Updated {matching_count} rows in Hudi table {table_id}")
        return matching_count

    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 rows using Hudi MERGE INTO statement.

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

        Returns:
            Number of rows affected (estimated from source size)
        """
        table_id = self._normalize_table_identifier(str(table_path))

        # Convert source to Spark DataFrame
        source_df = self._convert_to_spark_df(source_dataframe)
        source_count = source_df.count()

        # Register source as temp view
        source_view = f"_hudi_merge_source_{id(source_df)}"
        source_df.createOrReplaceTempView(source_view)

        try:
            # Build MERGE statement
            merge_sql = f"MERGE INTO {table_id} AS target\n"
            merge_sql += f"USING {source_view} AS source\n"
            merge_sql += f"ON {merge_condition}\n"

            if when_matched:
                set_clauses = ", ".join([f"target.{col} = {val}" for col, val in when_matched.items()])
                merge_sql += f"WHEN MATCHED THEN UPDATE SET {set_clauses}\n"

            if when_not_matched:
                columns = ", ".join(when_not_matched.keys())
                values = ", ".join([str(v) for v in when_not_matched.values()])
                merge_sql += f"WHEN NOT MATCHED THEN INSERT ({columns}) VALUES ({values})"

            # Execute MERGE
            self.spark.sql(merge_sql)

            self.logger.info(f"Merged {source_count} source rows into Hudi table {table_id}")
            return source_count

        finally:
            # Clean up temp view
            self.spark.catalog.dropTempView(source_view)


def get_hudi_maintenance_operations(
    spark_session: Any = None,
    record_key: str | None = None,
    precombine_field: str | None = None,
    table_type: str = "COPY_ON_WRITE",
    working_dir: str | Path | None = None,
) -> HudiMaintenanceOperations | None:
    """Get Hudi maintenance operations if PySpark is available.

    Args:
        spark_session: Active SparkSession with Hudi bundle (required)
        record_key: Field name for record key
        precombine_field: Field name for ordering during updates
        table_type: Hudi table type ("COPY_ON_WRITE" or "MERGE_ON_READ")
        working_dir: Optional working directory

    Returns:
        HudiMaintenanceOperations if PySpark is available and spark_session
        is provided, None otherwise
    """
    if not PYSPARK_AVAILABLE:
        logger.debug("Hudi maintenance not available (pyspark not installed)")
        return None

    if spark_session is None:
        logger.debug("No SparkSession provided, cannot create Hudi maintenance operations")
        return None

    return HudiMaintenanceOperations(
        spark_session=spark_session,
        record_key=record_key,
        precombine_field=precombine_field,
        table_type=table_type,
        working_dir=working_dir,
    )


__all__ = [
    "HudiMaintenanceOperations",
    "get_hudi_maintenance_operations",
    "PYSPARK_AVAILABLE",
]
