"""BigQuery connection helpers."""
from __future__ import annotations

import dataclasses as dc
from typing import Any
from typing import ClassVar

import pydantic.dataclasses as pydc
from pydantic import validate_call

from .base import _DBDatasetArgs
from .base import DatasetBase
from .filter import FilterMixin
from tempora.utils.dataclass import pydc_config
from tempora.utils.db import close_db_connection
from tempora.utils.db import ConnectionParams
from tempora.utils.db import open_db_connection
from tempora.utils.db import resolve_connection_params


@dc.dataclass(frozen=True)
class _BigQueryDatasetArgs(_DBDatasetArgs):
    project: str | None = None


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class BigQueryDataset(DatasetBase, _BigQueryDatasetArgs):
    """Class representing a BigQuery-backed dataset.

    Args:
        database: The name of the BigQuery dataset.
        table: The name of the table under the given dataset.
        db_schema: Optional schema name override.
        project: Optional GCP project identifier.
    """
    filtered_dataset_cls: ClassVar[type[FilteredBigQueryDataset]]

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)


@pydc.dataclass(frozen=True, eq=False, config=pydc_config)
class FilteredBigQueryDataset(FilterMixin, BigQueryDataset):
    """Class representing a filtered BigQuery-based dataset."""

    def __post_init__(self, _op_dict: dict[str, Any] | None) -> None:
        super().__post_init__(_op_dict=_op_dict)


BigQueryDataset.filtered_dataset_cls = FilteredBigQueryDataset


@validate_call
def connect_bigquery(
    project_id: str | None = None,
    auth_token: str | None = None,
) -> None:
    """Open a connection to BigQuery.

    Args:
        project_id: GCP project identifier for the BigQuery connection. If omitted,
            this must be specified using the BIGQUERY_PROJECT_ID env var.
        auth_token: Authentication token for the BigQuery connection. If omitted,
            this must be specified using the BIGQUERY_AUTH_TOKEN env var.

    Returns:
        None
    """
    params = _get_bigquery_connection_params(
        project_id=project_id,
        auth_token=auth_token,
    )
    open_db_connection(dbms_name='BigQuery', params=params)
    return None


def _get_bigquery_connection_params(
    project_id: str | None = None,
    auth_token: str | None = None,
) -> ConnectionParams:
    return resolve_connection_params(
        args={
            'project_id': project_id,
            'auth_token': auth_token,
        },
        env_vars={
            'project_id': 'BIGQUERY_PROJECT_ID',
            'auth_token': 'BIGQUERY_AUTH_TOKEN',
        },
        required_args={'project_id', 'auth_token'},
    )


def disconnect_bigquery() -> None:
    """Close the currently open BigQuery connection."""
    return close_db_connection(dbms_name='BigQuery')
