"""
Federation loader — single JDBC + adapter loader for all targets.

Loads DataFrames into target databases using:
- Spark JDBC (df.write.jdbc()) for bulk INSERT with parallel writers
- dbt adapter for DDL (CREATE/DROP/TRUNCATE) with proper dialect quoting

This replaces the previous per-adapter loader hierarchy (PostgresLoader,
DatabricksLoader, SnowflakeLoader, etc.) which had COPY/pipe paths that
added complexity without significant benefit for federation workloads.

Usage:
    from dvt.federation.loaders import get_loader

    loader = get_loader(on_progress=print)
    result = loader.load(df, config, adapter=adapter)
"""

import threading
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional


@dataclass
class LoadConfig:
    """Configuration for a single load operation."""

    table_name: str  # Fully qualified table name (schema.table or catalog.schema.table)
    mode: str = "overwrite"  # 'overwrite', 'append'
    truncate: bool = True  # Use TRUNCATE instead of DROP for overwrite
    full_refresh: bool = False  # --full-refresh: DROP + CREATE + INSERT
    connection_config: Optional[Dict[str, Any]] = None  # profiles.yml connection
    jdbc_config: Optional[Dict[str, Any]] = None  # jdbc_load settings from computes.yml
    # Phase 4 additions (placeholders):
    incremental_strategy: Optional[str] = None  # 'append', 'delete+insert', 'merge'
    unique_key: Optional[List[str]] = None  # For merge/delete+insert


@dataclass
class LoadResult:
    """Result of a load operation."""

    success: bool
    table_name: str
    row_count: int = 0
    error: Optional[str] = None
    load_method: str = "jdbc"  # 'jdbc', 'merge', 'delete+insert'
    elapsed_seconds: float = 0.0


