from typing import Dict

import pandas as pd
import pandas_gbq
from santoku.aws import SecretsManagerHandler

from google.oauth2 import service_account


class BigQueryHandler:
    """
    Manage Google BigQuery interactions, the simplest of which is to query a particular table.
    This makes use of the pandas-gbq API. More information at
    https://github.com/googleapis/python-bigquery-pandas

    """

    def __init__(self, credential_info: Dict[str, str]) -> None:
        """
        Base constructor that performs authentication against Google Cloud Platform.

        Parameters
        ----------
        credential_info : Dict[str, str]
            Credentials for a Google Cloud service account in dictionary form.

        Notes
        -----
        The retrieved secret must have the particular structure of the JSON file generated by the
        Google Cloud Console, which is associated with a service account. It looks like this:
        ```
        {
            "type": "service_account",
            "project_id": "<project name>",
            "private_key_id": "<private key id>",
            "private_key": "<private key>",
            "client_email": "<service account email>",
            "client_id": "<client id>",
            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
            "token_uri": "https://accounts.google.com/o/oauth2/token",
            "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
            "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/<email>"
        }
        ```
        See [1] for instructions on how to generate such file.

        More on authentication via service account in [2].
        To create a Credentials object, follow the google-auth guide in [3].

        References
        ----------
        [1] :
        https://cloud.google.com/iam/docs/creating-managing-service-account-keys#creating_service_account_keys
        [2] :
        https://cloud.google.com/docs/authentication/production#obtaining_and_providing_service_account_credentials_manually
        [3] :
        https://google-auth.readthedocs.io/en/latest/user-guide.html#service-account-private-key-files
        """

        credentials = service_account.Credentials.from_service_account_info(info=credential_info)

        pandas_gbq.context.credentials = credentials
        pandas_gbq.context.project = credential_info["project_id"]

    @classmethod
    def from_aws_secrets_manager(cls, secret_name: str):
        """
        Retrieve the necessary information to generate service account credentials from AWS Secrets
        Manager. Requires that AWS credentials with the appropriate permissions are located
        somewhere on the AWS credential chain in the local machine.

        Paramaters
        ----------
        secret_name : str
            Name or ARN for the secret containing the JSON needed for BigQuery authentication.

        """
        secrets_manager = SecretsManagerHandler()
        credential_info = secrets_manager.get_secret_value(secret_name=secret_name)
        return cls(credential_info=credential_info)

    def get_query_results_as_dataframe(
        self,
        query: str,
        use_bqstorage_api: bool = True,
        **kwargs,
    ) -> pd.DataFrame:
        """
        Run a SQL query and return the results as a pandas dataframe.

        Parameters
        ----------
        query : str
            SQL query to be executed. Using the standard SQL dialect by default.
        use_bqstorage_api: bool
            If set, use the BigQuery Storage API to download results more quickly.
        kwargs : dict
            Additional arguments for Pandas read_gbq method.

        Notes
        -----
        More on additional arguments for the read_gbq method: [1].

        References
        ----------
        [1] :
        https://pandas.pydata.org/docs/reference/api/pandas.read_gbq.html

        Returns
        -------
        pd.DataFrame
            The query result as a Pandas DataFrame.
        """

        return pd.read_gbq(query=query, use_bqstorage_api=use_bqstorage_api, **kwargs)
