"""
BigQuery Client Module for mlfastflow.

This module provides a high-level wrapper around the Google Cloud BigQuery client,
offering simplified methods for common operations like running queries,
converting between SQL and DataFrames (Pandas and Polars), managing tables,
and transferring data to/from Google Cloud Storage.

Both pandas-based (``sql2df``, ``df2table``) and Polars-based
(``sql2polars``, ``polars2table``) workflows are supported in this single module.
"""

import datetime
import io
import logging
import os
import threading
import time
from itertools import combinations
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import google.auth
import pandas as pd
import polars as pl
import pyarrow as pa
import pyarrow.parquet as pq
from google.api_core import exceptions as google_exceptions
from google.api_core.retry import Retry as _Retry, if_exception_type as _if_exc
from google.auth.credentials import Credentials as GoogleCredentials
from google.cloud import bigquery, storage
from google.oauth2 import service_account

# Optional dependency
try:
    import graphviz
except ImportError:
    graphviz = None

from mlfastflow.utils import timer_decorator

logger = logging.getLogger(__name__)

# Retry policy for transient Google API errors (5xx, 429, connection reset).
# Exponential back-off: 1s → 2s → 4s … capped at 60s, total deadline 300s.
_DEFAULT_RETRY = _Retry(
    predicate=_if_exc(
        google_exceptions.InternalServerError,
        google_exceptions.ServiceUnavailable,
        google_exceptions.TooManyRequests,
        google_exceptions.GatewayTimeout,
    ),
    initial=1.0,
    multiplier=2.0,
    maximum=60.0,
    deadline=300.0,
)

__all__ = ["BigQueryClient"]


