"""
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 json
import logging
import os
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
import pyarrow.parquet as pq
from google.api_core import exceptions as google_exceptions
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__)


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()
    """

    # ------------------------------------------------------------------
    # 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._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,
        )
        logger.info("BigQuery client initialized for project: %s", self.project_id)

    @property
    def storage_client(self) -> storage.Client:
        """Lazy initialization of the Cloud Storage client."""
        if self._storage_client is None:
            self._storage_client = storage.Client(
                project=self.project_id,
                credentials=self.credentials,
            )
        return self._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
        logger.info("Clients closed.")

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

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

    def show(self) -> None:
        """Print a summary of the client configuration."""
        config_info = {
            "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",
            },
        }
        for section, details in config_info.items():
            print(f"\n{section}:")
            print("-" * (len(section) + 1))
            for key, value in details.items():
                print(f"  {key:20s}: {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,
    ) -> 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.

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

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

        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=bigquery.QueryPriority.BATCH,
            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,
    ) -> 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.

        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.
            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")

        if dry_run:
            job_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
            query_job = self.client.query(sql, job_config=job_config)
            estimated_bytes = query_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
            )
            return {"estimated_bytes": estimated_bytes, "estimated_gb": estimated_gb}

        query_job = self.client.query(sql)
        df = query_job.to_dataframe(timeout=timeout)

        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,
    ) -> 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``.

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

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

        if dry_run:
            job_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
            query_job = self.client.query(sql, job_config=job_config)
            logger.info(
                "Dry run successful. Estimated bytes: %s",
                f"{query_job.total_bytes_processed:,}",
            )
            return None

        query_job = self.client.query(sql)
        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)

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

        job = self.client.load_table_from_dataframe(
            df, 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.

        Args:
            table_id: Target table identifier.

        Example::

            client.truncate_table("my_table")
            # ⚠️  You are about to DELETE ALL DATA from 'project.dataset.my_table'.
            # Type the table name to confirm: my_table
        """
        full_table_id = self._get_full_table_id(table_id)

        answer = input(
            f"⚠️  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 — operation cancelled.")
            return

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

    # ------------------------------------------------------------------
    # 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 and "*" 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)

        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:
            # Materialise query results into a temporary table
            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 temporary table %s", temp_table_ref)

            source_table = self.client.get_table(temp_table_ref)

            # Configure and start the extract job
            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 for status.",
                    extract_job.job_id,
                )

            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.
        """
        # Normalise input
        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 = None  # will use load_table_from_dataframe
        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 and "*" 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)

        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:
            # Load DataFrame into a temporary BigQuery table
            load_config = bigquery.LoadJobConfig(
                write_disposition="WRITE_TRUNCATE",
            )

            if arrow_table is not None:
                # Polars path: Arrow → Parquet buffer → BigQuery (no pandas)
                load_config.source_format = bigquery.SourceFormat.PARQUET
                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
                )
            else:
                # Pandas path
                load_job = self.client.load_table_from_dataframe(
                    df, temp_table_ref, job_config=load_config
                )

            load_job.result()
            logger.info("DataFrame loaded to temp table %s", temp_table_ref)

            # Extract from temp table to GCS
            source_table = self.client.get_table(temp_table_ref)

            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("Export completed to %s", destination_uri)
            else:
                logger.info(
                    "Extract job started (job_id=%s). Check BigQuery console.",
                    extract_job.job_id,
                )

            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:
            bucket.delete_blobs(blobs)
            logger.info(
                "Deleted %d objects from %s", len(blobs), gcs_folder_path
            )
            return True, len(blobs)
        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
                    try:
                        dt_converted = pd.to_datetime(df[col], errors="coerce", infer_datetime_format=True)
                        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 _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}")
