"""
BigQuery client wrapper.
Handles dry runs (cost estimation) and query execution.
"""

import os

# ── Hardcoded config ───────────────────────────────────────────────────────────
BIGQUERY_PROJECT_ID      = "wehealanalysis"
BQ_COST_THRESHOLD_BYTES  = 5 * 1024 ** 3   # 5 GB
# ──────────────────────────────────────────────────────────────────────────────


def _get_bq_client():
    """Lazy import so the agent works even without google-cloud-bigquery installed."""
    try:
        from google.cloud import bigquery
        return bigquery.Client(project=BIGQUERY_PROJECT_ID)
    except ImportError:
        raise ImportError(
            "google-cloud-bigquery is not installed.\n"
            "Run: pip install google-cloud-bigquery"
        )


class BigQueryClient:
    def __init__(self):
        self._client = None
        self.cost_threshold_bytes = int(
            os.getenv("BQ_COST_THRESHOLD_BYTES", BQ_COST_THRESHOLD_BYTES)
        )

    @property
    def client(self):
        if self._client is None:
            self._client = _get_bq_client()
        return self._client

    def dry_run(self, sql: str) -> tuple:
        """
        Runs a BigQuery dry run to validate SQL and estimate bytes processed.
        Returns (bytes_processed: int, error: str|None).
        """
        try:
            from google.cloud import bigquery

            job_config = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
            job = self.client.query(sql, job_config=job_config)
            bytes_processed = job.total_bytes_processed

            # Enforce cost threshold
            if bytes_processed > self.cost_threshold_bytes:
                gb = bytes_processed / (1024 ** 3)
                threshold_gb = self.cost_threshold_bytes / (1024 ** 3)
                return bytes_processed, (
                    f"Query would scan {gb:.1f} GB, exceeding the "
                    f"{threshold_gb:.0f} GB safety limit. "
                    f"Please add tighter filters (e.g. narrow the date range)."
                )

            return bytes_processed, None

        except Exception as e:
            return None, str(e)

    def execute(self, sql: str, max_rows: int = 1000) -> tuple:
        """
        Executes a BigQuery query and returns results as a list of dicts.
        Returns (data: list, row_count: int, error: str|None).
        """
        try:
            job = self.client.query(sql)
            results = job.result()

            rows = []
            for row in results:
                rows.append(dict(row))
                if len(rows) >= max_rows:
                    break

            return rows, results.total_rows, None

        except Exception as e:
            return None, 0, str(e)