class BigQueryClient:
    """A client for interacting with Google Cloud BigQuery.

    Supports both Pandas and Polars DataFrames, Google Cloud Storage operations,
    and schema introspection utilities like ERD generation.

    Args:
        project_id: The Google Cloud Project ID. If None, inferred from credentials.
        dataset_id: The default BigQuery Dataset ID.
        key_file: Path to a service account JSON key file, a dict of service
            account info, or a ``google.auth.credentials.Credentials`` object.
            When ``None``, Application Default Credentials (ADC) are used.

    Example::

        client = BigQueryClient(
            project_id="my-project",
            dataset_id="my_dataset",
            key_file="/path/to/key.json",
        )
        df = client.sql2df("SELECT * FROM my_table LIMIT 10")
        client.close()
    """

    # Refresh the underlying BQ connection every N queries to prevent staleness.
    _MAX_QUERIES_BEFORE_REFRESH: int = 100

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def __init__(
        self,
        project_id: Optional[str],
        dataset_id: str,
        key_file: Optional[Union[str, Path, Dict[str, Any], GoogleCredentials]] = None,
    ):
        self.project_id = project_id
        self.dataset_id = dataset_id
        self.key_file = key_file

        self.client: Optional[bigquery.Client] = None
        self.credentials: Optional[GoogleCredentials] = None
        self._storage_client: Optional[storage.Client] = None
        self._query_count: int = 0
        self._connection_created_at: Optional[float] = None
        self._lock = threading.Lock()  # guards connection refresh

        self._initialize_client()

    def _initialize_client(self) -> None:
        """Initialize the BigQuery client and credentials."""
        scopes = ["https://www.googleapis.com/auth/cloud-platform"]

        if isinstance(self.key_file, GoogleCredentials):
            self.credentials = self.key_file
        elif isinstance(self.key_file, dict):
            self.credentials = service_account.Credentials.from_service_account_info(
                self.key_file, scopes=scopes
            )
        elif isinstance(self.key_file, (str, Path)):
            self.credentials = service_account.Credentials.from_service_account_file(
                os.fspath(self.key_file), scopes=scopes
            )
        else:
            # Fallback to Application Default Credentials
            self.credentials, default_project = google.auth.default(scopes=scopes)
            if not self.project_id:
                self.project_id = default_project

        if not self.project_id and hasattr(self.credentials, "project_id"):
            self.project_id = self.credentials.project_id

        if not self.project_id:
            raise ValueError(
                "Project ID could not be determined. Pass project_id explicitly."
            )

        self.client = bigquery.Client(
            credentials=self.credentials,
            project=self.project_id,
        )
        self._connection_created_at = time.time()
        logger.info("BigQuery client initialized for project: %s", self.project_id)

    def _get_storage_client(self) -> storage.Client:
        """Return the Cloud Storage client, creating it lazily on first access."""
        if self._storage_client is None:
            self._storage_client = storage.Client(
                project=self.project_id,
                credentials=self.credentials,
            )
        return self._storage_client

    @property
    def storage_client(self) -> storage.Client:
        """Lazy initialization of the Cloud Storage client."""
        return self._get_storage_client()

    # -- Context manager support ----------------------------------------

    def __enter__(self) -> "BigQueryClient":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    def close(self) -> None:
        """Close the BigQuery and Storage clients and release resources."""
        if self.client:
            self.client.close()
            self.client = None
        if self._storage_client:
            self._storage_client.close()
            self._storage_client = None
        self.credentials = None
        self._connection_created_at = None
        self._query_count = 0
        logger.info("Clients closed.")

    def __del__(self) -> None:
        try:
            self.close()
        except Exception:
            pass  # Suppress errors during interpreter shutdown

    # -- Connection management ------------------------------------------

    def _refresh_connection(self) -> None:
        """Re-initialize the BigQuery connection to prevent staleness."""
        logger.info(
            "Refreshing BigQuery connection after %d queries.", self._query_count
        )
        if self.client:
            self.client.close()
        self.client = bigquery.Client(
            credentials=self.credentials,
            project=self.project_id,
        )
        self._connection_created_at = time.time()
        self._query_count = 0
        logger.info("BigQuery connection refreshed.")

    def _check_and_refresh(self) -> None:
        """Thread-safe: refresh connection if threshold reached, then count."""
        with self._lock:
            if self._query_count >= self._MAX_QUERIES_BEFORE_REFRESH:
                self._refresh_connection()
            self._query_count += 1

    # -- Inspect ---------------------------------------------------------

    def show(self) -> None:
        """Log a summary of the client configuration."""
        sections = {
            "GCP Configuration": {
                "Project ID": self.project_id,
                "Dataset ID": self.dataset_id,
            },
            "Client Status": {
                "BigQuery Client": "Initialized" if self.client else "Not initialized",
                "Storage Client": "Initialized" if self._storage_client else "Not initialized",
                "Credentials": "Set" if self.credentials else "Not set",
                "Query count": str(self._query_count),
            },
        }
        for section, details in sections.items():
            logger.info("%s:", section)
            for key, value in details.items():
                logger.info("  %-20s: %s", key, value)

    def __repr__(self) -> str:
        status = "connected" if self.client else "closed"
        return (
            f"BigQueryClient(project_id={self.project_id!r}, "
            f"dataset_id={self.dataset_id!r}, status={status!r})"
        )

    # ------------------------------------------------------------------
    # SQL Execution
    # ------------------------------------------------------------------

    @timer_decorator
    def run_sql(
        self,
        sql: str,
        timeout: int = 300,
        dry_run: bool = False,
        priority: str = "INTERACTIVE",
    ) -> Optional[str]:
        """Execute a SQL query for DDL/DML operations.

        Args:
            sql: SQL query to execute (CREATE, INSERT, TRUNCATE, DELETE, etc.).
            timeout: Timeout in seconds.
            dry_run: If ``True``, only validate the query without executing it.
            priority: BigQuery job priority — ``'INTERACTIVE'`` (default, runs
                immediately) or ``'BATCH'`` (may be delayed up to 24 h by GCP).

        Returns:
            The BigQuery job ID as a string, useful for tracking async jobs.

        Raises:
            ValueError: If *sql* is empty or *priority* is invalid.
        """
        if not sql or not sql.strip():
            raise ValueError("SQL must be a non-empty string")

        _validate_enum(priority.upper(), ("INTERACTIVE", "BATCH"), "priority")
        bq_priority = (
            bigquery.QueryPriority.INTERACTIVE
            if priority.upper() == "INTERACTIVE"
            else bigquery.QueryPriority.BATCH
        )

        sql_upper = sql.upper()
        if "DELETE" in sql_upper or "TRUNCATE" in sql_upper or "DROP" in sql_upper:
            logger.warning("Executing destructive operation (DELETE/TRUNCATE/DROP).")

        job_config = bigquery.QueryJobConfig(
            use_query_cache=False,
            priority=bq_priority,
            dry_run=dry_run,
        )

        query_job = self.client.query(sql, job_config=job_config)

        if dry_run:
            logger.info(
                "Dry run successful. Estimated bytes: %s",
                f"{query_job.total_bytes_processed:,}",
            )
            return None

        query_job.result(timeout=timeout)
        logger.info("Operation completed successfully (job_id=%s)", query_job.job_id)
        return query_job.job_id

    def check_job_status(self, job_id: str) -> Dict[str, Any]:
        """Check the status of a BigQuery job.

        Args:
            job_id: The BigQuery job ID returned by :meth:`run_sql`.

        Returns:
            A dict with keys ``state``, ``error_result``, ``created``, and ``ended``.
        """
        if not job_id:
            raise ValueError("job_id must be a non-empty string")

        job = self.client.get_job(job_id)
        return {
            "state": job.state,
            "error_result": job.error_result,
            "created": job.created,
            "ended": job.ended,
        }

    # ------------------------------------------------------------------
    # SQL → DataFrame
    # ------------------------------------------------------------------

    @timer_decorator
    def sql2df(
        self,
        sql: str,
        dry_run: bool = False,
        timeout: Optional[int] = 300,
        max_gb: Optional[float] = None,
    ) -> Optional[Union[pd.DataFrame, Dict[str, Any]]]:
        """Execute SQL and return results as a Pandas DataFrame.

        Args:
            sql: SQL query to execute.
            dry_run: If ``True``, only validate and return a dict with
                estimated cost info instead of executing the query.
            timeout: Maximum seconds to wait for the query to complete.
                Defaults to 300 (5 minutes). Set to ``None`` for no limit.
            max_gb: If set, a dry-run is performed first and
                ``ValueError`` is raised if the estimated scan exceeds this
                threshold (in GB). Use to prevent accidental expensive queries.

        Returns:
            A Pandas DataFrame, or a cost-estimate dict when *dry_run* is
            ``True``::

                {"estimated_bytes": 1048576, "estimated_gb": 0.001}

        Raises:
            ValueError: If *sql* is empty or estimated cost exceeds *max_gb*.
            google.api_core.exceptions.GoogleAPICallError: On query failure.
            concurrent.futures.TimeoutError: If the query exceeds *timeout*.
        """
        if not sql or not sql.strip():
            raise ValueError("SQL must be a non-empty string")

        logger.debug("sql2df SQL: %.300s%s", sql, "..." if len(sql) > 300 else "")

        if dry_run or max_gb is not None:
            dr_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
            dr_job = self.client.query(sql, job_config=dr_config)
            estimated_bytes = dr_job.total_bytes_processed
            estimated_gb = round(estimated_bytes / (1024**3), 4)
            logger.info("Dry run: %s bytes (~%s GB)", f"{estimated_bytes:,}", estimated_gb)
            if max_gb is not None and estimated_gb > max_gb:
                raise ValueError(
                    f"Query estimated to scan {estimated_gb:.2f} GB, which exceeds "
                    f"the limit of {max_gb} GB. Pass max_gb=None to disable this check."
                )
            if dry_run:
                return {"estimated_bytes": estimated_bytes, "estimated_gb": estimated_gb}

        self._check_and_refresh()
        query_job = self.client.query(sql, retry=_DEFAULT_RETRY)
        query_job.result(timeout=timeout)  # wait with timeout; to_dataframe() no longer accepts it (BQ ≥ 3.x)
        df = query_job.to_dataframe()

        logger.info(
            "Query returned %d rows, %d columns (%s bytes processed).",
            len(df),
            len(df.columns),
            f"{query_job.total_bytes_processed:,}" if query_job.total_bytes_processed else "N/A",
        )
        return df

    @timer_decorator
    def sql2polars(
        self,
        sql: str,
        lazy: bool = False,
        dry_run: bool = False,
        timeout: Optional[int] = 600,
        max_gb: Optional[float] = None,
    ) -> Optional[Union[pl.DataFrame, pl.LazyFrame]]:
        """Execute SQL and return results as a Polars DataFrame or LazyFrame.

        Uses Arrow for zero-copy data transfer from BigQuery.

        Args:
            sql: SQL query to execute.
            lazy: If ``True``, return a ``pl.LazyFrame`` instead.
            dry_run: If ``True``, only validate and return ``None``.
            timeout: Maximum seconds to wait for the query to complete.
                Defaults to 600 (10 min) — longer than ``sql2df`` because Arrow
                transfers are used for larger datasets. Set to ``None`` for no limit.
            max_gb: If set, raises ``ValueError`` before executing if the
                estimated scan exceeds this threshold (in GB).

        Returns:
            Polars DataFrame/LazyFrame, or ``None`` for dry runs.

        Raises:
            ValueError: If *sql* is empty or estimated cost exceeds *max_gb*.
            google.api_core.exceptions.GoogleAPICallError: On query failure.
        """
        if not sql or not sql.strip():
            raise ValueError("SQL must be a non-empty string")

        logger.debug("sql2polars SQL: %.300s%s", sql, "..." if len(sql) > 300 else "")

        if dry_run or max_gb is not None:
            dr_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
            dr_job = self.client.query(sql, job_config=dr_config)
            estimated_bytes = dr_job.total_bytes_processed
            estimated_gb = round(estimated_bytes / (1024**3), 4)
            logger.info(
                "Dry run: %s bytes (~%s GB)", f"{estimated_bytes:,}", estimated_gb
            )
            if max_gb is not None and estimated_gb > max_gb:
                raise ValueError(
                    f"Query estimated to scan {estimated_gb:.2f} GB, which exceeds "
                    f"the limit of {max_gb} GB. Pass max_gb=None to disable this check."
                )
            if dry_run:
                return None

        self._check_and_refresh()
        query_job = self.client.query(sql, retry=_DEFAULT_RETRY)
        query_job.result(timeout=timeout)   # wait with timeout before fetching
        arrow_table = query_job.to_arrow()
        df = pl.from_arrow(arrow_table)
        return df.lazy() if lazy else df

    @timer_decorator
    def sql2file(
        self,
        sql: str,
        file_path: Union[str, Path],
        compression: str = "snappy",
    ) -> bool:
        """Execute SQL and save results directly to a local file.

        Supported formats are inferred from the file extension:
        ``.parquet``, ``.csv``, and ``.json``.

        Args:
            sql: SQL query to execute.
            file_path: Destination file path.
            compression: Compression for Parquet files (e.g. ``snappy``, ``zstd``).

        Returns:
            ``True`` if the file was written successfully.

        Raises:
            ValueError: If *sql* is empty or the file extension is unsupported.
        """
        if not sql or not sql.strip():
            raise ValueError("SQL must be a non-empty string")

        file_path = Path(file_path)
        ext = file_path.suffix.lower()

        if ext not in (".parquet", ".csv", ".json"):
            raise ValueError(
                f"Unsupported file extension '{ext}'. Use .parquet, .csv, or .json"
            )

        file_path.parent.mkdir(parents=True, exist_ok=True)

        query_job = self.client.query(sql)
        arrow_table = query_job.to_arrow()
        df = pl.from_arrow(arrow_table)

        if ext == ".parquet":
            df.write_parquet(file_path, compression=compression)
        elif ext == ".csv":
            df.write_csv(file_path)
        elif ext == ".json":
            df.write_json(file_path)

        logger.info("Saved %d rows to %s", df.height, file_path)
        return True

    # ------------------------------------------------------------------
    # DataFrame → BigQuery
    # ------------------------------------------------------------------

    @timer_decorator
    def df2table(
        self,
        df: pd.DataFrame,
        table_id: str,
        if_exists: str = "fail",
        schema: Optional[List[bigquery.SchemaField]] = None,
    ) -> bool:
        """Upload a Pandas DataFrame to a BigQuery table.

        Args:
            df: Pandas DataFrame to upload.
            table_id: Target table name, ``dataset.table``, or fully-qualified
                ``project.dataset.table``.
            if_exists: ``'fail'``, ``'replace'``, or ``'append'``.
            schema: Optional explicit BigQuery schema.

        Returns:
            ``True`` if the upload completed successfully.

        Raises:
            ValueError: If *df* is empty or *if_exists* is invalid.
        """
        if df.empty:
            raise ValueError("DataFrame is empty. Nothing to upload.")

        full_table_id = self._get_full_table_id(table_id)

        # Convert Pandas → Arrow → Parquet buffer (same efficient path as polars2table)
        arrow_table = pa.Table.from_pandas(df)
        buffer = io.BytesIO()
        pq.write_table(arrow_table, buffer, compression="snappy")
        buffer.seek(0)

        job_config = bigquery.LoadJobConfig(
            source_format=bigquery.SourceFormat.PARQUET,
            schema=schema,
            write_disposition=self._get_write_disposition(if_exists),
        )

        job = self.client.load_table_from_file(
            buffer, full_table_id, job_config=job_config
        )
        job.result()

        logger.info("Uploaded %d rows to %s", len(df), full_table_id)
        return True

    @timer_decorator
    def polars2table(
        self,
        df: Union[pl.DataFrame, pl.LazyFrame],
        table_id: str,
        if_exists: str = "fail",
        schema: Optional[List[bigquery.SchemaField]] = None,
    ) -> bool:
        """Upload a Polars DataFrame to a BigQuery table.

        Converts to an Arrow table, writes to an in-memory Parquet buffer,
        and loads via the BigQuery ``load_table_from_file`` API for efficiency.

        Args:
            df: Polars DataFrame or LazyFrame to upload.
            table_id: Target table name, ``dataset.table``, or fully-qualified
                ``project.dataset.table``.
            if_exists: ``'fail'``, ``'replace'``, or ``'append'``.
            schema: Optional explicit BigQuery schema.

        Returns:
            ``True`` if the upload completed successfully.

        Raises:
            ValueError: If *df* is empty or *if_exists* is invalid.
        """
        if isinstance(df, pl.LazyFrame):
            df = df.collect()

        if df.height == 0:
            raise ValueError("DataFrame is empty. Nothing to upload.")

        full_table_id = self._get_full_table_id(table_id)

        # Convert Polars → Arrow → Parquet buffer (no pandas detour)
        arrow_table = df.rechunk().to_arrow()

        buffer = io.BytesIO()
        pq.write_table(arrow_table, buffer, compression="snappy")
        buffer.seek(0)

        job_config = bigquery.LoadJobConfig(
            source_format=bigquery.SourceFormat.PARQUET,
            schema=schema,
            write_disposition=self._get_write_disposition(if_exists),
        )

        job = self.client.load_table_from_file(
            buffer, full_table_id, job_config=job_config
        )
        job.result()

        logger.info("Uploaded %d rows to %s", df.height, full_table_id)
        return True

    # ------------------------------------------------------------------
    # Table Operations
    # ------------------------------------------------------------------

    def truncate_table(self, table_id: str) -> None:
        """Truncate a BigQuery table (remove all rows, preserve schema).

        This is a **destructive** operation. You will be prompted to type the
        table name to confirm before execution.
        """
        full_table_id = self._get_full_table_id(table_id)

        answer = input(
            f"\u26a0\ufe0f  You are about to DELETE ALL DATA from '{full_table_id}'.\n"
            f"Type the table name to confirm: "
        )
        if answer.strip() != table_id:
            print("Confirmation failed \u2014 operation cancelled.")
            return

        sql = f"TRUNCATE TABLE `{full_table_id}`"
        self.run_sql(sql)
        logger.info("Table %s truncated.", full_table_id)

    def list_tables(self) -> List[str]:
        """List all table names in the configured dataset.

        Returns:
            List of short table names (not fully qualified).
        """
        tables = self.client.list_tables(f"{self.project_id}.{self.dataset_id}")
        return [t.table_id for t in tables]

    def table_exists(self, table_id: str) -> bool:
        """Check whether a table exists in BigQuery.

        Args:
            table_id: Table name, ``dataset.table``, or fully-qualified.

        Returns:
            ``True`` if the table exists, ``False`` otherwise.
        """
        full_table_id = self._get_full_table_id(table_id)
        try:
            self.client.get_table(full_table_id)
            return True
        except google_exceptions.NotFound:
            return False

    def describe_table(self, table_id: str) -> List[Dict[str, str]]:
        """Return the schema of a BigQuery table as a list of dicts.

        Args:
            table_id: Table name, ``dataset.table``, or fully-qualified.

        Returns:
            A list of ``{"name", "type", "mode", "description"}`` dicts,
            one per column.

        Example::

            schema = client.describe_table("my_table")
            # [{"name": "id", "type": "INT64", "mode": "NULLABLE", ...}, ...]
        """
        full_table_id = self._get_full_table_id(table_id)
        table = self.client.get_table(full_table_id)
        return [
            {
                "name": f.name,
                "type": f.field_type,
                "mode": f.mode,
                "description": f.description or "",
            }
            for f in table.schema
        ]

    @classmethod
    def from_env(
        cls,
        project_id_env: str = "BQ_PROJECT_ID",
        dataset_id_env: str = "BQ_DATASET_ID",
        key_file_env: str = "BQ_KEY_FILE",
    ) -> "BigQueryClient":
        """Construct a ``BigQueryClient`` from environment variables.

        Reads ``BQ_PROJECT_ID``, ``BQ_DATASET_ID``, and optionally
        ``BQ_KEY_FILE`` (path to a service-account JSON file). When
        ``BQ_KEY_FILE`` is not set, Application Default Credentials are used.

        Args:
            project_id_env: Name of the env var holding the project ID.
            dataset_id_env: Name of the env var holding the dataset ID.
            key_file_env: Name of the env var holding the key-file path.

        Returns:
            An initialised ``BigQueryClient``.

        Raises:
            ValueError: If ``BQ_PROJECT_ID`` or ``BQ_DATASET_ID`` are unset.

        Example::

            # set in shell: export BQ_PROJECT_ID=my-project BQ_DATASET_ID=my_ds
            client = BigQueryClient.from_env()
        """
        project_id = os.environ.get(project_id_env)
        dataset_id = os.environ.get(dataset_id_env)
        key_file = os.environ.get(key_file_env)  # None → ADC

        missing = [v for v, k in [(project_id_env, project_id), (dataset_id_env, dataset_id)] if k is None]
        if missing:
            raise ValueError(
                f"Required environment variable(s) not set: {missing}. "
                "Set them or pass project_id/dataset_id explicitly."
            )

        return cls(project_id=project_id, dataset_id=dataset_id, key_file=key_file)

    # ------------------------------------------------------------------
    # Convenience & Exploration
    # ------------------------------------------------------------------

    def preview(
        self,
        table_id: str,
        n: int = 10,
        use_polars: bool = False,
    ) -> Union[pd.DataFrame, pl.DataFrame]:
        """Return the first *n* rows of a table.

        Args:
            table_id: Table name, ``dataset.table``, or fully-qualified.
            n: Number of rows to return. Defaults to 10.
            use_polars: If ``True``, returns a Polars DataFrame.

        Returns:
            A Pandas or Polars DataFrame with the first *n* rows.

        Example::

            df = client.preview("orders")           # first 10 rows as Pandas
            df = client.preview("orders", n=100, use_polars=True)
        """
        full_table_id = self._get_full_table_id(table_id)
        sql = f"SELECT * FROM `{full_table_id}` LIMIT {n}"
        return self.sql2polars(sql) if use_polars else self.sql2df(sql)

    def get_table_info(self, table_id: str) -> Dict[str, Any]:
        """Return metadata about a BigQuery table.

        Args:
            table_id: Table name, ``dataset.table``, or fully-qualified.

        Returns:
            A dict with keys: ``table_id``, ``rows``, ``size_bytes``,
            ``size_gb``, ``created``, ``modified``, ``expires``,
            ``partitioning``, ``clustering``, ``schema``.

        Example::

            info = client.get_table_info("orders")
            print(info["rows"], info["size_gb"])
        """
        full_table_id = self._get_full_table_id(table_id)
        table = self.client.get_table(full_table_id)
        size_bytes = table.num_bytes or 0
        return {
            "table_id": full_table_id,
            "rows": table.num_rows,
            "size_bytes": size_bytes,
            "size_gb": round(size_bytes / (1024 ** 3), 4),
            "created": table.created,
            "modified": table.modified,
            "expires": table.expires,
            "partitioning": str(table.time_partitioning) if table.time_partitioning else None,
            "clustering": table.clustering_fields,
            "schema": [
                {"name": f.name, "type": f.field_type, "mode": f.mode}
                for f in table.schema
            ],
        }

    def query_to_table(
        self,
        sql: str,
        dest_table_id: str,
        if_exists: str = "fail",
    ) -> bool:
        """Run a SQL query and materialise results into a BigQuery table.

        Uses BigQuery's native ``destination`` query config — no data passes
        through the client. Equivalent to ``CREATE TABLE ... AS SELECT ...``
        but with proper write-disposition control.

        Args:
            sql: SQL query whose results will be written to *dest_table_id*.
            dest_table_id: Destination table (short name, ``dataset.table``,
                or fully-qualified).
            if_exists: ``'fail'``, ``'replace'``, or ``'append'``.

        Returns:
            ``True`` if the operation completed successfully.

        Example::

            client.query_to_table(
                "SELECT user_id, COUNT(*) AS orders FROM orders GROUP BY 1",
                dest_table_id="summary",
                if_exists="replace",
            )
        """
        if not sql or not sql.strip():
            raise ValueError("SQL must be a non-empty string")

        full_dest = self._get_full_table_id(dest_table_id)
        job_config = bigquery.QueryJobConfig(
            destination=full_dest,
            write_disposition=self._get_write_disposition(if_exists),
            allow_large_results=True,
        )
        query_job = self.client.query(sql, job_config=job_config, retry=_DEFAULT_RETRY)
        query_job.result()
        logger.info("Query results written to %s (%s rows).", full_dest, query_job.num_dml_affected_rows)
        return True

    def copy_table(
        self,
        src_table_id: str,
        dest_table_id: str,
        if_exists: str = "fail",
    ) -> bool:
        """Copy a BigQuery table using the native server-side copy job.

        Faster and cheaper than re-reading data with ``SELECT *`` — BigQuery
        copies storage blocks directly without scanning the data.

        Args:
            src_table_id: Source table identifier.
            dest_table_id: Destination table identifier.
            if_exists: ``'fail'``, ``'replace'``, or ``'append'``.

        Returns:
            ``True`` if the copy completed successfully.

        Example::

            client.copy_table("orders_staging", "orders", if_exists="replace")
            client.copy_table("orders", "orders_backup_20260304")
        """
        src = self._get_full_table_id(src_table_id)
        dest = self._get_full_table_id(dest_table_id)
        job_config = bigquery.CopyJobConfig(
            write_disposition=self._get_write_disposition(if_exists),
        )
        copy_job = self.client.copy_table(src, dest, job_config=job_config)
        copy_job.result()
        logger.info("Copied %s → %s", src, dest)
        return True

    def wait_for_job(
        self,
        job_id: str,
        poll_interval: float = 5.0,
        timeout: Optional[int] = 600,
    ) -> Dict[str, Any]:
        """Block until a BigQuery job completes.

        The natural complement to :meth:`run_sql`, which returns a job ID
        immediately without waiting. Use this to wait for async jobs.

        Args:
            job_id: BigQuery job ID returned by :meth:`run_sql`.
            poll_interval: Seconds between status checks. Defaults to 5.
            timeout: Maximum seconds to wait. ``None`` = no limit.

        Returns:
            The final job status dict from :meth:`check_job_status`.

        Raises:
            RuntimeError: If the job finishes with an error.
            TimeoutError: If the job does not complete within *timeout*.

        Example::

            job_id = client.run_sql("CREATE TABLE archive AS SELECT ...")
            # … do other work …
            client.wait_for_job(job_id, timeout=300)
        """
        if not job_id:
            raise ValueError("job_id must be a non-empty string")

        deadline = time.time() + timeout if timeout else None

        while True:
            status = self.check_job_status(job_id)
            if status["state"] == "DONE":
                if status["error_result"]:
                    raise RuntimeError(
                        f"Job {job_id} failed: {status['error_result']}"
                    )
                logger.info("Job %s completed successfully.", job_id)
                return status

            if deadline and time.time() >= deadline:
                raise TimeoutError(
                    f"Job {job_id} did not finish within {timeout}s. "
                    "Check BigQuery console for current status."
                )
            time.sleep(poll_interval)

    @timer_decorator
    def gcs2df(
        self,
        gcs_uri: str,
        format: str = "PARQUET",
        use_polars: bool = False,
    ) -> Union[pd.DataFrame, pl.DataFrame]:
        """Load a file from Google Cloud Storage into a DataFrame.

        Downloads directly — no BigQuery query is executed.  Suitable for
        single files; for wildcard URIs (``gs://bucket/path/*.parquet``)
        route through :meth:`gcs2table` + :meth:`sql2df` instead.

        Args:
            gcs_uri: GCS URI of the file, e.g. ``gs://bucket/path/file.parquet``.
            format: ``'PARQUET'``, ``'CSV'``, or ``'JSON'``.
            use_polars: If ``True``, returns a Polars DataFrame.

        Returns:
            A Pandas or Polars DataFrame.

        Raises:
            ValueError: On invalid URI or unsupported format.

        Example::

            df = client.gcs2df("gs://my-bucket/exports/report.parquet")
            df = client.gcs2df("gs://my-bucket/data.csv", format="CSV", use_polars=True)
        """
        if not gcs_uri or not gcs_uri.startswith("gs://"):
            raise ValueError("gcs_uri must start with 'gs://'")

        fmt = format.upper()
        _validate_enum(fmt, ("PARQUET", "CSV", "JSON"), "format")

        parts = gcs_uri.replace("gs://", "").split("/", 1)
        bucket_name = parts[0]
        blob_path = parts[1] if len(parts) > 1 else ""

        bucket = self._get_storage_client().bucket(bucket_name)
        blob = bucket.blob(blob_path)

        buffer = io.BytesIO()
        blob.download_to_file(buffer)
        buffer.seek(0)

        if fmt == "PARQUET":
            return pl.read_parquet(buffer) if use_polars else pd.read_parquet(buffer)
        elif fmt == "CSV":
            return pl.read_csv(buffer) if use_polars else pd.read_csv(buffer)
        else:  # JSON / NDJSON
            return pl.read_ndjson(buffer) if use_polars else pd.read_json(buffer)

    def delete_table(self, table_id: str) -> None:
        """Permanently drop a BigQuery table.

        Prompts for confirmation before executing — the same safety pattern
        as :meth:`truncate_table`.  Unlike ``truncate_table``, this removes
        the table structure entirely and **cannot be undone**.

        Args:
            table_id: Table name, ``dataset.table``, or fully-qualified.

        Example::

            client.delete_table("old_staging_table")
            # ⚠️  You are about to DROP 'project.dataset.old_staging_table'.
            # This cannot be undone. Type the table name to confirm:
        """
        full_table_id = self._get_full_table_id(table_id)

        answer = input(
            f"\u26a0\ufe0f  You are about to DROP '{full_table_id}' permanently.\n"
            f"This cannot be undone. Type the table name to confirm: "
        )
        if answer.strip() != table_id:
            print("Confirmation failed \u2014 operation cancelled.")
            return

        self.client.delete_table(full_table_id)
        logger.info("Table %s deleted.", full_table_id)

    def execute_many(
        self,
        sql_list: List[str],
        parallel: bool = False,
    ) -> List[Optional[str]]:
        """Execute a list of SQL statements and return their job IDs.

        Args:
            sql_list: List of SQL strings to execute.
            parallel: If ``True``, submit all jobs concurrently using a
                thread pool. Each statement runs independently; errors in
                one do not cancel others.

        Returns:
            List of BigQuery job IDs (one per statement, in order).

        Raises:
            ValueError: If *sql_list* is empty.

        Example::

            # Sequential ETL steps
            client.execute_many([
                "TRUNCATE TABLE staging.orders",
                "INSERT INTO staging.orders SELECT ...",
                "CALL my_dataset.refresh_summary()",
            ])

            # Fire all jobs simultaneously, wait for all to finish
            client.execute_many([...], parallel=True)
        """
        if not sql_list:
            raise ValueError("sql_list must contain at least one SQL statement.")

        if parallel:
            import concurrent.futures

            with concurrent.futures.ThreadPoolExecutor() as pool:
                futures = [pool.submit(self.run_sql, sql) for sql in sql_list]
                return [f.result() for f in concurrent.futures.as_completed(futures)]

        results: List[Optional[str]] = []
        for i, sql in enumerate(sql_list, start=1):
            logger.info("Executing statement %d / %d", i, len(sql_list))
            results.append(self.run_sql(sql))
        return results

    # ------------------------------------------------------------------
    # BigQuery ↔ GCS
    # ------------------------------------------------------------------


    @timer_decorator
    def sql2gcs(
        self,
        sql: str,
        destination_uri: str,
        format: str = "PARQUET",
        compression: str = "SNAPPY",
        wait_for_completion: bool = True,
        timeout: int = 300,
        use_sharding: bool = True,
    ) -> bool:
        """Export BigQuery query results directly to Google Cloud Storage.

        Creates a temporary table for the query results, then uses BigQuery's
        extract job for server-side export (no data passes through the client).

        Args:
            sql: SQL query to execute.
            destination_uri: GCS URI, e.g. ``gs://bucket/path/to/file.parquet``.
            format: ``'PARQUET'``, ``'CSV'``, ``'JSON'``, or ``'AVRO'``.
            compression: ``'NONE'``, ``'GZIP'``, ``'SNAPPY'``, or ``'DEFLATE'``.
            wait_for_completion: Block until the export finishes.
            timeout: Timeout in seconds when *wait_for_completion* is ``True``.
            use_sharding: Automatically add ``-*`` wildcard for parallel export.

        Returns:
            ``True`` if the export succeeded.

        Raises:
            ValueError: On invalid arguments.
        """
        if not sql or not sql.strip():
            raise ValueError("SQL query cannot be None or empty")
        if not destination_uri or not destination_uri.startswith("gs://"):
            raise ValueError("destination_uri must start with 'gs://'")

        format = format.upper()
        compression = compression.upper()
        _validate_enum(format, ("PARQUET", "CSV", "JSON", "AVRO"), "format")
        _validate_enum(
            compression, ("NONE", "GZIP", "SNAPPY", "DEFLATE"), "compression"
        )

        if use_sharding:
            destination_uri = self._apply_sharding(destination_uri)

        temp_table_ref = (
            f"{self.project_id}.{self.dataset_id}"
            f".temp_export_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f')}"
        )

        try:
            query_config = bigquery.QueryJobConfig(
                destination=temp_table_ref,
                write_disposition="WRITE_TRUNCATE",
            )
            query_job = self.client.query(sql, job_config=query_config)
            query_job.result()
            logger.info("Query results stored in temp table %s", temp_table_ref)

            source_table = self.client.get_table(temp_table_ref)
            self._run_extract_job(
                source_table, destination_uri, format, compression,
                wait_for_completion, timeout,
            )
            return True

        finally:
            # Always clean up the temporary table
            try:
                self.client.delete_table(temp_table_ref, not_found_ok=True)
            except Exception as cleanup_err:
                logger.warning(
                    "Failed to clean up temp table %s: %s",
                    temp_table_ref,
                    cleanup_err,
                )

    @timer_decorator
    def gcs2table(
        self,
        gcs_uri: str,
        table_id: str,
        schema: Optional[List[bigquery.SchemaField]] = None,
        write_disposition: str = "WRITE_EMPTY",
        source_format: str = "PARQUET",
        allow_jagged_rows: bool = False,
        ignore_unknown_values: bool = False,
        max_bad_records: int = 0,
    ) -> bool:
        """Load data from Google Cloud Storage into a BigQuery table.

        Args:
            gcs_uri: Source GCS URI, e.g. ``gs://bucket/path/*.parquet``.
            table_id: Destination table identifier.
            schema: Optional explicit BigQuery schema.
            write_disposition: ``'WRITE_EMPTY'``, ``'WRITE_TRUNCATE'``, or
                ``'WRITE_APPEND'``.
            source_format: ``'PARQUET'``, ``'CSV'``, ``'JSON'``, ``'AVRO'``,
                or ``'ORC'``.
            allow_jagged_rows: Allow missing trailing optional columns (CSV only).
            ignore_unknown_values: Ignore extra values not in schema.
            max_bad_records: Maximum tolerated bad records before failure.

        Returns:
            ``True`` if the load completed successfully.

        Raises:
            ValueError: On invalid arguments.
        """
        if not gcs_uri or not gcs_uri.startswith("gs://"):
            raise ValueError("gcs_uri must start with 'gs://'")

        _validate_enum(
            source_format.upper(),
            ("PARQUET", "CSV", "JSON", "AVRO", "ORC"),
            "source_format",
        )
        _validate_enum(
            write_disposition,
            ("WRITE_EMPTY", "WRITE_TRUNCATE", "WRITE_APPEND"),
            "write_disposition",
        )

        full_table_id = self._get_full_table_id(table_id)

        job_config = bigquery.LoadJobConfig(
            schema=schema,
            source_format=source_format.upper(),
            write_disposition=write_disposition,
            allow_jagged_rows=allow_jagged_rows,
            ignore_unknown_values=ignore_unknown_values,
            max_bad_records=max_bad_records,
        )

        load_job = self.client.load_table_from_uri(
            gcs_uri, full_table_id, job_config=job_config
        )
        load_job.result()

        table = self.client.get_table(full_table_id)
        logger.info(
            "Loaded %d rows from %s into %s", table.num_rows, gcs_uri, full_table_id
        )
        return True

    @timer_decorator
    def df2gcs(
        self,
        df: Union[pd.DataFrame, pl.DataFrame, pl.LazyFrame],
        destination_uri: str,
        format: str = "PARQUET",
        compression: str = "SNAPPY",
        wait_for_completion: bool = True,
        timeout: int = 300,
        use_sharding: bool = True,
    ) -> bool:
        """Export a DataFrame to Google Cloud Storage via a temporary BigQuery table.

        Loads the DataFrame into a temporary BigQuery table, then uses
        BigQuery's server-side extract job to export to GCS. The temporary
        table is always cleaned up, even on failure.

        Accepts Pandas DataFrames, Polars DataFrames, and Polars LazyFrames.

        Args:
            df: The DataFrame to export.
            destination_uri: GCS URI, e.g. ``gs://bucket/path/to/file.parquet``.
            format: ``'PARQUET'``, ``'CSV'``, ``'JSON'``, or ``'AVRO'``.
            compression: ``'NONE'``, ``'GZIP'``, ``'SNAPPY'``, or ``'DEFLATE'``.
            wait_for_completion: Block until the export finishes.
            timeout: Timeout in seconds when *wait_for_completion* is ``True``.
            use_sharding: Automatically add ``-*`` wildcard for parallel export.

        Returns:
            ``True`` if the export succeeded.

        Raises:
            ValueError: On invalid arguments or empty DataFrame.
        """
        if isinstance(df, pl.LazyFrame):
            df = df.collect()

        if isinstance(df, pl.DataFrame):
            if df.height == 0:
                raise ValueError("DataFrame is empty. Nothing to export.")
            arrow_table = df.rechunk().to_arrow()
        elif isinstance(df, pd.DataFrame):
            if df.empty:
                raise ValueError("DataFrame is empty. Nothing to export.")
            arrow_table = pa.Table.from_pandas(df)
        else:
            raise TypeError(
                f"Expected pd.DataFrame, pl.DataFrame, or pl.LazyFrame, "
                f"got {type(df).__name__}"
            )

        if not destination_uri or not destination_uri.startswith("gs://"):
            raise ValueError("destination_uri must start with 'gs://'")

        format = format.upper()
        compression = compression.upper()
        _validate_enum(format, ("PARQUET", "CSV", "JSON", "AVRO"), "format")
        _validate_enum(
            compression, ("NONE", "GZIP", "SNAPPY", "DEFLATE"), "compression"
        )

        if use_sharding:
            destination_uri = self._apply_sharding(destination_uri)

        temp_table_ref = (
            f"{self.project_id}.{self.dataset_id}"
            f".temp_export_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f')}"
        )

        try:
            # All DataFrames already normalised to Arrow above
            load_config = bigquery.LoadJobConfig(
                source_format=bigquery.SourceFormat.PARQUET,
                write_disposition="WRITE_TRUNCATE",
            )
            buffer = io.BytesIO()
            pq.write_table(arrow_table, buffer, compression="snappy")
            buffer.seek(0)
            load_job = self.client.load_table_from_file(
                buffer, temp_table_ref, job_config=load_config
            )
            load_job.result()
            logger.info("DataFrame loaded to temp table %s", temp_table_ref)

            source_table = self.client.get_table(temp_table_ref)
            self._run_extract_job(
                source_table, destination_uri, format, compression,
                wait_for_completion, timeout,
            )
            return True

        finally:
            try:
                self.client.delete_table(temp_table_ref, not_found_ok=True)
            except Exception as cleanup_err:
                logger.warning(
                    "Failed to clean up temp table %s: %s",
                    temp_table_ref,
                    cleanup_err,
                )

    # ------------------------------------------------------------------
    # GCS Folder Management
    # ------------------------------------------------------------------

    def create_gcs_folder(self, gcs_folder_path: str) -> bool:
        """Create a folder marker in Google Cloud Storage.

        GCS folders are virtual constructs represented by zero-byte objects
        whose name ends with ``/``.

        Args:
            gcs_folder_path: GCS path, e.g. ``gs://bucket/folder/``.

        Returns:
            ``True`` if created successfully or already exists.
        """
        if not gcs_folder_path or not gcs_folder_path.startswith("gs://"):
            raise ValueError("gcs_folder_path must start with 'gs://'")

        # Ensure trailing slash
        if not gcs_folder_path.endswith("/"):
            gcs_folder_path += "/"

        parts = gcs_folder_path.replace("gs://", "").split("/", 1)
        bucket_name = parts[0]
        prefix = parts[1] if len(parts) > 1 else ""

        bucket = self.storage_client.bucket(bucket_name)
        blob = bucket.blob(prefix)

        if blob.exists():
            logger.info("Folder already exists: %s", gcs_folder_path)
            return True

        blob.upload_from_string("", content_type="application/x-directory")
        logger.info("Created folder: %s", gcs_folder_path)
        return True

    def delete_gcs_folder(
        self, gcs_folder_path: str, dry_run: bool = False
    ) -> Tuple[bool, int]:
        """Delete a folder and all its contents from Google Cloud Storage.

        Args:
            gcs_folder_path: GCS path to the folder, e.g.
                ``gs://bucket/folder-to-delete/``.
            dry_run: If ``True``, report how many objects *would* be deleted
                without actually deleting them.

        Returns:
            A tuple ``(success, count)`` where *count* is the number of objects
            deleted (or that would be deleted if *dry_run*).
        """
        if not gcs_folder_path or not gcs_folder_path.startswith("gs://"):
            raise ValueError("gcs_folder_path must start with 'gs://'")

        if not gcs_folder_path.endswith("/"):
            gcs_folder_path += "/"

        parts = gcs_folder_path.replace("gs://", "").split("/", 1)
        bucket_name = parts[0]
        prefix = parts[1] if len(parts) > 1 else ""

        bucket = self.storage_client.bucket(bucket_name)
        blobs = list(bucket.list_blobs(prefix=prefix))

        if not blobs:
            logger.info("No objects found at %s", gcs_folder_path)
            return True, 0

        if dry_run:
            logger.info(
                "Dry run: would delete %d objects from %s",
                len(blobs),
                gcs_folder_path,
            )
            return True, len(blobs)

        try:
            # Chunk deletions at 1000 to stay within GCS delete API limits
            _GCS_CHUNK = 1000
            total_deleted = 0
            for i in range(0, len(blobs), _GCS_CHUNK):
                chunk = blobs[i : i + _GCS_CHUNK]
                bucket.delete_blobs(chunk)
                total_deleted += len(chunk)
                logger.info(
                    "Deleted %d / %d objects from %s",
                    total_deleted, len(blobs), gcs_folder_path,
                )
            return True, total_deleted
        except Exception as e:
            logger.error("Error deleting folder %s: %s", gcs_folder_path, e)
            return False, 0

    # ------------------------------------------------------------------
    # Data Type Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def fix_mixed_types(
        df: pd.DataFrame,
        columns: Optional[List[str]] = None,
        strategy: str = "infer",
        numeric_errors: str = "coerce",
    ) -> pd.DataFrame:
        """Fix mixed data types in a Pandas DataFrame.

        Mixed-type columns often cause errors when uploading to BigQuery.

        Args:
            df: The DataFrame to process.
            columns: Columns to check. Defaults to all ``object``-dtype columns.
            strategy:
                - ``'infer'``: try numeric → datetime → fall back to string.
                - ``'to_string'``: unconditionally cast to string (nulls preserved).
            numeric_errors: ``'coerce'`` (non-numeric → NaN) or ``'raise'``.

        Returns:
            A new DataFrame with consistent column types.

        Raises:
            ValueError: If strategy is invalid.
        """
        if strategy not in ("infer", "to_string"):
            raise ValueError("strategy must be 'infer' or 'to_string'")

        df = df.copy()
        cols_to_check = (
            columns
            if columns is not None
            else [c for c in df.columns if df[c].dtype == object]
        )

        fixed: List[str] = []
        for col in cols_to_check:
            if col not in df.columns:
                logger.warning("Column '%s' not in DataFrame — skipping.", col)
                continue

            original_dtype = df[col].dtype
            null_count_before = df[col].isna().sum()

            if strategy == "infer":
                # 1. Try numeric
                converted = pd.to_numeric(df[col], errors="coerce")
                new_nulls = converted.isna().sum() - null_count_before

                if converted.notna().any() and new_nulls <= len(df) * 0.5:
                    # Numeric conversion is viable (didn't null out >50% of values)
                    if new_nulls > 0:
                        logger.warning(
                            "Column '%s': %d non-numeric values coerced to NaN.",
                            col,
                            new_nulls,
                        )
                    df[col] = converted
                else:
                    # 2. Try datetime (infer_datetime_format removed — deprecated in Pandas 2.0)
                    try:
                        dt_converted = pd.to_datetime(df[col], errors="coerce")
                        dt_nulls = dt_converted.isna().sum() - null_count_before
                        if dt_converted.notna().any() and dt_nulls <= len(df) * 0.5:
                            if dt_nulls > 0:
                                logger.warning(
                                    "Column '%s': %d values could not be parsed as datetime.",
                                    col,
                                    dt_nulls,
                                )
                            df[col] = dt_converted
                        else:
                            # 3. Fall back to string (preserve nulls)
                            df[col] = df[col].where(df[col].isna(), df[col].astype(str))
                    except Exception:
                        df[col] = df[col].where(df[col].isna(), df[col].astype(str))

            else:  # to_string — preserve nulls
                df[col] = df[col].where(df[col].isna(), df[col].astype(str))

            if df[col].dtype != original_dtype:
                fixed.append(col)
                logger.info(
                    "Column '%s': %s → %s", col, original_dtype, df[col].dtype
                )

        if fixed:
            logger.info("Fixed mixed types in columns: %s", fixed)
        else:
            logger.info("No columns required type changes.")

        return df

    # ------------------------------------------------------------------
    # ERD Generation
    # ------------------------------------------------------------------

    # Common columns that appear in many tables but don't indicate
    # real relationships — excluded from automatic edge inference.
    _ERD_IGNORE_COLUMNS: frozenset = frozenset({
        "id", "created_at", "updated_at", "modified_at", "deleted_at",
        "created_by", "updated_by", "is_active", "is_deleted",
        "name", "description", "status", "type", "metadata",
    })

    def erd(
        self,
        table_list: List[str],
        output_filename: str = "bq_erd",
        output_format: str = "png",
        view_diagram: bool = False,
        ignore_columns: Optional[frozenset] = None,
    ) -> Optional[str]:
        """Generate an Entity-Relationship Diagram for BigQuery tables.

        Matches columns by name across tables to infer relationships.
        Common columns like ``id``, ``created_at``, etc. are excluded by
        default to avoid false relationships.

        Args:
            table_list: Fully-qualified table identifiers.
            output_filename: Base name for the output file (no extension).
            output_format: ``'png'``, ``'svg'``, or ``'pdf'``.
            view_diagram: If ``True``, open the rendered diagram automatically.
            ignore_columns: Column names to exclude from relationship matching.
                Defaults to :attr:`_ERD_IGNORE_COLUMNS`. Pass ``frozenset()``
                to disable filtering.

        Returns:
            Path to the generated file, or ``None`` on failure.

        Raises:
            ImportError: If the ``graphviz`` package is not installed.
            ValueError: If *table_list* is empty.
        """
        if not graphviz:
            raise ImportError(
                "The 'graphviz' package is required for ERD generation. "
                "Install with: pip install graphviz"
            )

        if not table_list:
            raise ValueError("table_list must contain at least one table.")

        if ignore_columns is None:
            ignore_columns = self._ERD_IGNORE_COLUMNS

        dot = graphviz.Digraph("BigQuery ERD", strict=True)
        dot.attr(rankdir="LR", splines="ortho")
        dot.attr("node", shape="plaintext")

        column_map: Dict[str, List[str]] = {}

        for table_id in table_list:
            try:
                table = self.client.get_table(table_id)
                short_name = table_id.rsplit(".", 1)[-1]

                rows = [
                    f'<TR><TD COLSPAN="2" BGCOLOR="#4A90D9">'
                    f'<FONT COLOR="white"><B>{short_name}</B></FONT></TD></TR>'
                ]
                for field in table.schema:
                    col_lower = field.name.lower()
                    if col_lower not in ignore_columns:
                        column_map.setdefault(field.name, []).append(table_id)

                    rows.append(
                        f'<TR><TD BGCOLOR="#F0F0F0">{field.field_type}</TD>'
                        f'<TD PORT="{field.name}">{field.name}</TD></TR>'
                    )

                label = (
                    f'<<TABLE BORDER="1" CELLBORDER="1" CELLSPACING="0">'
                    f'{"" .join(rows)}</TABLE>>'
                )
                dot.node(table_id, label=label)

            except Exception as e:
                logger.warning("Could not process table %s: %s", table_id, e)

        # Draw edges for shared column names (likely foreign-key relationships)
        for col_name, tables in column_map.items():
            if len(tables) > 1:
                for t1, t2 in combinations(tables, 2):
                    dot.edge(f"{t1}:{col_name}", f"{t2}:{col_name}")

        try:
            output_path = dot.render(
                output_filename,
                format=output_format,
                cleanup=True,
                view=view_diagram,
            )
            logger.info("ERD generated at %s", output_path)
            return output_path
        except Exception as e:
            logger.error("Error rendering ERD: %s", e)
            return None

    # ------------------------------------------------------------------
    # Internal Helpers
    # ------------------------------------------------------------------

    def _apply_sharding(self, destination_uri: str) -> str:
        """Add a ``-*`` wildcard to a GCS URI for parallel sharded export."""
        if "*" not in destination_uri:
            base, ext = os.path.splitext(destination_uri)
            destination_uri = f"{base}-*{ext}" if ext else f"{destination_uri}-*"
            logger.info("Sharding enabled: %s", destination_uri)
        return destination_uri

    def _run_extract_job(
        self,
        source_table: bigquery.Table,
        destination_uri: str,
        format: str,
        compression: str,
        wait_for_completion: bool,
        timeout: int,
    ) -> None:
        """Run a BigQuery extract job from a table to GCS."""
        extract_config = bigquery.ExtractJobConfig(destination_format=format)
        if compression != "NONE":
            extract_config.compression = compression

        extract_job = self.client.extract_table(
            source_table, destination_uri, job_config=extract_config
        )
        if wait_for_completion:
            extract_job.result(timeout=timeout)
            logger.info("Extract completed to %s", destination_uri)
        else:
            logger.info(
                "Extract job started (job_id=%s). Check BigQuery console.",
                extract_job.job_id,
            )

    def _get_full_table_id(self, table_id: str) -> str:
        """Construct a fully-qualified ``project.dataset.table`` identifier."""
        dot_count = table_id.count(".")
        if dot_count == 0:
            return f"{self.project_id}.{self.dataset_id}.{table_id}"
        if dot_count == 1:
            return f"{self.project_id}.{table_id}"
        return table_id

    @staticmethod
    def _get_write_disposition(if_exists: str) -> str:
        """Map a user-friendly string to a BigQuery WriteDisposition constant."""
        mapping = {
            "fail": bigquery.WriteDisposition.WRITE_EMPTY,
            "replace": bigquery.WriteDisposition.WRITE_TRUNCATE,
            "append": bigquery.WriteDisposition.WRITE_APPEND,
        }
        if if_exists not in mapping:
            raise ValueError(
                f"Invalid if_exists value: {if_exists!r}. "
                f"Choose from {list(mapping.keys())}"
            )
        return mapping[if_exists]


# ------------------------------------------------------------------
# Module-level helpers
# ------------------------------------------------------------------

def _validate_enum(value: str, valid: tuple, name: str) -> None:
    """Raise ``ValueError`` if *value* is not among *valid* choices."""
    if value not in valid:
        raise ValueError(f"{name} must be one of {list(valid)}, got {value!r}")
