"""
BigQuery extractor for EL layer.

Extraction method: Spark JDBC (parallel reads).

Legacy native EXPORT DATA method (_extract_native_parallel) is retained
for potential future opt-in use but is NOT called by default.
"""

import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from dvt.federation.extractors.base import (
    BaseExtractor,
    ExtractionConfig,
    ExtractionResult,
)


class BigQueryExtractor(BaseExtractor):
    """BigQuery-specific extractor with native cloud export support.

    Extraction priority:
    1. EXPORT DATA for GCS buckets - fastest, parallel
    2. Spark JDBC for local filesystem - parallel reads
    """

    adapter_types = ["bigquery"]

    def _get_connection(self, config: ExtractionConfig = None) -> Any:
        """Get or create a BigQuery client connection.

        If self.connection is None but connection_config is available,
        creates a new connection using google.cloud.bigquery.

        Args:
            config: Optional extraction config with connection_config

        Returns:
            BigQuery client (with cursor() method wrapper)
        """
        if self.connection is not None:
            return self.connection

        # Return cached lazy connection if available
        if self._lazy_connection is not None:
            return self._lazy_connection

        # Try to get connection_config from config or instance
        conn_config = None
        if config and config.connection_config:
            conn_config = config.connection_config
        elif self.connection_config:
            conn_config = self.connection_config

        if not conn_config:
            raise ValueError(
                "No connection provided and no connection_config available. "
                "Either provide a connection to the extractor or include "
                "connection_config in ExtractionConfig."
            )

        try:
            from google.cloud import bigquery
            from google.oauth2 import service_account
        except ImportError:
            raise ImportError(
                "google-cloud-bigquery is required for BigQuery extraction. "
                "Install with: pip install google-cloud-bigquery"
            )

        from dvt.federation.auth.bigquery import BigQueryAuthHandler

        handler = BigQueryAuthHandler()
        connect_kwargs = handler.get_native_connection_kwargs(conn_config)

        # Build credentials
        credentials = None
        if connect_kwargs.get("keyfile"):
            credentials = service_account.Credentials.from_service_account_file(
                connect_kwargs["keyfile"]
            )
        elif connect_kwargs.get("keyfile_json"):
            credentials = service_account.Credentials.from_service_account_info(
                connect_kwargs["keyfile_json"]
            )

        # Create BigQuery client
        client = bigquery.Client(
            project=connect_kwargs.get("project"),
            credentials=credentials,
        )

        # Wrap client to provide cursor() interface for compatibility
        class BigQueryConnectionWrapper:
            def __init__(self, bq_client):
                self._client = bq_client

            def cursor(self):
                return BigQueryCursorWrapper(self._client)

        class BigQueryCursorWrapper:
            def __init__(self, bq_client):
                self._client = bq_client
                self._rows = None  # Materialized list of tuples

            def execute(self, query, params=None):
                result = self._client.query(query).result()
                self._rows = [tuple(row.values()) for row in result]
                self._row_idx = 0

            def fetchone(self):
                if self._rows and self._row_idx < len(self._rows):
                    row = self._rows[self._row_idx]
                    self._row_idx += 1
                    return row
                return None

            def fetchmany(self, size=1):
                if not self._rows:
                    return []
                end = min(self._row_idx + size, len(self._rows))
                batch = self._rows[self._row_idx : end]
                self._row_idx = end
                return batch

            def fetchall(self):
                if self._rows:
                    remaining = self._rows[self._row_idx :]
                    self._row_idx = len(self._rows)
                    return remaining
                return []

            def close(self):
                self._rows = None
                self._row_idx = 0

        self._lazy_connection = BigQueryConnectionWrapper(client)
        return self._lazy_connection

    def supports_native_export(self, bucket_type: str) -> bool:
        """BigQuery supports native export to GCS."""
        return bucket_type == "gcs"

    def get_native_export_bucket_types(self) -> List[str]:
        return ["gcs"]

    def extract(
        self,
        config: ExtractionConfig,
        output_path: Path,
    ) -> ExtractionResult:
        """Extract data from BigQuery to Parquet via Spark JDBC."""
        return self._extract_jdbc(config, output_path)

    def _extract_native_parallel(
        self,
        config: ExtractionConfig,
        bucket_config: Dict[str, Any],
        output_path: Path,
    ) -> ExtractionResult:
        """Extract using BigQuery EXPORT DATA."""
        start_time = time.time()

        from dvt.federation.cloud_storage import CloudStorageHelper

        helper = CloudStorageHelper(bucket_config)
        query = self.build_export_query(config)

        # Generate unique staging path
        staging_suffix = helper.generate_staging_path(config.source_name)
        # BigQuery EXPORT DATA uses gs:// paths with wildcard for output
        native_path = helper.get_native_path(staging_suffix, dialect="bigquery")
        gcs_path = f"{native_path}*.parquet"

        export_sql = f"""
            EXPORT DATA OPTIONS(
                uri='{gcs_path}',
                format='PARQUET',
                overwrite=true,
                compression='ZSTD'
            ) AS
            {query}
        """

        cursor = self._get_connection(config).cursor()
        cursor.execute(export_sql)
        cursor.close()

        # Get row count
        count_query = f"SELECT COUNT(*) FROM ({query})"
        cursor = self._get_connection(config).cursor()
        cursor.execute(count_query)
        row_count = cursor.fetchone()[0]
        cursor.close()

        elapsed = time.time() - start_time
        self._log(
            f"Exported {row_count:,} rows from {config.source_name} "
            f"to GCS in {elapsed:.1f}s (parallel)"
        )

        return ExtractionResult(
            success=True,
            source_name=config.source_name,
            row_count=row_count,
            output_path=output_path,
            extraction_method="native_parallel",
            elapsed_seconds=elapsed,
        )

    def extract_hashes(
        self,
        config: ExtractionConfig,
    ) -> Dict[str, str]:
        """Extract row hashes using BigQuery MD5 function."""
        if not config.pk_columns:
            raise ValueError("pk_columns required for hash extraction")

        pk_expr = (
            config.pk_columns[0]
            if len(config.pk_columns) == 1
            else "CONCAT(" + ", '|', ".join(config.pk_columns) + ")"
        )

        if config.columns:
            cols = config.columns
        else:
            col_info = self.get_columns(config.schema, config.table)
            cols = [c["name"] for c in col_info]

        # BigQuery uses TO_HEX(MD5(...)) — insert '|' separators between columns
        col_exprs = [f"IFNULL(CAST({c} AS STRING), '')" for c in cols]
        concat_hash = ", '|', ".join(col_exprs)
        hash_expr = f"TO_HEX(MD5(CONCAT({concat_hash})))"

        query = f"""
            SELECT
                CAST({pk_expr} AS STRING) as _pk,
                {hash_expr} as _hash
            FROM `{config.schema}.{config.table}`
        """

        if config.predicates:
            where_clause = " AND ".join(config.predicates)
            query += f" WHERE {where_clause}"

        cursor = self._get_connection(config).cursor()
        cursor.execute(query)

        hashes = {}
        while True:
            batch = cursor.fetchmany(config.batch_size)
            if not batch:
                break
            for row in batch:
                hashes[row[0]] = row[1]

        cursor.close()
        return hashes

    def get_row_count(
        self,
        schema: str,
        table: str,
        predicates: Optional[List[str]] = None,
        config: ExtractionConfig = None,
    ) -> int:
        """Get row count using COUNT(*)."""
        query = f"SELECT COUNT(*) FROM `{schema}.{table}`"
        if predicates:
            query += f" WHERE {' AND '.join(predicates)}"

        cursor = self._get_connection(config).cursor()
        cursor.execute(query)
        count = cursor.fetchone()[0]
        cursor.close()
        return count

    def get_columns(
        self,
        schema: str,
        table: str,
        config: ExtractionConfig = None,
    ) -> List[Dict[str, str]]:
        """Get column metadata from INFORMATION_SCHEMA."""
        safe_table = table.replace("'", "")
        query = f"""
            SELECT column_name, data_type
            FROM `{schema}.INFORMATION_SCHEMA.COLUMNS`
            WHERE table_name = '{safe_table}'
            ORDER BY ordinal_position
        """
        cursor = self._get_connection(config).cursor()
        cursor.execute(query)

        columns = []
        for row in cursor.fetchall():
            columns.append({"name": row[0], "type": row[1]})

        cursor.close()
        return columns

    def detect_primary_key(
        self,
        schema: str,
        table: str,
        config: ExtractionConfig = None,
    ) -> List[str]:
        """Detect primary key from BigQuery table constraints.

        Note: BigQuery primary keys are relatively new and optional.
        """
        safe_table = table.replace("'", "")
        query = f"""
            SELECT column_name
            FROM `{schema}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE`
            WHERE table_name = '{safe_table}'
            AND constraint_name LIKE '%_pk'
            ORDER BY ordinal_position
        """
        cursor = self._get_connection(config).cursor()
        try:
            cursor.execute(query)
            pk_cols = [row[0] for row in cursor.fetchall()]
        except Exception:
            pk_cols = []
        finally:
            cursor.close()

        return pk_cols
