from typing import Optional

import pandas as pd
from google.cloud import storage

from energy_consumption_forecasting.utils import get_env_var

# Google cloud platform environment variables for project, bucket and service account
GCS_PROJECT = get_env_var(key="GOOGLE_CLOUD_PROJECT")
GCS_BUCKET_NAME = get_env_var(key="GOOGLE_CLOUD_BUCKET_NAME")
GCS_SERVICE_ACCOUNT_JSON_PATH = get_env_var(
    key="GOOGLE_CLOUD_SERVICE_ACCOUNT_JSON_PATH"
)


def get_gcs_bucket(
    project: str = GCS_PROJECT,
    bucket_name: str = GCS_BUCKET_NAME,
    json_credentials_path: str = GCS_SERVICE_ACCOUNT_JSON_PATH,
) -> storage.Bucket:
    """
    This function gets the Google cloud storage(GCS) bucket from GCP, using this bucket
    object you can upload and download data from the Google cloud storage.

    Parameters
    ----------
    project: str
        Connects to the GCS using the project name, by default uses the
        environment variable called GOOGLE_CLOUD_PROJECT.

    bucket_name: str
        Connects to the GCS using the bucket name, by default uses the
        environment variable called GOOGLE_CLOUD_BUCKET_NAME.

    json_credentials_path: str
        Connects to the GCS using the JSON filepath that was generated when an
        service account was created with roles to manipulate the bucket and objects
        in GCS, by default uses the environment variable called
        GCS_SERVICE_ACCOUNT_JSON_PATH.

    Returns
    -------
    storage.Bucket
        A bucket object received from Google cloud storage to upload and download data.
    """
    client = storage.Client.from_service_account_json(
        json_credentials_path=json_credentials_path, project=project
    )

    bucket = client.bucket(bucket_name=bucket_name)

    return bucket


def write_blob_to_bucket(bucket: storage.Bucket, blob_name: str, data: pd.DataFrame):
    """
    This function gets the binary large object(blob) from the Google cloud storage(GCS)
    bucket. Using the blob object, dataframe can be stored as binary file for
    which the data is transformed into apache parquet.

    Parameters
    ----------
    bucket: storage.Bucket
        A bucket object that is generated by GCS, can be used for uploading and
        downloading data.

    blob_name: str
        The name of the blob to be instantiated.

    data: pd.DataFrame
        The data that needs to be stored in the GCS.
    """
    blob_obj = bucket.blob(blob_name=blob_name)

    with blob_obj.open(mode="wb") as file:
        data.to_parquet(path=file)


def read_blob_from_bucket(
    bucket: storage.Bucket, blob_name: str
) -> Optional[pd.DataFrame]:
    """
    This function gets the binary large object(blob) from the Google cloud storage(GCS)
    bucket. Using the blob object, the binary data is transformed back to pandas
    dataframe.

    Parameters
    ----------
    bucket: storage.Bucket
        A bucket object that is generated by GCS, can be used for uploading and
        downloading data.

    blob_name: str
        The name of the blob to be instantiated.

    Returns
    -------
    pd.DataFrame or None
        A dataframe containing the data within the blob or None if the blob does not
        exists.
    """
    blob_obj = bucket.blob(blob_name=blob_name)

    if not blob_obj.exists():
        return None

    with blob_obj.open("rb") as file:
        return pd.read_parquet(path=file)
