import asyncio
import itertools
import json
import math
import os
import urllib.request
import warnings
from datetime import datetime, timedelta, timezone
from typing import Optional

warnings.filterwarnings(
    "ignore", message="Your application has authenticated using end user credentials"
)
import pandas as pd
import panel as pn
import param

if pn.state._is_pyodide:
    CACHE_TO_DISK = False
else:
    CACHE_TO_DISK = True
    import diskcache
CACHE = True

BQ_BASE_URL = "https://bigquery.googleapis.com/bigquery/v2"
NUMERIC_TYPES = {"INTEGER", "INT64", "FLOAT", "FLOAT64", "NUMERIC", "BIGNUMERIC"}
TIMESTAMP_TYPES = {"TIMESTAMP", "DATETIME", "DATE"}


# @apply_to_all(log_class_methods, except_for=[])
class BigQuerier(param.Parameterized):

    job_project = param.String(precedence=-1)
    data_project = param.String(precedence=-1)
    dataset = param.String(precedence=-1)
    table = param.String(precedence=-1)
    access_token = param.String(default="", precedence=-1, doc="OAuth 2.0 access token. Required in Pyodide where google-auth is unavailable.")
    _access_token: str = ""  # class-level store for static method access

    columns = param.Dict(default={}, doc="Column name -> BigQuery type")
    ranges = param.Dict(default={}, doc="Numeric/timestamp column name -> (min, max)")
    data = param.DataFrame()
    is_loading = param.Boolean(
        default=False, doc="True while async requests are in flight"
    )
    nqueued = param.Integer(default=0, doc="Number of async requests currently queued")
    max_concurrent = param.Integer(
        default=5, bounds=(1, None), doc="Maximum number of concurrent async requests"
    )

    selected_columns = param.ListSelector(
        default=[], objects=set(), allow_None=True, doc="Selected column name"
    )
    indexer_columns = param.ListSelector(
        default=[], objects=set(), allow_None=True, doc="Selected column name"
    )

    index_chunker = param.Dict()

    def __init__(
        self,
        ignore_columns: list[str] | None = None,
        keep_columns: list[str] | None = None,
        categorical_columns: list[str] | None = None,
        **params,
    ):
        """Authenticate, fetch column schema and numeric ranges from the table.

        Parameters
        ----------
        ignore_columns : list[str], optional
            Column names to exclude from the schema entirely.
            Mutually exclusive with ``keep_columns``.
        keep_columns : list[str], optional
            Column names to keep; all other columns are ignored.
            Mutually exclusive with ``ignore_columns``.
        categorical_columns : list[str], optional
            Column names to treat as categorical (ListSelector) regardless of
            their BigQuery type.
        **params
            Keyword arguments passed to ``param.Parameterized``, including
            ``job_project``, ``data_project``, ``dataset``, ``table``, and
            ``ranges``.  Pre-supplied ``ranges`` entries skip the corresponding
            ``_fetch_ranges`` or ``_fetch_unique_values`` call.  For
            numeric/timestamp columns supply ``(min, max)`` tuples; for
            categorical columns supply a list of allowed values.
        """
        if ignore_columns is not None and keep_columns is not None:
            raise ValueError(
                "ignore_columns and keep_columns are mutually exclusive; "
                "specify one or the other, not both."
            )
        indexer_columns = params.pop("indexer_columns", [])
        selected_columns = params.pop("selected_columns", [])
        user_ranges = params.pop("ranges", {})
        super().__init__(**params)
        if pn.state._is_pyodide and not self.access_token:
            raise ValueError("An access_token is required in Pyodide environment.")
        if self.access_token:
            BigQuerier._access_token = self.access_token

        self._indexer_params = set()
        self._indexer_watchers: dict[str, param.parameterized.Watcher] = {}
        self._categorical_columns: set[str] = set(categorical_columns or [])
        self._semaphore = asyncio.Semaphore(self.max_concurrent)
        self.columns = self._fetch_columns(
            ignore=ignore_columns or [], keep=keep_columns,
        )
        # Fetch ranges only for indexer columns not already supplied by the user.
        fetch_cols = [c for c in indexer_columns if c not in user_ranges]
        ranges = {**self._fetch_ranges(fetch_cols), **user_ranges}
        
        self.param["selected_columns"].objects = [None] + list(self.columns.keys())
        self.param["indexer_columns"].objects = [None] + list(self.columns.keys())

        self.param.update(
            ranges = ranges,
            selected_columns=selected_columns,
            indexer_columns=indexer_columns,
        )

        self.view = self.make_view()

    def make_view(self):

        indexer_widgets = list()
        for column in self.indexer_columns:
            column_type = self.columns.get(column)
            if column_type in NUMERIC_TYPES and column not in self._categorical_columns:
                widget = pn.widgets.RangeSlider.from_param(
                    self.param[f"indexer_{column}"],
                    name=column,
                )
            elif column_type in TIMESTAMP_TYPES and column not in self._categorical_columns:
                widget = pn.widgets.DateRangeSlider.from_param(
                    self.param[f"indexer_{column}"],
                    name=column,
                )
            else:
                widget = pn.widgets.MultiChoice.from_param(
                    self.param[f"indexer_{column}"],
                    name=column,
                )
            indexer_widgets.append(widget)

        return pn.Row(
            pn.Column(
                pn.widgets.MultiChoice.from_param(self.param["selected_columns"]),
                pn.widgets.MultiChoice.from_param(self.param["indexer_columns"]),
                pn.indicators.LoadingSpinner.from_param(
                    self.param["is_loading"],
                    name="Test",
                    size=22,
                    styles={"margin-left": "-31px"},
                ),
                pn.indicators.Number.from_param(self.param["nqueued"]),
                pn.widgets.Tabulator.from_param(self.param["data"], pagination="remote", page_size=10),
            ),
            pn.Column(
                *indexer_widgets
            ),
            sizing_mode="stretch_width",
        )

    @param.depends("indexer_columns", watch=True)
    def _on_indexer_columns_changed(self, param_name_format="indexer_{name}"):
        """Add/remove Range or DateRange parameters based on indexer column selection."""
        current = {c for c in self.indexer_columns if c is not None}
        previous = self._indexer_params

        # Remove parameters for deselected columns
        cls = type(self)
        for name in previous - current:
            formatted_name = param_name_format.format(name=name)
            if formatted_name in self._indexer_watchers:
                self.param.unwatch(self._indexer_watchers.pop(formatted_name))
            if hasattr(cls, formatted_name):
                delattr(cls, formatted_name)
            self._indexer_params.discard(name)

        # Add parameters for newly selected columns
        for name in current - previous:
            formatted_name = param_name_format.format(name=name)
            col_type = self.columns.get(name)
            bounds = self.ranges.get(name)
            if (
                col_type in NUMERIC_TYPES
                and bounds is not None
                and name not in self._categorical_columns
            ):
                self.param.add_parameter(
                    formatted_name, param.Range(default=bounds, bounds=bounds, allow_refs=True)
                )
            elif (
                col_type in TIMESTAMP_TYPES
                and bounds is not None
                and name not in self._categorical_columns
            ):
                self.param.add_parameter(
                    formatted_name, param.DateRange(default=bounds, bounds=bounds, allow_refs=True)
                )
            else:
                if name in self.ranges:
                    unique_vals = self.ranges[name]
                else:
                    unique_vals = self._fetch_unique_values(name)
                self.param.add_parameter(
                    formatted_name,
                    param.ListSelector(default=[], objects=unique_vals, allow_refs=True),
                )
            self._indexer_watchers[formatted_name] = self.param.watch(
                self._on_indexer_value_changed, [formatted_name]
            )
            self._indexer_params.add(name)

    async def _on_indexer_value_changed(self, event: param.parameterized.Event) -> None:
        """Re-fetch selected columns when an indexer parameter value changes."""
        await self._on_selected_columns_changed()

    def _build_where_clause(
        self,
        param_name_format: str = "indexer_{name}",
        overrides: dict | None = None,
    ) -> str | None:
        """Build a SQL WHERE clause from current indexer parameter values.

        Parameters
        ----------
        param_name_format : str
            Format string mapping column names to parameter names.
        overrides : dict, optional
            Mapping of column name to override value.  For numeric/timestamp
            columns the value is a ``(lo, hi)`` tuple; for categorical
            columns it is a list of values.  Overridden range columns use
            strict ``<`` for the upper bound (chunk semantics).

        Returns
        -------
        str or None
            A ``" WHERE ..."`` string, ``""`` if no filters are active,
            or ``None`` if the filters guarantee zero results.
        """
        conditions: list[str] = []
        for name in self._indexer_params:
            is_overridden = overrides is not None and name in overrides
            if is_overridden:
                value = overrides[name]
            else:
                formatted_name = param_name_format.format(name=name)
                value = getattr(self, formatted_name, None)
            if value is None:
                continue
            col_type = self.columns.get(name)
            # Use strict '<' for upper bound on chunk overrides
            upper_op = "<" if is_overridden else "<="
            if col_type in NUMERIC_TYPES and name not in self._categorical_columns:
                lo, hi = value
                if lo is not None:
                    conditions.append(f"`{name}` >= {lo}")
                if hi is not None:
                    conditions.append(f"`{name}` {upper_op} {hi}")
            elif col_type in TIMESTAMP_TYPES and name not in self._categorical_columns:
                lo, hi = value
                if lo is not None:
                    conditions.append(f"`{name}` >= TIMESTAMP('{lo.isoformat()}')")
                if hi is not None:
                    conditions.append(f"`{name}` {upper_op} TIMESTAMP('{hi.isoformat()}')")
            else:
                if not value:
                    return None
                elif col_type in NUMERIC_TYPES:
                    joined = ", ".join(str(v) for v in value)
                    conditions.append(f"`{name}` IN ({joined})")
                else:
                    joined = ", ".join(f"'{v}'" for v in value)
                    conditions.append(f"`{name}` IN ({joined})")
        if not conditions:
            return ""
        return " WHERE " + " AND ".join(conditions)

    def _compute_chunks(self, column: str, chunk_size) -> list | None:
        """Compute aligned chunk boundaries for a given indexer column.

        Parameters
        ----------
        column : str
            Name of the indexer column.
        chunk_size
            For numeric columns: a number (same unit as the column values).
            For timestamp columns: a number of days or a ``timedelta``.
            For categorical columns: an integer count of values per chunk.

        Returns
        -------
        list or None
            List of ``(lo, hi)`` tuples for numeric/timestamp columns,
            list of value-lists for categorical columns, or ``None`` if the
            indexer has no current value.
        """
        formatted_name = f"indexer_{column}"
        value = getattr(self, formatted_name, None)
        if value is None:
            return None

        col_type = self.columns.get(column)

        if col_type in NUMERIC_TYPES and column not in self._categorical_columns:
            lo, hi = value
            if lo is None or hi is None:
                return None
            aligned_lo = math.floor(lo / chunk_size) * chunk_size
            aligned_hi = math.ceil(hi / chunk_size) * chunk_size
            if aligned_hi <= hi:
                aligned_hi += chunk_size
            chunks = []
            current = aligned_lo
            while current < aligned_hi:
                chunks.append((current, current + chunk_size))
                current += chunk_size
            return chunks

        elif col_type in TIMESTAMP_TYPES and column not in self._categorical_columns:
            lo, hi = value
            if lo is None or hi is None:
                return None
            if isinstance(chunk_size, (int, float)):
                chunk_delta = timedelta(days=chunk_size)
            else:
                chunk_delta = chunk_size
            # Align to chunk boundaries from epoch
            epoch = datetime(1970, 1, 1, tzinfo=lo.tzinfo)
            total_seconds = chunk_delta.total_seconds()
            lo_offset = (lo - epoch).total_seconds()
            aligned_lo_seconds = math.floor(lo_offset / total_seconds) * total_seconds
            aligned_lo = epoch + timedelta(seconds=aligned_lo_seconds)
            chunks = []
            current = aligned_lo
            while current < hi:
                next_val = current + chunk_delta
                chunks.append((current, next_val))
                current = next_val
            return chunks or None

        else:
            # Categorical: chunk by number of items
            if not value:
                return None
            chunks = []
            for i in range(0, len(value), chunk_size):
                chunks.append(value[i:i + chunk_size])
            return chunks or None

    @param.depends("selected_columns", watch=True)
    async def _on_selected_columns_changed(self):
        """Fetch each selected column individually (for effective caching) and combine.

        When ``index_chunker`` contains a key matching an indexer column, that
        column's range is split into aligned chunks of the given size.  Each
        chunk is queried separately so that results are cached independently,
        and results are concatenated and post-filtered to the actual selected
        range.
        """
        cols = [c for c in self.selected_columns if c is not None]
        if not cols:
            self.data = None
            return

        all_cols = cols + [c for c in self._indexer_params if c not in cols]

        # Determine which indexer columns need chunking
        chunk_specs: dict[str, list] = {}
        if self.index_chunker:
            for col_name in self._indexer_params:
                if col_name in self.index_chunker:
                    chunks = self._compute_chunks(
                        col_name, self.index_chunker[col_name]
                    )
                    if chunks is not None:
                        chunk_specs[col_name] = chunks

        if not chunk_specs:
            # Original behaviour – no chunking
            where = self._build_where_clause()
            if where is None:
                self.data = pd.DataFrame(columns=all_cols)
                return
            dfs = await asyncio.gather(
                *(
                    self.query(f"SELECT `{c}` FROM `{self.full_table_id}`{where}")
                    for c in all_cols
                )
            )
            self.param.update(data=pd.concat(dfs, axis=1))
        else:
            # Query each chunk combination separately
            chunk_columns = list(chunk_specs.keys())
            chunk_lists = [chunk_specs[c] for c in chunk_columns]
            all_combos = list(itertools.product(*chunk_lists))

            all_dfs: list[pd.DataFrame] = []
            for combo in all_combos:
                overrides = dict(zip(chunk_columns, combo))

                # Shortcut: categorical single-value chunks whose column is
                # being selected don't need a query — the result is trivially
                # the value itself.  We fill them in after querying the rest.
                shortcut_cols: dict[str, object] = {}
                for col_name, chunk_val in overrides.items():
                    col_type = self.columns.get(col_name)
                    is_categorical = (
                        col_type not in NUMERIC_TYPES
                        and col_type not in TIMESTAMP_TYPES
                    ) or col_name in self._categorical_columns
                    if (
                        is_categorical
                        and isinstance(chunk_val, list)
                        and len(chunk_val) == 1
                        and col_name in all_cols
                    ):
                        shortcut_cols[col_name] = chunk_val[0]

                query_cols = [c for c in all_cols if c not in shortcut_cols]

                where = self._build_where_clause(overrides=overrides)
                if where is None:
                    continue

                if query_cols:
                    dfs = await asyncio.gather(
                        *(
                            self.query(
                                f"SELECT `{c}` FROM `{self.full_table_id}`{where}"
                            )
                            for c in query_cols
                        )
                    )
                    chunk_df = pd.concat(dfs, axis=1)
                else:
                    # Every column is shortcutted — we still need the row
                    # count, so run a COUNT(*) query.
                    count_df = await self.query(
                        f"SELECT COUNT(*) AS _n FROM `{self.full_table_id}`{where}"
                    )
                    n_rows = int(count_df.iloc[0, 0])
                    chunk_df = pd.DataFrame(index=range(n_rows))

                # Fill in shortcutted columns
                for col_name, val in shortcut_cols.items():
                    chunk_df[col_name] = val

                if not chunk_df.empty:
                    all_dfs.append(chunk_df)

            if not all_dfs:
                self.param.update(data=pd.DataFrame(columns=all_cols))
                return

            result = pd.concat(all_dfs, ignore_index=True)

            # Post-filter to the actual selected range (chunks may overshoot)
            for col_name in chunk_columns:
                actual = getattr(self, f"indexer_{col_name}", None)
                if actual is None:
                    continue
                col_type = self.columns.get(col_name)
                if (
                    col_type in (NUMERIC_TYPES | TIMESTAMP_TYPES)
                    and col_name not in self._categorical_columns
                ):
                    lo, hi = actual
                    if lo is not None:
                        result = result[result[col_name] >= lo]
                    if hi is not None:
                        result = result[result[col_name] <= hi]

            self.param.update(data=result.reset_index(drop=True))

    @property
    def full_table_id(self) -> str:
        """Fully qualified BigQuery table ID: ``data_project.dataset.table``."""
        return f"{self.data_project}.{self.dataset}.{self.table}"

    @staticmethod
    @pn.cache(
        to_disk=CACHE_TO_DISK,
        cache_path="./big_querier_cache",
        per_session=False,
        ttl=5 * 86400,
        max_items={True: None, False:0}[CACHE],
    )
    def _request(method: str, url: str, body: Optional[dict] = None) -> dict:
        """Make a synchronous HTTP request to the BigQuery REST API.

        Token is fetched internally so it does not affect the cache key.

        Parameters
        ----------
        method : str
            HTTP method (e.g. ``"GET"``, ``"POST"``).
        url : str
            Full request URL.
        body : dict, optional
            JSON-serializable request body.

        Returns
        -------
        dict
            Parsed JSON response.
        """
        print(f"Running {body}")
        token = BigQuerier._get_access_token()
        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }
        data = json.dumps(body).encode() if body else None
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read())

    @staticmethod
    @pn.cache(
        to_disk=CACHE_TO_DISK, cache_path="./big_querier_cache", per_session=False, max_items={True: None, False:0}[CACHE],
    )
    async def _async_request(
        method: str, url: str, body: Optional[dict] = None
    ) -> dict:
        """Make an asynchronous HTTP request to the BigQuery REST API.

        Token is fetched internally so it does not affect the cache key.

        Parameters
        ----------
        method : str
            HTTP method (e.g. ``"GET"``, ``"POST"``).
        url : str
            Full request URL.
        body : dict, optional
            JSON-serializable request body.

        Returns
        -------
        dict
            Parsed JSON response.
        """
        print(f"Running {body}")
        if pn.state._is_pyodide:
            token = BigQuerier._get_access_token()
        else:
            token = await asyncio.to_thread(BigQuerier._get_access_token)
        headers = {"Authorization": f"Bearer {token}"}
        if pn.state._is_pyodide:
            from pyodide.http import (  # pyright: ignore [reportMissingImports] # noqa
                pyfetch,
            )

            kwargs = {"method": method, "headers": headers}
            if body is not None:
                kwargs["body"] = json.dumps(body)
            resp = await pyfetch(url, **kwargs)
            return await resp.json()
        else:
            import aiohttp
            async with aiohttp.ClientSession() as session:
                req = getattr(session, method.lower())
                async with req(url, json=body, headers=headers) as resp:
                    resp.raise_for_status()
                    return await resp.json()

    def _fetch_columns(
        self,
        ignore: list[str] | None = None,
        keep: list[str] | None = None,
    ) -> dict[str, str]:
        """Fetch column names and BigQuery types via the tables.get API.

        Parameters
        ----------
        ignore : list[str], optional
            Column names to exclude from the result.
        keep : list[str], optional
            Column names to include; all others are excluded.

        Returns
        -------
        dict[str, str]
            Mapping of column name to BigQuery type string (e.g. ``"FLOAT"``).
        """
        skip = set(ignore or [])
        allow = set(keep) if keep is not None else None
        url = (
            f"{BQ_BASE_URL}/projects/{self.data_project}"
            f"/datasets/{self.dataset}/tables/{self.table}"
        )
        response = self._request("GET", url)
        return {
            f["name"]: f["type"]
            for f in response["schema"]["fields"]
            if f["name"] not in skip
            and (allow is None or f["name"] in allow)
        }

    def _fetch_ranges(
        self,
        only: list[str] | None = None,
    ) -> dict[
        str, tuple[float | None, float | None] | tuple[datetime | None, datetime | None]
    ]:
        """Fetch (min, max) ranges for numeric and timestamp columns via a SQL query.

        Parameters
        ----------
        only : list[str], optional
            If given, only fetch ranges for these column names.  Columns not
            present in the table or not numeric/timestamp are silently skipped.

        Returns
        -------
        dict[str, tuple]
            Mapping of column name to ``(min, max)`` tuple. Values are
            ``float`` for numeric columns and ``datetime`` for timestamp columns.
        """
        allow = set(only) if only is not None else None
        numeric_cols = [
            name for name, dtype in self.columns.items()
            if dtype in NUMERIC_TYPES
            and name not in self._categorical_columns
            and (allow is None or name in allow)
        ]
        timestamp_cols = [
            name for name, dtype in self.columns.items()
            if dtype in TIMESTAMP_TYPES
            and name not in self._categorical_columns
            and (allow is None or name in allow)
        ]
        range_cols = numeric_cols + timestamp_cols
        if not range_cols:
            return {}

        agg_parts = []
        for col in range_cols:
            agg_parts.append(f"MIN(`{col}`) AS `{col}__min`")
            agg_parts.append(f"MAX(`{col}`) AS `{col}__max`")

        query = f"SELECT {', '.join(agg_parts)} FROM `{self.full_table_id}`"
        url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
        response = self._request("POST", url, {"query": query, "useLegacySql": False})

        row = response["rows"][0]["f"]
        ranges = {}
        for i, col in enumerate(range_cols):
            min_val = row[i * 2]["v"]
            max_val = row[i * 2 + 1]["v"]
            if col in timestamp_cols:
                ranges[col] = (
                    (
                        datetime.fromtimestamp(float(min_val), tz=timezone.utc)
                        if min_val is not None
                        else None
                    ),
                    (
                        datetime.fromtimestamp(float(max_val), tz=timezone.utc)
                        if max_val is not None
                        else None
                    ),
                )
            else:
                ranges[col] = (
                    float(min_val) if min_val is not None else None,
                    float(max_val) if max_val is not None else None,
                )
        return ranges

    def _fetch_unique_values(self, column: str) -> list[str]:
        """Fetch unique values for a column via a synchronous SQL query.

        Parameters
        ----------
        column : str
            Column name to query distinct values for.

        Returns
        -------
        list[str]
            Sorted list of unique values.
        """
        self.nqueued += 1
        self.is_loading = True
        try:
            query = f"SELECT DISTINCT `{column}` FROM `{self.full_table_id}` ORDER BY `{column}`"
            url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
            response = BigQuerier._request(
                "POST", url, {"query": query, "useLegacySql": False}
            )
            if "rows" not in response:
                return []
            return [
                row["f"][0]["v"]
                for row in response["rows"]
                if row["f"][0]["v"] is not None
            ]
        finally:
            self.nqueued -= 1
            self.is_loading = self.nqueued > 0

    async def _async_fetch_unique_values(self, column: str) -> list[str]:
        """Fetch unique values for a column via an async SQL query.

        Parameters
        ----------
        column : str
            Column name to query distinct values for.

        Returns
        -------
        list[str]
            Sorted list of unique values.
        """
        self.nqueued += 1
        self.is_loading = True
        try:
            async with self._semaphore:
                query = f"SELECT DISTINCT `{column}` FROM `{self.full_table_id}` ORDER BY `{column}`"
                url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
                response = await self._async_request(
                    "POST", url, {"query": query, "useLegacySql": False}
                )
            if "rows" not in response:
                return []
            return [
                row["f"][0]["v"]
                for row in response["rows"]
                if row["f"][0]["v"] is not None
            ]
        finally:
            self.nqueued -= 1
            self.is_loading = self.nqueued > 0

    @staticmethod
    def _get_access_token() -> str:
        """Get an access token using Application Default Credentials.

        In Pyodide, supply a token via ``BigQuerier(access_token="...")``.

        Returns
        -------
        str
            OAuth 2.0 bearer token.
        """
        if BigQuerier._access_token:
            return BigQuerier._access_token
        import google.auth
        import google.auth.transport.requests
        credentials, _ = google.auth.default(
            scopes=["https://www.googleapis.com/auth/bigquery"]
        )
        credentials.refresh(google.auth.transport.requests.Request())
        return credentials.token

    @staticmethod
    def _rows_to_dataframe(response: dict) -> pd.DataFrame:
        """Convert a BigQuery query response JSON to a pandas DataFrame.

        Parameters
        ----------
        response : dict
            Raw JSON response from the BigQuery queries API.

        Returns
        -------
        pd.DataFrame
            DataFrame with column names from the response schema.
        """
        fields = response["schema"]["fields"]
        col_names = [f["name"] for f in fields]

        if "rows" not in response:
            return pd.DataFrame(columns=col_names)

        rows = [[cell["v"] for cell in row["f"]] for row in response["rows"]]
        df = pd.DataFrame(rows, columns=col_names)
        for field in fields:
            col = field["name"]
            if field["type"] in TIMESTAMP_TYPES:
                df[col] = pd.to_datetime(df[col].astype(float), unit="s")
            elif field["type"] in NUMERIC_TYPES:
                df[col] = pd.to_numeric(df[col])
        return df

    async def query(self, query: str) -> pd.DataFrame:
        """Execute an async SQL query and return the results as a DataFrame.

        Parameters
        ----------
        query : str
            SQL query string (Standard SQL).

        Returns
        -------
        pd.DataFrame
            Query results. Also stored in ``self.data``.
        """
        self.nqueued += 1
        self.is_loading = True
        try:
            async with self._semaphore:
                url = f"{BQ_BASE_URL}/projects/{self.job_project}/queries"
                data = await self._async_request(
                    "POST", url, {"query": query, "useLegacySql": False}
                )
            data_ = self._rows_to_dataframe(data)
            return data_
        finally:
            self.nqueued -= 1
            self.is_loading = self.nqueued > 0