class FederationLoader:
    """Single loader for all federation targets.

    Uses Spark JDBC for data transfer + dbt adapter for DDL.
    All adapters (Postgres, Databricks, Snowflake, BigQuery, etc.)
    use the same JDBC + adapter pattern.
    """

    def __init__(
        self,
        on_progress: Optional[Callable[[str], None]] = None,
    ):
        """Initialize loader.

        Args:
            on_progress: Optional callback for progress messages
        """
        self.on_progress = on_progress or (lambda msg: None)

    def _log(self, message: str) -> None:
        """Log a progress message."""
        self.on_progress(message)

    # Adapters whose JDBC drivers have limited session pools or are known
    # to deadlock under concurrent batch inserts.
    #
    # dict value = (max_partitions_per_model, max_concurrent_writes).
    #
    # max_partitions_per_model: caps the Spark repartition count so each
    #   model opens at most this many JDBC writer connections.
    #
    # max_concurrent_writes: a threading.Semaphore that serialises the
    #   df.write.jdbc() calls across ALL concurrent models targeting this
    #   adapter.  Oracle XE, for example, has a hard PROCESSES limit (~20)
    #   and deadlocks when 4 models x 4 partitions = 16 simultaneous
    #   writers compete for sessions.
    #
    # Users can override max_partitions via computes.yml
    # ``jdbc_load.num_partitions``.
    _SESSION_LIMITED_ADAPTERS: Dict[str, tuple] = {
        "oracle": (1, 1),  # 1 partition per model, 1 write at a time
    }

    # Class-level semaphores keyed by adapter type.  Lazily created.
    _write_locks: Dict[str, threading.Semaphore] = {}
    _write_locks_init = threading.Lock()

    @classmethod
    def _get_write_semaphore(cls, adapter_type: str) -> Optional[threading.Semaphore]:
        """Return a Semaphore for session-limited adapters, or None."""
        key = adapter_type.lower()
        cap_tuple = cls._SESSION_LIMITED_ADAPTERS.get(key)
        if cap_tuple is None:
            return None

        _, max_concurrent = cap_tuple
        with cls._write_locks_init:
            if key not in cls._write_locks:
                cls._write_locks[key] = threading.Semaphore(max_concurrent)
            return cls._write_locks[key]

    @classmethod
    def _cap_partitions(
        cls,
        adapter_type: str,
        num_partitions: int,
        jdbc_settings: Dict[str, Any],
    ) -> int:
        """Cap JDBC writer partitions for session-limited adapters.

        Oracle (especially XE) has a limited session pool and deadlocks
        when multiple concurrent models each open several JDBC writer
        connections.  This method applies a safe per-adapter cap unless
        the user explicitly set ``num_partitions`` in ``computes.yml``.

        Args:
            adapter_type: Target adapter type (e.g. "oracle", "postgres")
            num_partitions: Currently resolved partition count
            jdbc_settings: Raw jdbc_load config from computes.yml

        Returns:
            Adjusted partition count
        """
        adapter_key = adapter_type.lower()

        # If the user explicitly configured num_partitions, respect it.
        if "num_partitions" in jdbc_settings:
            return num_partitions

        cap_tuple = cls._SESSION_LIMITED_ADAPTERS.get(adapter_key)
        if cap_tuple is not None:
            max_partitions, _ = cap_tuple
            if num_partitions > max_partitions:
                return max_partitions

        return num_partitions

    def load(
        self,
        df: Any,  # pyspark.sql.DataFrame
        config: LoadConfig,
        adapter: Optional[Any] = None,
    ) -> LoadResult:
        """Load DataFrame into target database.

        Dispatches to the appropriate load strategy:
        - Default: JDBC INSERT (append or overwrite with DDL via adapter)
        - merge: Temp table + MERGE SQL (upsert via unique_key)
        - delete+insert: Temp table + DELETE matching + INSERT

        Args:
            df: PySpark DataFrame to load
            config: Load configuration
            adapter: Optional dbt adapter for DDL operations

        Returns:
            LoadResult with success status and metadata
        """
        strategy = config.incremental_strategy

        if strategy == "merge" and config.unique_key and adapter:
            return self._load_merge(df, config, adapter)
        elif strategy == "delete+insert" and config.unique_key and adapter:
            return self._load_delete_insert(df, config, adapter)
        else:
            # Default: append (incremental), overwrite (table/full-refresh)
            return self._load_jdbc(df, config, adapter)

    # =========================================================================
    # DDL Operations via Adapter
    # =========================================================================

    def _table_exists(self, adapter: Any, config: LoadConfig) -> bool:
        """Check if the target table already exists in the database.

        Uses a lightweight SELECT 1 query that works across all SQL dialects.
        This avoids information_schema queries which vary between databases.

        Args:
            adapter: dbt adapter instance
            config: Load configuration with table_name

        Returns:
            True if the table exists, False otherwise
        """
        from dvt.federation.adapter_manager import get_quoted_table_name

        quoted_table = get_quoted_table_name(adapter, config.table_name)

        with adapter.connection_named("dvt_loader"):
            try:
                adapter.execute(
                    f"SELECT 1 FROM {quoted_table} WHERE 1=0", auto_begin=True
                )
                self._commit(adapter)
                return True
            except Exception:
                return False

    def _commit(self, adapter: Any) -> None:
        """Commit the adapter connection.

        All adapter.execute() calls in this loader use auto_begin=True, which
        makes dbt properly track the transaction (transaction_open=True).
        This means adapter.connections.commit() works correctly on all adapters:

        - PostgreSQL: sends COMMIT, sets transaction_open=False
        - Databricks/Spark: no-op (auto-commit at DB level)
        - Snowflake/BigQuery: handled by adapter internals

        No fallbacks, no raw handle access. One path, works right.
        """
        adapter.connections.commit()

    def _execute_ddl(
        self,
        adapter: Any,
        config: LoadConfig,
    ) -> None:
        """Execute DDL operations via adapter.

        Handles table preparation (TRUNCATE or DROP+CREATE) using
        properly quoted identifiers from the adapter.

        Args:
            adapter: dbt adapter instance
            config: Load configuration
        """
        from dvt.federation.adapter_manager import get_quoted_table_name

        quoted_table = get_quoted_table_name(adapter, config.table_name)

        with adapter.connection_named("dvt_loader"):
            ddl_executed = False
            if config.full_refresh:
                # Full refresh: DROP + CREATE (Spark will create)
                self._log(f"Dropping {config.table_name}...")
                adapter_type = adapter.type().lower()

                # Build dialect-appropriate DROP TABLE statement.
                # Oracle: no IF EXISTS, uses CASCADE CONSTRAINTS
                # SQL Server: no CASCADE, use IF EXISTS
                # Postgres/others: standard IF EXISTS + CASCADE
                if adapter_type == "oracle":
                    drop_variants = [
                        f"DROP TABLE {quoted_table} CASCADE CONSTRAINTS",
                        f"DROP TABLE {quoted_table}",
                    ]
                elif adapter_type in ("sqlserver", "synapse", "fabric"):
                    drop_variants = [
                        f"DROP TABLE IF EXISTS {quoted_table}",
                    ]
                else:
                    drop_variants = [
                        f"DROP TABLE IF EXISTS {quoted_table} CASCADE",
                        f"DROP TABLE IF EXISTS {quoted_table}",
                    ]

                for drop_sql in drop_variants:
                    try:
                        adapter.execute(drop_sql, auto_begin=True)
                        ddl_executed = True
                        break
                    except Exception:
                        continue  # Try next variant
            elif config.truncate:
                # Truncate: faster than DROP+CREATE, preserves structure
                self._log(f"Truncating {config.table_name}...")
                try:
                    adapter.execute(f"TRUNCATE TABLE {quoted_table}", auto_begin=True)
                    ddl_executed = True
                except Exception:
                    # Table might not exist - will be created by Spark
                    pass
            # Only commit if DDL was actually executed — committing after
            # a failed TRUNCATE (no open transaction) would itself error.
            if ddl_executed:
                self._commit(adapter)

    def _ensure_schema_exists(
        self,
        adapter: Any,
        config: LoadConfig,
    ) -> None:
        """Ensure the target schema exists.

        Args:
            adapter: dbt adapter instance
            config: Load configuration
        """
        from dvt.federation.adapter_manager import parse_table_name

        parts = parse_table_name(config.table_name)
        schema = parts.get("schema")

        if schema:
            with adapter.connection_named("dvt_loader"):
                try:
                    adapter.execute(
                        f"CREATE SCHEMA IF NOT EXISTS {schema}", auto_begin=True
                    )
                    self._commit(adapter)
                except Exception:
                    pass  # Schema might already exist or we might not have permissions

    def _create_table_with_adapter(
        self,
        adapter: Any,
        df: Any,  # pyspark.sql.DataFrame
        config: LoadConfig,
    ) -> None:
        """Create target table via adapter DDL with properly quoted columns.

        Generates a CREATE TABLE IF NOT EXISTS statement using dialect-aware
        quoting for both the table name and all column names. This preserves
        original column names (including spaces, special characters) by
        wrapping them in dialect-appropriate quotes.

        Args:
            adapter: dbt adapter instance
            df: PySpark DataFrame whose schema defines the table structure
            config: Load configuration with table_name
        """
        from dvt.federation.adapter_manager import get_quoted_table_name
        from dvt.utils.identifiers import build_create_table_sql

        adapter_type = adapter.type()
        quoted_table = get_quoted_table_name(adapter, config.table_name)
        create_sql = build_create_table_sql(df, adapter_type, quoted_table)

        self._log(f"Creating table {config.table_name} via adapter DDL...")
        with adapter.connection_named("dvt_loader"):
            try:
                adapter.execute(create_sql, auto_begin=True)
                # Only commit when DDL succeeded — a failed execute() causes
                # the adapter's exception_handler to rollback, setting
                # transaction_open=False. Committing after rollback raises
                # DbtInternalError on non-PG adapters.
                self._commit(adapter)
            except Exception as e:
                # Table might already exist (IF NOT EXISTS not supported
                # everywhere, e.g. Oracle, SQL Server)
                self._log(f"Create table note: {e}")

    # =========================================================================
    # Spark JDBC Load
    # =========================================================================

    def _load_jdbc(
        self,
        df: Any,  # pyspark.sql.DataFrame
        config: LoadConfig,
        adapter: Optional[Any] = None,
    ) -> LoadResult:
        """Load using Spark JDBC with parallel writes.

        If an adapter is provided, DDL is executed via adapter first,
        then data is loaded via Spark JDBC in append mode.

        If no adapter is provided, falls back to pure Spark JDBC
        which handles DDL internally.

        Args:
            df: PySpark DataFrame to load
            config: Load configuration
            adapter: Optional dbt adapter for DDL

        Returns:
            LoadResult with success status and metadata
        """
        start_time = time.time()

        try:
            if not config.connection_config:
                raise ValueError("connection_config required for JDBC load")

            from dvt.federation.auth import get_auth_handler
            from dvt.federation.spark_manager import SparkManager

            adapter_type = config.connection_config.get("type", "")

            spark_manager = SparkManager.get_instance()
            jdbc_url = spark_manager.get_jdbc_url(config.connection_config)
            jdbc_driver = spark_manager.get_jdbc_driver(adapter_type)

            if not jdbc_driver:
                raise ValueError(f"No JDBC driver for adapter: {adapter_type}")

            # Get auth handler and validate
            auth_handler = get_auth_handler(adapter_type)
            is_valid, error_msg = auth_handler.validate(config.connection_config)
            if not is_valid:
                raise ValueError(error_msg)

            # Get JDBC auth properties
            jdbc_props = auth_handler.get_jdbc_properties(config.connection_config)

            # Get JDBC load settings from config
            jdbc_settings = config.jdbc_config or {}
            num_partitions = jdbc_settings.get("num_partitions", 4)
            batch_size = jdbc_settings.get("batch_size", 10000)

            # Cap partitions for adapters with limited session pools.
            # Oracle (especially XE) deadlocks when multiple models each
            # open 4 parallel JDBC writers.  Use 1 writer per model so
            # that threads: N never creates more than N Oracle sessions.
            num_partitions = self._cap_partitions(
                adapter_type, num_partitions, jdbc_settings
            )

            # Build JDBC properties
            properties = {
                **jdbc_props,
                "driver": jdbc_driver,
                "batchsize": str(batch_size),
            }

            # Determine mode and DDL handling
            #
            # Databricks targets: always use adapter SQL Connector path.
            # Spark's JDBC writer generates INSERT with double-quoted column
            # names (e.g., "customer_code") but Databricks SQL uses backticks
            # (`customer_code`). This causes syntax errors for ALL loads, not
            # just those with special characters.  The adapter path uses the
            # adapter's native SQL Connector which handles quoting correctly.
            if adapter and adapter_type in ("databricks", "spark"):
                return self._load_via_adapter(df, config, adapter)

            # Repartition for parallel writes (before acquiring semaphore
            # so data shuffle happens concurrently for other models).
            if num_partitions > 1:
                df = df.repartition(num_partitions)
                self._log(f"Using {num_partitions} parallel JDBC writers")

            # Acquire write semaphore for session-limited adapters.
            # Oracle XE deadlocks when multiple concurrent models write
            # simultaneously, so we serialize DDL + JDBC writes.
            sem = self._get_write_semaphore(adapter_type)
            if sem is not None:
                self._log(
                    f"Waiting for {adapter_type} write slot ({config.table_name})..."
                )
                sem.acquire()

            try:
                if adapter and config.mode == "overwrite":
                    # DDL contract for overwrite mode:
                    #
                    # 1. --full-refresh: DROP CASCADE + CREATE + INSERT
                    #    Recreates table from scratch using Spark-inferred types.
                    #
                    # 2. Table exists (subsequent run): TRUNCATE + INSERT
                    #    Preserves the existing table schema (adapter-created
                    #    datatypes like text, float8, integer) — Spark just
                    #    inserts data. No CREATE TABLE call.
                    #
                    # 3. Table does not exist (first run): CREATE + INSERT
                    #    Creates table from Spark DataFrame schema since no
                    #    prior schema exists to preserve.
                    #
                    # This contract ensures `dvt seed --spark` preserves
                    # adapter-created types when the table already exists
                    # (e.g., after `dvt seed` created it with agate types).
                    if config.full_refresh:
                        # Path 1: DROP + CREATE + INSERT
                        self._execute_ddl(adapter, config)
                        self._create_table_with_adapter(adapter, df, config)
                    elif self._table_exists(adapter, config):
                        # Path 2: TRUNCATE + INSERT (preserve existing schema)
                        self._execute_ddl(adapter, config)
                        # Intentionally skip _create_table_with_adapter() —
                        # the table already exists with the correct datatypes.
                    else:
                        # Path 3: CREATE + INSERT (first run, no table yet)
                        self._create_table_with_adapter(adapter, df, config)
                    write_mode = "append"  # Table prepared by DDL, just append
                elif adapter and config.mode == "append":
                    # Incremental append: ensure table exists, then append.
                    if not self._table_exists(adapter, config):
                        self._create_table_with_adapter(adapter, df, config)
                    write_mode = "append"
                else:
                    # Pure Spark JDBC mode (no adapter)
                    write_mode = config.mode
                    if (
                        config.mode == "overwrite"
                        and config.truncate
                        and not config.full_refresh
                    ):
                        properties["truncate"] = "true"

                self._log(f"Loading {config.table_name} via Spark JDBC...")
                df.write.jdbc(
                    url=jdbc_url,
                    table=config.table_name,
                    mode=write_mode,
                    properties=properties,
                )
            finally:
                if sem is not None:
                    sem.release()

            row_count = df.count()
            elapsed = time.time() - start_time

            self._log(f"Loaded {row_count:,} rows via JDBC in {elapsed:.1f}s")

            return LoadResult(
                success=True,
                table_name=config.table_name,
                row_count=row_count,
                load_method="jdbc",
                elapsed_seconds=elapsed,
            )

        except Exception as e:
            elapsed = time.time() - start_time
            self._log(f"JDBC load failed: {e}")
            return LoadResult(
                success=False,
                table_name=config.table_name,
                error=str(e),
                elapsed_seconds=elapsed,
            )

    # =========================================================================
    # Adapter SQL Connector Load (Databricks + special column names)
    # =========================================================================

    def _load_via_adapter(
        self,
        df: Any,  # pyspark.sql.DataFrame
        config: LoadConfig,
        adapter: Any,
    ) -> LoadResult:
        """Load using the adapter's SQL Connector for both DDL and data.

        Used when Spark JDBC cannot handle special-character column names
        on Databricks. The Databricks server rejects CREATE TABLE with
        spaces in column names unless Delta Column Mapping is enabled,
        and the JDBC driver cannot INSERT into column-mapped tables.

        This method bypasses Spark JDBC entirely:
        1. Creates the table via adapter DDL with TBLPROPERTIES
        2. Collects the DataFrame to Python rows
        3. Generates batched INSERT statements with quoted column names
        4. Executes each batch via adapter.execute() (SQL Connector)

        Performance note: This path collects data to the driver. It's
        suitable for federation models (typically < 100K rows). For very
        large datasets, users should avoid special-character columns on
        Databricks targets.

        Args:
            df: PySpark DataFrame to load
            config: Load configuration
            adapter: dbt adapter instance (required, must be Databricks)

        Returns:
            LoadResult with success status and metadata
        """
        start_time = time.time()

        try:
            from dvt.federation.adapter_manager import get_quoted_table_name
            from dvt.utils.identifiers import build_create_table_sql, quote_identifier

            adapter_type = adapter.type()
            quoted_table = get_quoted_table_name(adapter, config.table_name)

            self._log(
                f"Special column names detected — using adapter SQL Connector "
                f"for {config.table_name} (bypassing Spark JDBC)"
            )

            # Step 1: DDL contract (same as _load_jdbc)
            if config.full_refresh:
                # DROP + CREATE
                self._execute_ddl(adapter, config)
                self._create_table_with_adapter(adapter, df, config)
            elif self._table_exists(adapter, config):
                # TRUNCATE only (preserve existing schema)
                self._execute_ddl(adapter, config)
            else:
                # First run: CREATE
                self._create_table_with_adapter(adapter, df, config)

            # Step 2: Collect DataFrame to Python rows
            self._log("Collecting data for adapter-based INSERT...")
            rows = df.collect()
            row_count = len(rows)

            if row_count == 0:
                elapsed = time.time() - start_time
                self._log(f"No rows to insert (0 rows, {elapsed:.1f}s)")
                return LoadResult(
                    success=True,
                    table_name=config.table_name,
                    row_count=0,
                    load_method="adapter",
                    elapsed_seconds=elapsed,
                )

            # Step 3: Build column list (quoted for Databricks)
            col_names = [field.name for field in df.schema.fields]
            quoted_cols = ", ".join(
                quote_identifier(c, adapter_type) for c in col_names
            )

            # Step 4: Batch INSERT via adapter
            batch_size = (config.jdbc_config or {}).get("batch_size", 1000)
            total_batches = (row_count + batch_size - 1) // batch_size

            with adapter.connection_named("dvt_loader"):
                for batch_num in range(total_batches):
                    start_idx = batch_num * batch_size
                    end_idx = min(start_idx + batch_size, row_count)
                    batch = rows[start_idx:end_idx]

                    # Build VALUES rows with proper SQL escaping
                    value_rows = []
                    for row in batch:
                        vals = []
                        for val in row:
                            if val is None:
                                vals.append("NULL")
                            elif isinstance(val, str):
                                # Escape single quotes by doubling them
                                escaped = val.replace("'", "''")
                                vals.append(f"'{escaped}'")
                            elif isinstance(val, bool):
                                vals.append("TRUE" if val else "FALSE")
                            elif isinstance(val, (int, float)):
                                vals.append(str(val))
                            else:
                                # datetime, date, decimal, etc. — stringify
                                escaped = str(val).replace("'", "''")
                                vals.append(f"'{escaped}'")
                        value_rows.append(f"({', '.join(vals)})")

                    values_sql = ",\n".join(value_rows)
                    insert_sql = (
                        f"INSERT INTO {quoted_table} ({quoted_cols})\n"
                        f"VALUES\n{values_sql}"
                    )

                    adapter.execute(insert_sql, auto_begin=True)

                    if total_batches > 1:
                        self._log(
                            f"  Batch {batch_num + 1}/{total_batches}: "
                            f"{len(batch)} rows"
                        )

                self._commit(adapter)

            elapsed = time.time() - start_time
            self._log(
                f"Loaded {row_count:,} rows via adapter SQL Connector in {elapsed:.1f}s"
            )

            return LoadResult(
                success=True,
                table_name=config.table_name,
                row_count=row_count,
                load_method="adapter",
                elapsed_seconds=elapsed,
            )

        except Exception as e:
            elapsed = time.time() - start_time
            self._log(f"Adapter load failed: {e}")
            return LoadResult(
                success=False,
                table_name=config.table_name,
                error=str(e),
                elapsed_seconds=elapsed,
            )

    # =========================================================================
    # MERGE Load (Phase 4)
    # =========================================================================

    def _get_staging_table_name(self, config: LoadConfig) -> str:
        """Generate a staging table name for merge/delete+insert operations.

        Creates a temp table name in the same schema as the target table.
        Uses a deterministic name based on the target table to avoid collisions.

        Args:
            config: Load configuration with table_name

        Returns:
            Staging table name like 'schema._dvt_staging_my_model'
        """
        from dvt.federation.adapter_manager import parse_table_name

        parts = parse_table_name(config.table_name)
        schema = parts.get("schema") or ""
        identifier = parts.get("identifier") or "model"
        staging_id = f"_dvt_staging_{identifier}"
        if schema:
            return f"{schema}.{staging_id}"
        return staging_id

    def _build_unique_index_sql(
        self,
        adapter: Any,
        table_name: str,
        unique_key: List[str],
    ) -> str:
        """Build CREATE UNIQUE INDEX SQL for PostgreSQL ON CONFLICT support.

        PostgreSQL's INSERT ... ON CONFLICT requires an existing unique index
        or constraint. This builds the DDL; the caller should execute it in
        the same connection/transaction as the MERGE itself.

        Args:
            adapter: dbt adapter instance
            table_name: Target table name (may be schema-qualified)
            unique_key: List of column names forming the unique key

        Returns:
            CREATE UNIQUE INDEX IF NOT EXISTS SQL string
        """
        from dvt.federation.adapter_manager import (
            get_quoted_table_name,
            parse_table_name,
        )
        from dvt.utils.identifiers import quote_identifier

        adapter_type = adapter.type()
        parts = parse_table_name(table_name)
        identifier = parts.get("identifier") or "model"
        key_suffix = "_".join(unique_key)[:40]
        index_name = f"_dvt_uq_{identifier}_{key_suffix}"

        quoted_cols = ", ".join(
            quote_identifier(col, adapter_type) for col in unique_key
        )
        quoted_table = get_quoted_table_name(adapter, table_name)

        return (
            f"CREATE UNIQUE INDEX IF NOT EXISTS {index_name} "
            f"ON {quoted_table} ({quoted_cols})"
        )

    def _write_to_staging_table(
        self,
        df: Any,
        staging_table: str,
        config: LoadConfig,
        adapter: Any = None,
    ) -> None:
        """Write DataFrame to a staging table in the target database via JDBC.

        Creates (or overwrites) a staging table that MERGE/DELETE+INSERT
        can reference. Both the target table and staging table must be in
        the same database for the MERGE SQL to work natively.

        For Databricks/Spark targets, uses the adapter SQL Connector instead
        of Spark JDBC to avoid double-quoted column names in INSERT statements.

        Args:
            df: PySpark DataFrame to write
            staging_table: Fully qualified staging table name
            config: Load configuration (for connection + JDBC settings)
            adapter: Optional dbt adapter (required for Databricks path)
        """
        from dvt.federation.auth import get_auth_handler
        from dvt.federation.spark_manager import SparkManager

        adapter_type = config.connection_config.get("type", "")

        # Databricks/Spark: use adapter SQL Connector to avoid Spark JDBC's
        # double-quoted column names which Databricks rejects.
        if adapter and adapter_type in ("databricks", "spark"):
            staging_config = LoadConfig(
                table_name=staging_table,
                mode="overwrite",
                truncate=False,
                full_refresh=True,  # Always DROP+CREATE staging tables
                connection_config=config.connection_config,
                jdbc_config=config.jdbc_config,
            )
            result = self._load_via_adapter(df, staging_config, adapter)
            if not result.success:
                raise RuntimeError(
                    f"Failed to write staging table via adapter: {result.error}"
                )
            return

        spark_manager = SparkManager.get_instance()
        jdbc_url = spark_manager.get_jdbc_url(config.connection_config)
        jdbc_driver = spark_manager.get_jdbc_driver(adapter_type)

        auth_handler = get_auth_handler(adapter_type)
        jdbc_props = auth_handler.get_jdbc_properties(config.connection_config)

        jdbc_settings = config.jdbc_config or {}
        num_partitions = jdbc_settings.get("num_partitions", 4)
        batch_size = jdbc_settings.get("batch_size", 10000)

        # Cap partitions for session-limited adapters (Oracle, etc.)
        num_partitions = self._cap_partitions(
            adapter_type, num_partitions, jdbc_settings
        )

        properties = {
            **jdbc_props,
            "driver": jdbc_driver,
            "batchsize": str(batch_size),
        }

        if num_partitions > 1:
            df = df.repartition(num_partitions)

        # Snowflake case-folding fix: Spark JDBC creates staging table columns
        # with double-quoted lowercase names (e.g., "customer_code"), making
        # them case-sensitive in Snowflake.  But our MERGE/DELETE+INSERT SQL
        # uses unquoted column names which Snowflake uppercases to
        # CUSTOMER_CODE.  This mismatch causes "SQL compilation error".
        # Fix: uppercase all DataFrame column names for Snowflake before
        # writing the staging table, so they match unquoted SQL references.
        if adapter_type == "snowflake":
            for col_name in df.columns:
                df = df.withColumnRenamed(col_name, col_name.upper())

        self._log(f"Writing to staging table {staging_table}...")
        df.write.jdbc(
            url=jdbc_url,
            table=staging_table,
            mode="overwrite",
            properties=properties,
        )

    def _drop_staging_table(self, adapter: Any, staging_table: str) -> None:
        """Drop the staging table after merge/delete+insert.

        Args:
            adapter: dbt adapter instance
            staging_table: Fully qualified staging table name
        """
        from dvt.federation.adapter_manager import get_quoted_table_name

        quoted = get_quoted_table_name(adapter, staging_table)
        with adapter.connection_named("dvt_loader"):
            try:
                adapter.execute(f"DROP TABLE IF EXISTS {quoted}", auto_begin=True)
            except Exception:
                pass  # Best-effort cleanup
            self._commit(adapter)

    def _load_merge(
        self,
        df: Any,
        config: LoadConfig,
        adapter: Any,
    ) -> LoadResult:
        """Load using MERGE via temp table pattern.

        1. Ensure target table exists (CREATE IF NOT EXISTS)
        2. Write result DataFrame to a staging table via Spark JDBC
        3. Execute dialect-specific MERGE SQL via the adapter
        4. Drop the staging table

        The MERGE executes entirely inside the target database — both
        the target table and staging table are in the same DB.

        Args:
            df: PySpark DataFrame to load
            config: Load configuration (must have unique_key)
            adapter: dbt adapter for MERGE SQL execution

        Returns:
            LoadResult with load_method='merge'
        """
        start_time = time.time()
        staging_table = self._get_staging_table_name(config)

        try:
            # 1. Ensure target table exists
            self._create_table_with_adapter(adapter, df, config)

            adapter_type = config.connection_config.get("type", "")

            # 2. Write to staging table
            self._write_to_staging_table(df, staging_table, config, adapter=adapter)

            # 3. Execute MERGE SQL (with unique index creation in same transaction for PG)
            columns = [f.name for f in df.schema.fields]
            merge_sql = self._build_merge_sql(
                adapter_type=adapter_type,
                adapter=adapter,
                target_table=config.table_name,
                staging_table=staging_table,
                unique_key=config.unique_key,
                columns=columns,
            )

            self._log(f"Executing MERGE into {config.table_name}...")
            with adapter.connection_named("dvt_loader"):
                # Ensure unique index exists for Postgres ON CONFLICT
                # Must be in the same connection as the MERGE to avoid
                # the index being lost when a separate connection closes.
                if adapter_type == "postgres" and config.unique_key:
                    index_sql = self._build_unique_index_sql(
                        adapter, config.table_name, config.unique_key
                    )
                    try:
                        adapter.execute(index_sql, auto_begin=True)
                    except Exception:
                        pass  # Index may already exist
                adapter.execute(merge_sql, auto_begin=True)
                self._commit(adapter)

            row_count = df.count()
            elapsed = time.time() - start_time
            self._log(f"Merged {row_count:,} rows in {elapsed:.1f}s")

            return LoadResult(
                success=True,
                table_name=config.table_name,
                row_count=row_count,
                load_method="merge",
                elapsed_seconds=elapsed,
            )

        except Exception as e:
            elapsed = time.time() - start_time
            self._log(f"MERGE load failed: {e}")
            return LoadResult(
                success=False,
                table_name=config.table_name,
                error=str(e),
                elapsed_seconds=elapsed,
            )
        finally:
            # Always clean up staging table
            try:
                self._drop_staging_table(adapter, staging_table)
            except Exception:
                pass

    def _build_merge_sql(
        self,
        adapter_type: str,
        adapter: Any,
        target_table: str,
        staging_table: str,
        unique_key: List[str],
        columns: List[str],
    ) -> str:
        """Generate dialect-specific MERGE SQL.

        Different databases have different MERGE syntax:
        - Postgres: INSERT ... ON CONFLICT ... DO UPDATE SET
        - MySQL: INSERT ... ON DUPLICATE KEY UPDATE
        - Redshift: DELETE + INSERT (no native MERGE)
        - Databricks/Snowflake/BigQuery/SQL Server: MERGE INTO ... USING ... ON

        Args:
            adapter_type: Target database adapter type
            adapter: dbt adapter for quoting
            target_table: Target table name
            staging_table: Staging table name
            unique_key: Columns for match condition
            columns: All column names in the DataFrame

        Returns:
            SQL string for the merge operation
        """
        from dvt.federation.adapter_manager import get_quoted_table_name
        from dvt.utils.identifiers import quote_identifier

        qt = get_quoted_table_name(adapter, target_table)
        qs = get_quoted_table_name(adapter, staging_table)

        # Quote column names and unique keys
        qcols = [quote_identifier(c, adapter_type) for c in columns]
        qkeys = [quote_identifier(k, adapter_type) for k in unique_key]
        non_key_cols = [c for c in columns if c not in unique_key]
        q_non_key = [quote_identifier(c, adapter_type) for c in non_key_cols]

        at = adapter_type.lower()

        if at == "postgres":
            # INSERT ... ON CONFLICT (...) DO UPDATE SET ...
            col_list = ", ".join(qcols)
            key_list = ", ".join(qkeys)
            select_list = ", ".join(qcols)
            update_set = (
                ", ".join(f"{c} = EXCLUDED.{c}" for c in q_non_key)
                if q_non_key
                else qkeys[0] + " = EXCLUDED." + qkeys[0]
            )
            return (
                f"INSERT INTO {qt} ({col_list}) "
                f"SELECT {select_list} FROM {qs} "
                f"ON CONFLICT ({key_list}) DO UPDATE SET {update_set}"
            )

        elif at == "mysql":
            # INSERT ... ON DUPLICATE KEY UPDATE ...
            col_list = ", ".join(qcols)
            select_list = ", ".join(qcols)
            update_set = (
                ", ".join(f"{c} = VALUES({c})" for c in q_non_key)
                if q_non_key
                else qkeys[0] + " = VALUES(" + qkeys[0] + ")"
            )
            return (
                f"INSERT INTO {qt} ({col_list}) "
                f"SELECT {select_list} FROM {qs} "
                f"ON DUPLICATE KEY UPDATE {update_set}"
            )

        elif at == "redshift":
            # Redshift has no native MERGE — use DELETE + INSERT
            on_clause = " AND ".join(f"{qt}.{k} = {qs}.{k}" for k in qkeys)
            col_list = ", ".join(qcols)
            return (
                f"DELETE FROM {qt} USING {qs} WHERE {on_clause};\n"
                f"INSERT INTO {qt} ({col_list}) SELECT {col_list} FROM {qs}"
            )

        else:
            # Standard MERGE INTO for Databricks, Snowflake, BigQuery, SQL Server
            on_clause = " AND ".join(f"target.{k} = stg.{k}" for k in qkeys)

            if q_non_key:
                update_set = ", ".join(f"target.{c} = stg.{c}" for c in q_non_key)
            else:
                # All columns are keys — just update the first key (no-op update)
                update_set = f"target.{qkeys[0]} = stg.{qkeys[0]}"

            col_list = ", ".join(qcols)
            val_list = ", ".join(f"stg.{c}" for c in qcols)

            return (
                f"MERGE INTO {qt} AS target "
                f"USING {qs} AS stg "
                f"ON {on_clause} "
                f"WHEN MATCHED THEN UPDATE SET {update_set} "
                f"WHEN NOT MATCHED THEN INSERT ({col_list}) VALUES ({val_list})"
            )

    # =========================================================================
    # DELETE+INSERT Load (Phase 4)
    # =========================================================================

    def _load_delete_insert(
        self,
        df: Any,
        config: LoadConfig,
        adapter: Any,
    ) -> LoadResult:
        """Load using DELETE+INSERT via temp table pattern.

        1. Ensure target table exists (CREATE IF NOT EXISTS)
        2. Write result DataFrame to a staging table via Spark JDBC
        3. DELETE matching rows from target (by unique_key)
        4. INSERT all staging rows into target
        5. Drop the staging table

        Args:
            df: PySpark DataFrame to load
            config: Load configuration (must have unique_key)
            adapter: dbt adapter for SQL execution

        Returns:
            LoadResult with load_method='delete+insert'
        """
        start_time = time.time()
        staging_table = self._get_staging_table_name(config)

        try:
            # 1. Ensure target table exists
            self._create_table_with_adapter(adapter, df, config)

            # 2. Write to staging table
            self._write_to_staging_table(df, staging_table, config, adapter=adapter)

            # 3. Execute DELETE + INSERT
            adapter_type = config.connection_config.get("type", "")
            columns = [f.name for f in df.schema.fields]
            delete_sql, insert_sql = self._build_delete_insert_sql(
                adapter_type=adapter_type,
                adapter=adapter,
                target_table=config.table_name,
                staging_table=staging_table,
                unique_key=config.unique_key,
                columns=columns,
            )

            self._log(f"Executing DELETE+INSERT into {config.table_name}...")
            with adapter.connection_named("dvt_loader"):
                adapter.execute(delete_sql, auto_begin=True)
                adapter.execute(insert_sql, auto_begin=True)
                self._commit(adapter)

            row_count = df.count()
            elapsed = time.time() - start_time
            self._log(f"Delete+inserted {row_count:,} rows in {elapsed:.1f}s")

            return LoadResult(
                success=True,
                table_name=config.table_name,
                row_count=row_count,
                load_method="delete+insert",
                elapsed_seconds=elapsed,
            )

        except Exception as e:
            elapsed = time.time() - start_time
            self._log(f"DELETE+INSERT load failed: {e}")
            return LoadResult(
                success=False,
                table_name=config.table_name,
                error=str(e),
                elapsed_seconds=elapsed,
            )
        finally:
            # Always clean up staging table
            try:
                self._drop_staging_table(adapter, staging_table)
            except Exception:
                pass

    def _build_delete_insert_sql(
        self,
        adapter_type: str,
        adapter: Any,
        target_table: str,
        staging_table: str,
        unique_key: List[str],
        columns: List[str],
    ) -> tuple:
        """Generate DELETE + INSERT SQL for delete+insert strategy.

        Args:
            adapter_type: Target database adapter type
            adapter: dbt adapter for quoting
            target_table: Target table name
            staging_table: Staging table name
            unique_key: Columns for match condition
            columns: All column names

        Returns:
            Tuple of (delete_sql, insert_sql)
        """
        from dvt.federation.adapter_manager import get_quoted_table_name
        from dvt.utils.identifiers import quote_identifier

        qt = get_quoted_table_name(adapter, target_table)
        qs = get_quoted_table_name(adapter, staging_table)

        qkeys = [quote_identifier(k, adapter_type) for k in unique_key]
        qcols = [quote_identifier(c, adapter_type) for c in columns]

        at = adapter_type.lower()

        if at in ("postgres", "redshift"):
            # DELETE ... USING ... WHERE
            on_clause = " AND ".join(f"{qt}.{k} = {qs}.{k}" for k in qkeys)
            delete_sql = f"DELETE FROM {qt} USING {qs} WHERE {on_clause}"
        else:
            # Standard: DELETE WHERE unique_key IN (SELECT ... FROM staging)
            if len(qkeys) == 1:
                delete_sql = (
                    f"DELETE FROM {qt} WHERE {qkeys[0]} IN "
                    f"(SELECT {qkeys[0]} FROM {qs})"
                )
            else:
                # Multi-column key: use EXISTS with correlated subquery
                corr_clause = " AND ".join(f"{qt}.{k} = stg.{k}" for k in qkeys)
                delete_sql = (
                    f"DELETE FROM {qt} WHERE EXISTS "
                    f"(SELECT 1 FROM {qs} AS stg WHERE {corr_clause})"
                )

        col_list = ", ".join(qcols)
        insert_sql = f"INSERT INTO {qt} ({col_list}) SELECT {col_list} FROM {qs}"

        return delete_sql, insert_sql


# Backward compatibility alias
BaseLoader = FederationLoader
