"""Databricks SQL engine utilities and helpers.

This module provides a thin “do the right thing” layer over:
- Databricks SQL Statement Execution API (warehouse)
- Spark SQL / Delta Lake (when running inside a Spark-enabled context)

It includes helpers to:
- Build fully-qualified table names
- Execute SQL via Spark or Databricks SQL API
- Insert Arrow/Spark data into Delta tables (append/overwrite/merge)
- Generate DDL from Arrow schemas
"""

import dataclasses
import logging
import random
import string
import time
from typing import Optional, Union, Any, Dict, List, Literal

import pyarrow as pa

from .statement_result import StatementResult
from .types import column_info_to_arrow_field
from .. import DatabricksPathKind, DatabricksPath
from ..workspaces import WorkspaceService
from ...libs.databrickslib import databricks_sdk
from ...libs.sparklib import SparkSession, SparkDataFrame, pyspark
from ...types import is_arrow_type_string_like, is_arrow_type_binary_like
from ...types.cast.cast_options import CastOptions
from ...types.cast.registry import convert
from ...types.cast.spark_cast import cast_spark_dataframe

try:
    from delta.tables import DeltaTable as SparkDeltaTable
except ImportError:
    class SparkDeltaTable:
        @classmethod
        def forName(cls, *args, **kwargs):
            from delta.tables import DeltaTable
            return DeltaTable.forName(*args, **kwargs)


if databricks_sdk is not None:
    from databricks.sdk.service.sql import (
        StatementResponse, Disposition, Format,
        ExecuteStatementRequestOnWaitTimeout, StatementParameterListItem
    )
    StatementResponse = StatementResponse
else:
    class StatementResponse:  # pragma: no cover
        pass


logger = logging.getLogger(__name__)

if pyspark is not None:
    import pyspark.sql.functions as F

__all__ = ["SQLEngine", "StatementResult"]


@dataclasses.dataclass
class CreateTablePlan:
    sql: str
    properties: dict[str, Any]
    warnings: list[str]
    result: Any = None  # StatementResult when executed


_INVALID_COL_CHARS = set(" ,;{}()\n\t=")


def _escape_sql_string(s: str) -> str:
    return s.replace("'", "''")


def _quote_ident(ident: str) -> str:
    # Always quote to be safe; also allows reserved keywords
    escaped = ident.replace("`", "``")
    return f"`{escaped}`"


def _needs_column_mapping(col_name: str) -> bool:
    return any(ch in _INVALID_COL_CHARS for ch in col_name)


@dataclasses.dataclass
class SQLEngine(WorkspaceService):
    """Execute SQL statements and manage tables via Databricks SQL / Spark."""
    warehouse_id: Optional[str] = None
    catalog_name: Optional[str] = None
    schema_name: Optional[str] = None

    def table_full_name(
        self,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        safe_chars: bool = True,
    ) -> str:
        """Build a fully qualified table name (catalog.schema.table).

        Args:
            catalog_name: Optional catalog override (defaults to engine.catalog_name).
            schema_name: Optional schema override (defaults to engine.schema_name).
            table_name: Table name to qualify.
            safe_chars: Whether to wrap each identifier in backticks.

        Returns:
            Fully qualified table name string.
        """
        catalog_name = catalog_name or self.catalog_name
        schema_name = schema_name or self.schema_name

        assert catalog_name, "No catalog name given"
        assert schema_name, "No schema name given"
        assert table_name, "No table name given"

        if safe_chars:
            return f"`{catalog_name}`.`{schema_name}`.`{table_name}`"
        return f"{catalog_name}.{schema_name}.{table_name}"

    def _catalog_schema_table_names(self, full_name: str) -> tuple[Optional[str], Optional[str], Optional[str]]:
        """Parse a catalog.schema.table string into components.

        Supports partial names:
        - table
        - schema.table
        - catalog.schema.table

        Backticks are stripped.

        Args:
            full_name: Fully qualified or partial table name.

        Returns:
            Tuple of (catalog_name, schema_name, table_name).
        """
        parts = [_.strip("`") for _ in full_name.split(".")]

        if len(parts) == 0:
            return self.catalog_name, self.schema_name, None
        if len(parts) == 1:
            return self.catalog_name, self.schema_name, parts[0]
        if len(parts) == 2:
            return self.catalog_name, parts[0], parts[1]

        catalog_name, schema_name, table_name = parts[-3], parts[-2], parts[-1]
        catalog_name = catalog_name or self.catalog_name
        schema_name = schema_name or self.schema_name
        return catalog_name, schema_name, table_name

    def _default_warehouse(
        self,
        cluster_size: str = "Small"
    ):
        """Pick a default SQL warehouse (best-effort) matching the desired size.

        Args:
            cluster_size: Desired warehouse size (Databricks "cluster_size"), e.g. "Small".
                If empty/None, returns the first warehouse encountered.

        Returns:
            Warehouse object.

        Raises:
            ValueError: If no warehouses exist in the workspace.
        """
        wk = self.workspace.sdk()
        existing = list(wk.warehouses.list())
        first = None

        for warehouse in existing:
            if first is None:
                first = warehouse

            if cluster_size:
                if getattr(warehouse, "cluster_size", None) == cluster_size:
                    logger.debug("Default warehouse match found: id=%s cluster_size=%s", warehouse.id, warehouse.cluster_size)
                    return warehouse
            else:
                logger.debug("Default warehouse selected (first): id=%s", warehouse.id)
                return warehouse

        if first is not None:
            logger.info(
                "No warehouse matched cluster_size=%s; falling back to first warehouse id=%s cluster_size=%s",
                cluster_size,
                getattr(first, "id", None),
                getattr(first, "cluster_size", None),
            )
            return first

        raise ValueError(f"No default warehouse found in {wk.config.host}")

    def _get_or_default_warehouse_id(self, cluster_size: str = "Small") -> str:
        """Return configured warehouse_id or resolve a default one.

        Args:
            cluster_size: Desired warehouse size filter used when resolving defaults.

        Returns:
            Warehouse id string.
        """
        if not self.warehouse_id:
            dft = self._default_warehouse(cluster_size=cluster_size)
            self.warehouse_id = dft.id
            logger.info("Resolved default warehouse_id=%s (cluster_size=%s)", self.warehouse_id, cluster_size)

        return self.warehouse_id

    @staticmethod
    def _random_suffix(prefix: str = "") -> str:
        """Generate a unique suffix for temporary resources."""
        unique = "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(8))
        timestamp = int(time.time() * 1000)
        return f"{prefix}{timestamp}_{unique}"

    @staticmethod
    def _sql_preview(sql: str, limit: int = 220) -> str:
        """Short, single-line preview for logs (avoids spewing giant SQL)."""
        if not sql:
            return ""
        return sql[:limit] + ("…" if len(sql) > limit else "")

    def execute(
        self,
        statement: Optional[str] = None,
        *,
        engine: Optional[Literal["spark", "api"]] = None,
        warehouse_id: Optional[str] = None,
        byte_limit: Optional[int] = None,
        disposition: Optional["Disposition"] = None,
        format: Optional["Format"] = None,
        on_wait_timeout: Optional["ExecuteStatementRequestOnWaitTimeout"] = None,
        parameters: Optional[List["StatementParameterListItem"]] = None,
        row_limit: Optional[int] = None,
        wait_timeout: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        wait_result: bool = True,
    ) -> "StatementResult":
        """Execute a SQL statement via Spark or Databricks SQL Statement Execution API.

        Engine resolution:
        - If `engine` is not provided and a Spark session is active -> uses Spark.
        - Otherwise uses Databricks SQL API (warehouse).

        Waiting behavior (`wait_result`):
        - If True (default): returns a StatementResult in terminal state (SUCCEEDED/FAILED/CANCELED).
        - If False: returns immediately with the initial handle (caller can `.wait()` later).

        Args:
            statement: SQL statement to execute. If None, a `SELECT *` is generated from the table params.
            engine: "spark" or "api".
            warehouse_id: Warehouse override (for API engine).
            byte_limit: Optional byte limit for results.
            disposition: Result disposition mode (API engine).
            format: Result format (API engine).
            on_wait_timeout: Timeout behavior for waiting (API engine).
            parameters: Optional statement parameters (API engine).
            row_limit: Optional row limit for results (API engine).
            wait_timeout: API wait timeout value.
            catalog_name: Optional catalog override for API engine.
            schema_name: Optional schema override for API engine.
            table_name: Optional table override used when `statement` is None.
            wait_result: Whether to block until completion (API engine).

        Returns:
            StatementResult.
        """
        # --- Engine auto-detection ---
        if not engine:
            if pyspark is not None:
                spark_session = SparkSession.getActiveSession()
                if spark_session is not None:
                    engine = "spark"

        # --- Spark path ---
        if engine == "spark":
            spark_session = SparkSession.getActiveSession()
            if spark_session is None:
                raise ValueError("No spark session found to run sql query")

            df: SparkDataFrame = spark_session.sql(statement)

            if row_limit:
                df = df.limit(row_limit)

            logger.debug(
                "SPARK SQL executed query:\n%s",
                statement
            )

            # Avoid Disposition dependency if SDK imports are absent
            spark_disp = disposition if disposition is not None else getattr(globals().get("Disposition", object), "EXTERNAL_LINKS", None)

            return StatementResult(
                engine=self,
                statement_id="sparksql",
                disposition=spark_disp,
                _spark_df=df,
            )

        # --- API path defaults ---
        if format is None:
            format = Format.ARROW_STREAM

        if (disposition is None or disposition == Disposition.INLINE) and format in [Format.CSV, Format.ARROW_STREAM]:
            disposition = Disposition.EXTERNAL_LINKS

        if not statement:
            full_name = self.table_full_name(catalog_name=catalog_name, schema_name=schema_name, table_name=table_name)
            statement = f"SELECT * FROM {full_name}"

        if not warehouse_id:
            warehouse_id = self._get_or_default_warehouse_id()

        response = self.workspace.sdk().statement_execution.execute_statement(
            statement=statement,
            warehouse_id=warehouse_id,
            byte_limit=byte_limit,
            disposition=disposition,
            format=format,
            on_wait_timeout=on_wait_timeout,
            parameters=parameters,
            row_limit=row_limit,
            wait_timeout=wait_timeout,
            catalog=catalog_name or self.catalog_name,
            schema=schema_name or self.schema_name,
        )

        execution = StatementResult(
            engine=self,
            statement_id=response.statement_id,
            _response=response,
            disposition=disposition,
        )

        logger.info(
            "API SQL executed statement '%s'",
            execution.statement_id
        )
        logger.debug(
            "API SQL executed query:\n%s",
            statement
        )

        return execution.wait() if wait_result else execution

    def spark_table(
        self,
        full_name: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
    ):
        """Return a DeltaTable handle for a given table name (Spark context required)."""
        if not full_name:
            full_name = self.table_full_name(
                catalog_name=catalog_name,
                schema_name=schema_name,
                table_name=table_name,
            )
        return SparkDeltaTable.forName(
            sparkSession=SparkSession.getActiveSession(),
            tableOrViewName=full_name,
        )

    def insert_into(
        self,
        data: Union[pa.Table, pa.RecordBatch, pa.RecordBatchReader, SparkDataFrame],
        location: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        mode: str = "auto",
        cast_options: Optional[CastOptions] = None,
        overwrite_schema: bool | None = None,
        match_by: Optional[list[str]] = None,
        zorder_by: Optional[list[str]] = None,
        optimize_after_merge: bool = False,
        vacuum_hours: int | None = None,
        spark_session: Optional[SparkSession] = None,
        spark_options: Optional[Dict[str, Any]] = None,
    ):
        """Insert data into a Delta table using Spark when available; otherwise stage Arrow.

        Strategy:
        - If Spark is available and we have an active session (or Spark DF input) -> use `spark_insert_into`.
        - Otherwise -> use `arrow_insert_into` (stages Parquet to a temp volume + runs SQL INSERT/MERGE).

        Args:
            data: Arrow or Spark data to insert.
            location: Fully qualified table name override.
            catalog_name: Optional catalog override.
            schema_name: Optional schema override.
            table_name: Optional table name override.
            mode: Insert mode ("auto", "append", "overwrite").
            cast_options: Optional casting options.
            overwrite_schema: Whether to overwrite schema (Spark path).
            match_by: Merge keys for upserts (MERGE semantics). When set, mode affects behavior.
            zorder_by: Z-ORDER columns (SQL path uses OPTIMIZE ZORDER; Spark path uses Delta optimize API).
            optimize_after_merge: Whether to run OPTIMIZE after a merge (SQL path) / after merge+zorder (Spark path).
            vacuum_hours: Optional VACUUM retention window.
            spark_session: Optional SparkSession override.
            spark_options: Optional Spark write options.

        Returns:
            None (mutates the destination table).
        """

        if pyspark is not None:
            spark_session = SparkSession.getActiveSession() if spark_session is None else spark_session

            if spark_session is not None or isinstance(data, SparkDataFrame):
                return self.spark_insert_into(
                    data=data,
                    location=location,
                    catalog_name=catalog_name,
                    schema_name=schema_name,
                    table_name=table_name,
                    mode=mode,
                    cast_options=cast_options,
                    overwrite_schema=overwrite_schema,
                    match_by=match_by,
                    zorder_by=zorder_by,
                    optimize_after_merge=optimize_after_merge,
                    vacuum_hours=vacuum_hours,
                    spark_options=spark_options,
                )

        return self.arrow_insert_into(
            data=data,
            location=location,
            catalog_name=catalog_name,
            schema_name=schema_name,
            table_name=table_name,
            mode=mode,
            cast_options=cast_options,
            overwrite_schema=overwrite_schema,
            match_by=match_by,
            zorder_by=zorder_by,
            optimize_after_merge=optimize_after_merge,
            vacuum_hours=vacuum_hours,
        )

    def arrow_insert_into(
        self,
        data: Union[pa.Table, pa.RecordBatch, pa.RecordBatchReader],
        location: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        mode: str = "auto",
        cast_options: Optional[CastOptions] = None,
        overwrite_schema: bool | None = None,
        match_by: Optional[list[str]] = None,
        zorder_by: Optional[list[str]] = None,
        optimize_after_merge: bool = False,
        vacuum_hours: int | None = None,
        existing_schema: pa.Schema | None = None,
        temp_volume_path: Optional[Union[str, DatabricksPath]] = None,
    ):
        """Insert Arrow data by staging Parquet to a temp volume and running Databricks SQL.

        Notes:
        - If the table does not exist, it is created from the input Arrow schema (best-effort).
        - If `match_by` is provided, uses MERGE INTO (upsert).
        - Otherwise uses INSERT INTO / INSERT OVERWRITE depending on mode.

        Args:
            data: Arrow table/batch data to insert.
            location: Fully qualified table name override.
            catalog_name: Optional catalog override.
            schema_name: Optional schema override.
            table_name: Optional table name override.
            mode: Insert mode ("auto", "append", "overwrite"). ("auto" behaves like append here.)
            cast_options: Optional casting options.
            overwrite_schema: Reserved for parity with Spark path (unused here).
            match_by: Merge keys for MERGE INTO upserts.
            zorder_by: Columns for OPTIMIZE ZORDER BY.
            optimize_after_merge: Run OPTIMIZE after MERGE (in addition to ZORDER optimization).
            vacuum_hours: Optional VACUUM retention window in hours.
            existing_schema: Optional pre-fetched destination schema (Arrow).
            temp_volume_path: Optional temp volume path override.

        Returns:
            None.
        """
        location, catalog_name, schema_name, table_name = self._check_location_params(
            location=location,
            catalog_name=catalog_name,
            schema_name=schema_name,
            table_name=table_name,
            safe_chars=True,
        )

        with self.connect() as connected:
            if existing_schema is None:
                try:
                    existing_schema = connected.get_table_schema(
                        catalog_name=catalog_name,
                        schema_name=schema_name,
                        table_name=table_name,
                        to_arrow_schema=True,
                    )
                except ValueError as exc:
                    data_tbl = convert(data, pa.Table)
                    existing_schema = data_tbl.schema
                    logger.warning(
                        "Table %s not found (%s). Creating it from input schema (columns=%s)",
                        location,
                        exc,
                        existing_schema.names,
                    )

                    connected.create_table(
                        field=existing_schema,
                        catalog_name=catalog_name,
                        schema_name=schema_name,
                        table_name=table_name,
                        if_not_exists=True,
                    )

                    try:
                        return connected.arrow_insert_into(
                            data=data_tbl,
                            location=location,
                            catalog_name=catalog_name,
                            schema_name=schema_name,
                            table_name=table_name,
                            mode="overwrite",
                            cast_options=cast_options,
                            overwrite_schema=overwrite_schema,
                            match_by=match_by,
                            zorder_by=zorder_by,
                            optimize_after_merge=optimize_after_merge,
                            vacuum_hours=vacuum_hours,
                            existing_schema=existing_schema,
                        )
                    except Exception:
                        logger.exception("Arrow insert failed after auto-creating %s; attempting cleanup (DROP TABLE)", location)
                        try:
                            connected.drop_table(location=location)
                        except Exception:
                            logger.exception("Failed to drop table %s after auto creation error", location)
                        raise

            transaction_id = self._random_suffix()

            data_tbl = convert(
                data, pa.Table,
                options=cast_options, target_field=existing_schema
            )
            num_rows = data_tbl.num_rows

            logger.debug(
                "Arrow inserting %s rows into %s (mode=%s, match_by=%s, zorder_by=%s)",
                num_rows,
                location,
                mode,
                match_by,
                zorder_by,
            )

            # Write in temp volume
            temp_volume_path = connected.dbfs_path(
                kind=DatabricksPathKind.VOLUME,
                parts=[catalog_name, schema_name, "tmp", "sql", transaction_id],
            ) if temp_volume_path is None else DatabricksPath.parse(obj=temp_volume_path, workspace=connected.workspace)

            logger.debug("Staging Parquet to temp volume: %s", temp_volume_path)
            temp_volume_path.mkdir()
            temp_volume_path.write_arrow_table(data_tbl)

            columns = list(existing_schema.names)
            cols_quoted = ", ".join([f"`{c}`" for c in columns])

            statements: list[str] = []

            if match_by:
                on_condition = " AND ".join([f"T.`{k}` = S.`{k}`" for k in match_by])

                update_cols = [c for c in columns if c not in match_by]
                if update_cols:
                    update_set = ", ".join([f"T.`{c}` = S.`{c}`" for c in update_cols])
                    update_clause = f"WHEN MATCHED THEN UPDATE SET {update_set}"
                else:
                    update_clause = ""

                insert_clause = (
                    f"WHEN NOT MATCHED THEN INSERT ({cols_quoted}) "
                    f"VALUES ({', '.join([f'S.`{c}`' for c in columns])})"
                )

                merge_sql = f"""MERGE INTO {location} AS T
USING (
  SELECT {cols_quoted} FROM parquet.`{temp_volume_path}`
) AS S
ON {on_condition}
{update_clause}
{insert_clause}"""
                statements.append(merge_sql)
            else:
                if mode.lower() in ("overwrite",):
                    insert_sql = f"""INSERT OVERWRITE {location}
SELECT {cols_quoted}
FROM parquet.`{temp_volume_path}`"""
                else:
                    insert_sql = f"""INSERT INTO {location} ({cols_quoted})
SELECT {cols_quoted}
FROM parquet.`{temp_volume_path}`"""
                statements.append(insert_sql)

            try:
                for stmt in statements:
                    connected.execute(stmt.strip())
            finally:
                try:
                    temp_volume_path.rmdir(recursive=True)
                except Exception:
                    logger.exception("Failed cleaning temp volume: %s", temp_volume_path)

            logger.info(
                "Arrow inserted %s rows into %s (mode=%s, match_by=%s, zorder_by=%s)",
                num_rows,
                location,
                mode,
                match_by,
                zorder_by,
            )

            if zorder_by:
                zcols = ", ".join([f"`{c}`" for c in zorder_by])
                optimize_sql = f"OPTIMIZE {location} ZORDER BY ({zcols})"
                logger.info("Running OPTIMIZE ZORDER BY: %s", zorder_by)
                connected.execute(optimize_sql)

            if optimize_after_merge and match_by:
                logger.info("Running OPTIMIZE after MERGE")
                connected.execute(f"OPTIMIZE {location}")

            if vacuum_hours is not None:
                logger.info("Running VACUUM retain=%s hours", vacuum_hours)
                connected.execute(f"VACUUM {location} RETAIN {vacuum_hours} HOURS")

        return None

    def spark_insert_into(
        self,
        data: SparkDataFrame,
        *,
        location: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        mode: str = "auto",
        cast_options: Optional[CastOptions] = None,
        overwrite_schema: bool | None = None,
        match_by: Optional[list[str]] = None,
        zorder_by: Optional[list[str]] = None,
        optimize_after_merge: bool = False,
        vacuum_hours: int | None = None,
        spark_options: Optional[Dict[str, Any]] = None,
    ):
        """Insert a Spark DataFrame into a Delta table (append/overwrite/merge).

        Behavior:
        - If the table does not exist: creates it via `saveAsTable(location)` (overwrite).
        - If `match_by` is provided: uses Delta MERGE for upserts.
          - If mode == "overwrite": deletes matching keys first, then appends the batch (fast-ish overwrite-by-key).
          - Else: updates matching rows + inserts new ones.
        - Else: uses `DataFrameWriter.saveAsTable` with mode.

        Args:
            data: Spark DataFrame to insert.
            location: Fully qualified table name override.
            catalog_name: Optional catalog override.
            schema_name: Optional schema override.
            table_name: Optional table name override.
            mode: Insert mode ("auto", "append", "overwrite").
            cast_options: Optional casting options (align to destination schema).
            overwrite_schema: Whether to overwrite schema on write (when supported).
            match_by: Merge keys for upserts.
            zorder_by: Z-ORDER columns (used only if `optimize_after_merge` is True).
            optimize_after_merge: Whether to run Delta optimize (and z-order) after merge.
            vacuum_hours: Optional VACUUM retention window in hours.
            spark_options: Optional Spark write options.

        Returns:
            None.
        """
        location, catalog_name, schema_name, table_name = self._check_location_params(
            location=location,
            catalog_name=catalog_name,
            schema_name=schema_name,
            table_name=table_name,
            safe_chars=True,
        )

        logger.info(
            "Spark insert into %s (mode=%s, match_by=%s, overwrite_schema=%s)",
            location,
            mode,
            match_by,
            overwrite_schema,
        )

        spark_options = spark_options if spark_options else {}
        if overwrite_schema:
            spark_options["overwriteSchema"] = "true"

        try:
            existing_schema = self.get_table_schema(
                catalog_name=catalog_name,
                schema_name=schema_name,
                table_name=table_name,
                to_arrow_schema=False,
            )
        except ValueError:
            logger.warning("Destination table missing; creating table %s via overwrite write", location)
            data = convert(data, pyspark.sql.DataFrame)
            data.write.mode("overwrite").options(**spark_options).saveAsTable(location)
            return

        if not isinstance(data, pyspark.sql.DataFrame):
            data = convert(data, pyspark.sql.DataFrame, target_field=existing_schema)
        else:
            cast_options = CastOptions.check_arg(options=cast_options, target_field=existing_schema)
            data = cast_spark_dataframe(data, options=cast_options)

        logger.debug("Incoming Spark columns: %s", data.columns)

        if match_by:
            notnull = None
            for k in match_by:
                if k not in data.columns:
                    raise ValueError(f"Missing match key '{k}' in DataFrame columns: {data.columns}")
                notnull = data[k].isNotNull() if notnull is None else notnull & data[k].isNotNull()

            data = data.filter(notnull)
            logger.debug("Filtered null keys for match_by=%s", match_by)

        target = self.spark_table(full_name=location)

        if match_by:
            cond = " AND ".join([f"t.`{k}` <=> s.`{k}`" for k in match_by])

            if mode.casefold() == "overwrite":
                data = data.cache()
                distinct_keys = data.select([f"`{k}`" for k in match_by]).distinct()

                (
                    target.alias("t")
                    .merge(distinct_keys.alias("s"), cond)
                    .whenMatchedDelete()
                    .execute()
                )

                data.write.format("delta").mode("append").options(**spark_options).saveAsTable(location)
            else:
                update_cols = [c for c in data.columns if c not in match_by]
                set_expr = {c: F.expr(f"s.`{c}`") for c in update_cols}

                (
                    target.alias("t")
                    .merge(data.alias("s"), cond)
                    .whenMatchedUpdate(set=set_expr)
                    .whenNotMatchedInsertAll()
                    .execute()
                )
        else:
            if mode == "auto":
                mode = "append"
            logger.info("Spark write saveAsTable mode=%s", mode)
            data.write.mode(mode).options(**spark_options).saveAsTable(location)

        if optimize_after_merge and zorder_by:
            logger.info("Delta optimize + zorder (%s)", zorder_by)
            target.optimize().executeZOrderBy(*zorder_by)

        if vacuum_hours is not None:
            logger.info("Delta vacuum retain=%s hours", vacuum_hours)
            target.vacuum(vacuum_hours)

    def get_table_schema(
        self,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        to_arrow_schema: bool = True,
    ) -> Union[pa.Field, pa.Schema]:
        """Fetch a table schema from Unity Catalog and convert it to Arrow types.

        Args:
            catalog_name: Optional catalog override.
            schema_name: Optional schema override.
            table_name: Optional table name override.
            to_arrow_schema: If True returns pa.Schema; else returns a pa.Field(STRUCT<...>).

        Returns:
            Arrow Schema or a STRUCT Field representing the table.
        """
        full_name = self.table_full_name(
            catalog_name=catalog_name,
            schema_name=schema_name,
            table_name=table_name,
            safe_chars=False,
        )

        wk = self.workspace.sdk()

        try:
            table = wk.tables.get(full_name)
        except Exception as e:
            raise ValueError(f"Table %s not found, {type(e)} {e}" % full_name)

        fields = [column_info_to_arrow_field(_) for _ in table.columns]

        if to_arrow_schema:
            return pa.schema(fields, metadata={b"name": table_name})
        return pa.field(table.name, pa.struct(fields))

    def drop_table(
        self,
        location: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
    ):
        """Drop a table if it exists."""
        location, _, _, _ = self._check_location_params(
            location=location,
            catalog_name=catalog_name,
            schema_name=schema_name,
            table_name=table_name,
            safe_chars=True,
        )
        logger.info("Dropping table if exists: %s", location)
        return self.execute(f"DROP TABLE IF EXISTS {location}")

    def create_table(
        self,
        field: Union[pa.Field, pa.Schema],
        table_fqn: Optional[str] = None,            # e.g. catalog.schema.table
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        storage_location: Optional[str] = None,     # external location path
        partition_by: Optional[list[str]] = None,
        cluster_by: Optional[bool | list[str]] = True,
        comment: Optional[str] = None,
        tblproperties: Optional[dict[str, Any]] = None,
        if_not_exists: bool = True,
        or_replace: bool = False,
        using: str = "DELTA",
        optimize_write: bool = True,
        auto_compact: bool = True,
        # perf-ish optional knobs (don’t hard-force)
        enable_cdf: Optional[bool] = None,
        enable_deletion_vectors: Optional[bool] = None,
        target_file_size: Optional[int] = None,     # bytes
        # column mapping: None=auto, "none"/"name"/"id" explicit
        column_mapping_mode: Optional[str] = None,
        execute: bool = True,
        wait_result: bool = True,
        return_plan: bool = False,
    ) -> Union[str, CreateTablePlan, "StatementResult"]:
        """
        Generate (and optionally execute) a Databricks/Delta `CREATE TABLE` statement from an Apache Arrow
        schema/field, with integration-friendly safety checks and performance-oriented defaults.

        This helper is meant to be "team-safe":
          - Quotes identifiers (catalog/schema/table/columns) to avoid SQL keyword/name edge cases.
          - Validates `partition_by` / `cluster_by` columns exist in the Arrow schema before generating SQL.
          - Supports managed or external tables via `storage_location`.
          - Optionally enables Delta Column Mapping (name/id) and applies the required protocol upgrades.

        Parameters
        ----------
        field:
            Arrow schema or field describing the table.
            - If `pa.Schema`, all schema fields are used as columns.
            - If `pa.Field` with struct type, its children become columns.
            - If `pa.Field` non-struct, it becomes a single-column table.
        table_fqn:
            Fully-qualified table name, e.g. `"catalog.schema.table"`.
            If provided, it takes precedence over `catalog_name`/`schema_name`/`table_name`.
            Parts are quoted as needed.
        catalog_name, schema_name, table_name:
            Used to build the table identifier when `table_fqn` is not provided.
            All three must be provided together.
        storage_location:
            If set, emits `LOCATION '<path>'` to create an external Delta table at the given path.
            (Path string is SQL-escaped.)
        partition_by:
            List of partition column names. Must exist in the schema.
            Note: Partitioning is a physical layout choice; only use for low-cardinality columns.
        cluster_by:
            Controls clustering / liquid clustering:
            - True  -> emits `CLUSTER BY AUTO`
            - False -> emits no clustering clause
            - list[str] -> emits `CLUSTER BY (<cols...>)` (all cols must exist in schema)
        comment:
            Optional table comment. If not provided and Arrow metadata contains `b"comment"`, that is used.
        tblproperties:
            Additional/override Delta table properties (final say).
            Example: `{"delta.enableChangeDataFeed": "true"}` or `{"delta.logRetentionDuration": "30 days"}`
        if_not_exists:
            If True, generates `CREATE TABLE IF NOT EXISTS ...`.
            Mutually exclusive with `or_replace`.
        or_replace:
            If True, generates `CREATE OR REPLACE TABLE ...`.
            Mutually exclusive with `if_not_exists`.
        using:
            Storage format keyword. Defaults to `"DELTA"`.
        optimize_write:
            Sets `delta.autoOptimize.optimizeWrite` table property.
        auto_compact:
            Sets `delta.autoOptimize.autoCompact` table property.
        enable_cdf:
            If set, adds `delta.enableChangeDataFeed` property.
            Useful for CDC pipelines; avoid enabling by default if you don't need it.
        enable_deletion_vectors:
            If set, adds `delta.enableDeletionVectors` property.
            Can improve performance for updates/deletes in some workloads (subject to platform support).
        target_file_size:
            If set, adds `delta.targetFileSize` (bytes). Helps guide file sizing and reduce small files.
        column_mapping_mode:
            Delta column mapping mode:
            - None  -> auto-detect: enables `"name"` only if invalid column names are present, else `"none"`
            - "none" -> do not enable column mapping (max compatibility)
            - "name" -> enable name-based column mapping
            - "id"   -> enable id-based column mapping

            When enabled (name/id), this method also sets the required protocol properties:
            `delta.minReaderVersion=2` and `delta.minWriterVersion=5`.
        execute:
            If True, executes the generated SQL via `self.execute(...)`.
            If False, returns the SQL (or plan) without executing.
        wait_result:
            Passed to `self.execute(...)`. If True, blocks until the statement finishes.
        return_plan:
            If True, returns a `CreateTablePlan` containing SQL + applied properties + warnings (+ result if executed).
            If False:
              - returns SQL string when `execute=False`
              - returns `StatementResult` when `execute=True`

        Returns
        -------
        Union[str, CreateTablePlan, StatementResult]
            - If `execute=False` and `return_plan=False`: the SQL string.
            - If `execute=False` and `return_plan=True`: `CreateTablePlan(sql=..., properties=..., warnings=...)`.
            - If `execute=True` and `return_plan=False`: `StatementResult`.
            - If `execute=True` and `return_plan=True`: `CreateTablePlan` with `result` populated.

        Raises
        ------
        ValueError
            If required naming params are missing, if `or_replace` and `if_not_exists` conflict,
            if `column_mapping_mode` is invalid, or if partition/cluster columns are not present.

        Notes
        -----
        - Column mapping is primarily a metadata feature; performance impact is usually negligible vs IO,
          but enabling it affects compatibility with older readers.
        - Partitioning and clustering are workload-dependent: partition for selective pruning on low-cardinality
          columns; cluster for speeding up common filter/join patterns.

        Examples
        --------
        Create a managed Delta table with auto clustering and auto column mapping:
            >>> plan = client.create_table(schema, table_fqn="main.analytics.events", execute=False, return_plan=True)
            >>> print(plan.sql)

        External table with explicit partitioning and CDF:
            >>> client.create_table(
            ...     schema,
            ...     table_fqn="main.analytics.events",
            ...     storage_location="abfss://.../events",
            ...     partition_by=["event_date"],
            ...     enable_cdf=True,
            ... )
        """

        # ---- Normalize Arrow input ----
        if isinstance(field, pa.Schema):
            arrow_fields = list(field)
            schema_metadata = field.metadata or {}
        else:
            # pa.Field
            schema_metadata = field.metadata or {}
            if pa.types.is_struct(field.type):
                arrow_fields = list(field.type)
            else:
                arrow_fields = [field]

        # ---- Resolve table FQN ----
        # Prefer explicit table_fqn. Else build from catalog/schema/table_name.
        if table_fqn is None:
            if not (catalog_name and schema_name and table_name):
                raise ValueError("Provide table_fqn or (catalog_name, schema_name, table_name).")
            table_fqn = ".".join(map(_quote_ident, [catalog_name, schema_name, table_name]))
        else:
            # If caller passes raw "cat.schema.table", quote each part safely
            parts = table_fqn.split(".")
            table_fqn = ".".join(_quote_ident(p) for p in parts)

        # ---- Comments ----
        if comment is None and schema_metadata:
            c = schema_metadata.get(b"comment")
            if isinstance(c, bytes):
                comment = c.decode("utf-8")

        # ---- Detect invalid column names -> column mapping auto ----
        any_invalid = any(_needs_column_mapping(f.name) for f in arrow_fields)
        warnings: list[str] = []
        if column_mapping_mode is None:
            column_mapping_mode = "name" if any_invalid else "none"

        if column_mapping_mode not in ("none", "name", "id"):
            raise ValueError("column_mapping_mode must be one of: None, 'none', 'name', 'id'.")

        # ---- Validate partition/cluster columns exist ----
        col_names = {f.name for f in arrow_fields}
        for cols, label in ((partition_by, "partition_by"),):
            if cols:
                missing = [c for c in cols if c not in col_names]
                if missing:
                    raise ValueError(f"{label} contains unknown columns: {missing}")

        if isinstance(cluster_by, list):
            missing = [c for c in cluster_by if c not in col_names]
            if missing:
                raise ValueError(f"cluster_by contains unknown columns: {missing}")

        # ---- Column DDL ----
        # IMPORTANT: your _field_to_ddl should quote names with backticks if needed.
        # I’d recommend it ALWAYS quotes via _quote_ident internally.
        column_definitions = [self._field_to_ddl(child) for child in arrow_fields]

        # ---- Build CREATE TABLE ----
        if or_replace and if_not_exists:
            raise ValueError("Use either or_replace or if_not_exists, not both.")

        create_kw = "CREATE OR REPLACE TABLE" if or_replace else "CREATE TABLE"
        if if_not_exists and not or_replace:
            create_kw = "CREATE TABLE IF NOT EXISTS"

        sql_parts: list[str] = [
            f"{create_kw} {table_fqn} (",
            "  " + ",\n  ".join(column_definitions),
            ")",
            f"USING {using}",
        ]

        if partition_by:
            sql_parts.append("PARTITIONED BY (" + ", ".join(_quote_ident(c) for c in partition_by) + ")")
        elif cluster_by:
            if isinstance(cluster_by, bool):
                sql_parts.append("CLUSTER BY AUTO")
            else:
                sql_parts.append("CLUSTER BY (" + ", ".join(_quote_ident(c) for c in cluster_by) + ")")

        if comment:
            sql_parts.append(f"COMMENT '{_escape_sql_string(comment)}'")

        if storage_location:
            sql_parts.append(f"LOCATION '{_escape_sql_string(storage_location)}'")

        # ---- Table properties (defaults + overrides) ----
        props: dict[str, Any] = {
            "delta.autoOptimize.optimizeWrite": bool(optimize_write),
            "delta.autoOptimize.autoCompact": bool(auto_compact)
        }

        if enable_cdf is not None:
            props["delta.enableChangeDataFeed"] = bool(enable_cdf)

        if enable_deletion_vectors is not None:
            props["delta.enableDeletionVectors"] = bool(enable_deletion_vectors)

        if target_file_size is not None:
            props["delta.targetFileSize"] = int(target_file_size)

        # Column mapping + required protocol bumps
        if column_mapping_mode != "none":
            props["delta.columnMapping.mode"] = column_mapping_mode
            props["delta.minReaderVersion"] = 2
            props["delta.minWriterVersion"] = 5
        else:
            # only set explicitly if user wants; otherwise leave unset for max compatibility
            pass

        # Let caller override anything (final say)
        if tblproperties:
            props.update(tblproperties)

        if any_invalid and column_mapping_mode == "none":
            warnings.append(
                "Schema has invalid column names but column_mapping_mode='none'. "
                "This will fail unless you rename/escape columns."
            )

        if props:
            def fmt(k: str, v: Any) -> str:
                if isinstance(v, str):
                    return f"'{k}' = '{_escape_sql_string(v)}'"
                if isinstance(v, bool):
                    return f"'{k}' = '{'true' if v else 'false'}'"
                return f"'{k}' = {v}"

            sql_parts.append("TBLPROPERTIES (" + ", ".join(fmt(k, v) for k, v in props.items()) + ")")

        statement = "\n".join(sql_parts)

        plan = CreateTablePlan(sql=statement, properties=props, warnings=warnings)

        if not execute:
            return plan if return_plan else statement

        res = self.execute(statement, wait_result=wait_result)
        plan.result = res
        return plan if return_plan else res

    def _check_location_params(
        self,
        location: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        safe_chars: bool = True,
    ) -> tuple[str, Optional[str], Optional[str], Optional[str]]:
        """Resolve (location OR catalog/schema/table) into a fully-qualified name."""
        if location:
            c, s, t = self._catalog_schema_table_names(location)
            catalog_name, schema_name, table_name = catalog_name or c, schema_name or s, table_name or t

        location = self.table_full_name(
            catalog_name=catalog_name,
            schema_name=schema_name,
            table_name=table_name,
            safe_chars=safe_chars,
        )

        return location, catalog_name or self.catalog_name, schema_name or self.schema_name, table_name

    @staticmethod
    def _field_to_ddl(
        field: pa.Field,
        put_name: bool = True,
        put_not_null: bool = True,
        put_comment: bool = True,
    ) -> str:
        """Convert an Arrow Field to a Databricks SQL column DDL fragment."""
        name = field.name
        nullable_str = " NOT NULL" if put_not_null and not field.nullable else ""
        name_str = f"`{name}` " if put_name else ""

        comment_str = ""
        if put_comment and field.metadata and b"comment" in field.metadata:
            comment = field.metadata[b"comment"].decode("utf-8")
            comment_str = f" COMMENT '{comment}'"

        if not pa.types.is_nested(field.type):
            sql_type = SQLEngine._arrow_to_sql_type(field.type)
            return f"{name_str}{sql_type}{nullable_str}{comment_str}"

        if pa.types.is_struct(field.type):
            child_defs = [SQLEngine._field_to_ddl(child) for child in field.type]
            struct_body = ", ".join(child_defs)
            return f"{name_str}STRUCT<{struct_body}>{nullable_str}{comment_str}"

        if pa.types.is_map(field.type):
            map_type: pa.MapType = field.type
            key_type = SQLEngine._field_to_ddl(map_type.key_field, put_name=False, put_comment=False, put_not_null=False)
            val_type = SQLEngine._field_to_ddl(map_type.item_field, put_name=False, put_comment=False, put_not_null=False)
            return f"{name_str}MAP<{key_type}, {val_type}>{nullable_str}{comment_str}"

        if pa.types.is_list(field.type) or pa.types.is_large_list(field.type):
            list_type: pa.ListType = field.type
            elem_type = SQLEngine._field_to_ddl(list_type.value_field, put_name=False, put_comment=False, put_not_null=False)
            return f"{name_str}ARRAY<{elem_type}>{nullable_str}{comment_str}"

        raise TypeError(f"Cannot make ddl field from {field}")

    @staticmethod
    def _arrow_to_sql_type(arrow_type: Union[pa.DataType, pa.Decimal128Type]) -> str:
        """Convert an Arrow data type to a Databricks SQL type string."""
        if pa.types.is_boolean(arrow_type):
            return "BOOLEAN"
        if pa.types.is_int8(arrow_type):
            return "TINYINT"
        if pa.types.is_int16(arrow_type):
            return "SMALLINT"
        if pa.types.is_int32(arrow_type):
            return "INT"
        if pa.types.is_int64(arrow_type):
            return "BIGINT"
        if pa.types.is_float32(arrow_type):
            return "FLOAT"
        if pa.types.is_float64(arrow_type):
            return "DOUBLE"
        if is_arrow_type_string_like(arrow_type):
            return "STRING"
        if is_arrow_type_binary_like(arrow_type):
            return "BINARY"
        if pa.types.is_timestamp(arrow_type):
            tz = getattr(arrow_type, "tz", None)
            return "TIMESTAMP" if tz else "TIMESTAMP_NTZ"
        if pa.types.is_date(arrow_type):
            return "DATE"
        if pa.types.is_decimal(arrow_type):
            return f"DECIMAL({arrow_type.precision}, {arrow_type.scale})"
        if pa.types.is_null(arrow_type):
            return "STRING"
        raise ValueError(f"Cannot make ddl type for {arrow_type}")
