"""
BigQuery Data Warehouse classes
"""

import logging
import os
from typing import List, Optional, Tuple, Union

import json
import pandas as pd

from rasgoql.data.base import DataWarehouse, DWCredentials
from rasgoql.errors import (
    DWCredentialsWarning,
    DWConnectionError,
    DWQueryError,
    PackageDependencyWarning,
    SQLWarning,
    TableAccessError,
    TableConflictException,
)
from rasgoql.imports import bq, gcp_exc, gcp_flow, gcp_svc
from rasgoql.primitives.enums import check_response_type, check_write_method, check_write_table_type
from rasgoql.utils.creds import load_env, save_env
from rasgoql.utils.df import cleanse_sql_dataframe
from rasgoql.utils.messaging import verbose_message
from rasgoql.utils.sql import is_scary_sql

logging.basicConfig()
logger = logging.getLogger('BigQuery DataWarehouse')
logger.setLevel(logging.INFO)


class BigQueryCredentials(DWCredentials):
    """
    BigQuery Credentials
    """

    dw_type = 'bigquery'

    def __init__(
        self,
        json_filepath: str,
        project: str = None,
        dataset: str = None,
    ):
        self.json_filepath = json_filepath
        self.project = project
        self.dataset = dataset

    def __repr__(self) -> str:
        return json.dumps(
            {
                "json_filepath": self.json_filepath,
                "project": self.project,
                "dataset": self.dataset,
            }
        )

    @classmethod
    def from_env(
        cls,
        filepath: str = None,
    ) -> 'BigQueryCredentials':
        """
        Creates an instance of this Class from a .env file on your machine
        """
        load_env(filepath)
        json_filepath = os.getenv('BIGQUERY_JSON_FILEPATH', os.getenv('GOOGLE_APPLICATION_CREDENTIALS'))
        project = os.getenv('BIGQUERY_PROJECT')
        dataset = os.getenv('BIGQUERY_DATASET')
        if not all([json_filepath, project, dataset]):
            raise DWCredentialsWarning(
                'Your env file is missing expected credentials. Consider running '
                'BigQueryCredentials(*args).to_env() to repair this.'
            )
        return cls(
            json_filepath,
            project,
            dataset,
        )

    def to_dict(self) -> dict:
        """
        Returns a dict of the credentials
        """
        return {
            "json_filepath": self.json_filepath,
            "project": self.project,
            "dataset": self.dataset,
        }

    def to_env(
        self,
        filepath: str = None,
        overwrite: bool = False,
    ):
        """
        Saves credentials to a .env file on your machine
        """
        creds = {
            "BIGQUERY_JSON_FILEPATH": self.json_filepath,
            "BIGQUERY_PROJECT": self.project,
            "BIGQUERY_DATASET": self.dataset,
        }
        return save_env(creds, filepath, overwrite)


class BigQueryDataWarehouse(DataWarehouse):
    """
    Google BigQuery DataWarehouse
    """

    dw_type = 'bigquery'
    credentials_class = BigQueryCredentials

    def __init__(self):
        if bq is None:
            raise PackageDependencyWarning(
                'Missing a required python package to run BigQuery. '
                'Please download the BigQuery package by running: '
                'pip install rasgoql[bigquery]'
            )

        super().__init__()
        self.credentials: dict = None
        self.connection: bq.Client = None
        self.default_project = None
        self.default_dataset = None

    # ---------------------------
    # Core Data Warehouse methods
    # ---------------------------
    def change_namespace(
        self,
        namespace: str,
    ) -> None:
        """
        Changes the default namespace of your connection

        Params:
        `namespace`: str:
            namespace (project.dataset)
        """
        namespace = self.validate_namespace(namespace)
        project, dataset = self.parse_namespace(namespace)
        try:
            self.execute_query(f'USE PROJECT {project}')
            self.execute_query(f'USE DATASET {dataset}')
            self.default_namespace = namespace
            self.default_project = project
            self.default_dataset = dataset
            verbose_message(
                f"Namespace reset to {self.default_namespace}",
                logger,
            )
        except Exception as e:
            self._error_handler(e)

    def connect(
        self,
        credentials: Union[dict, BigQueryCredentials],
    ):
        """
        Connect to BigQuery

        Params:
        `credentials`: dict:
            dict (or DWCredentials class) holding the connection credentials
        """
        if isinstance(credentials, BigQueryCredentials):
            credentials = credentials.to_dict()

        try:
            self.default_project = credentials.get('project')
            self.default_dataset = credentials.get('dataset')
            if not self.credentials:
                self.credentials = self._get_credentials(credentials.get('json_filepath'))
            self.connection = bq.Client(
                credentials=self.credentials,
                project=self.default_project,
            )
            verbose_message(
                "Connected to BigQuery",
                logger,
            )
        except Exception as e:
            self._error_handler(e)

    def close_connection(self):
        """
        Close connection to BigQuery
        """
        try:
            if self.connection:
                self.connection.close()
            self.connection = None
            verbose_message(
                "Connection to BigQuery closed",
                logger,
            )
        except Exception as e:
            self._error_handler(e)

    def create(
        self,
        sql: str,
        fqtn: str,
        table_type: str = 'VIEW',
        overwrite: bool = False,
    ):
        """
        Create a view or table from given SQL

        Params:
        `sql`: str:
            query that returns data (i.e. can be wrapped in a CREATE TABLE statement)
        `fqtn`: str:
            Fully-qualified table name (database.schema.table)
            Name for the new table
        `table_type`: str:
            One of values: [view, table]
        `overwrite`: bool
            pass True when this table name already exists in your DataWarehouse
            and you know you want to overwrite it
            WARNING: This will completely overwrite data in the existing table
        """
        table_type = check_write_table_type(table_type)
        fqtn = self.magic_fqtn_handler(fqtn, self.default_namespace)
        if self._table_exists(fqtn) and not overwrite:
            msg = (
                f'A table or view named {fqtn} already exists. '
                'If you are sure you want to overwrite it, '
                'pass in overwrite=True and run this function again'
            )
            raise TableConflictException(msg)
        query = f'CREATE OR REPLACE {table_type} {fqtn} AS {sql}'
        self.execute_query(query, acknowledge_risk=True, response='None')
        return fqtn

    @property
    def default_namespace(self) -> str:
        """
        Returns the default project.dataset of this connection
        """
        return f'{self.default_project}.{self.default_dataset}'

    @default_namespace.setter
    def default_namespace(
        self,
        new_namespace: str,
    ):
        """
        Setter method for the `default_namespace` property
        """
        namespace = self.validate_namespace(new_namespace)
        project, dataset = self.parse_namespace(namespace)
        self.default_project = project
        self.default_dataset = dataset

    def execute_query(
        self,
        sql: str,
        response: str = 'tuple',
        acknowledge_risk: bool = False,
        **kwargs,
    ):
        """
        Run a query against BigQuery and return all results

        `sql`: str:
            query text to execute
        `response`: str:
            Possible values: [dict, df, None]
        `acknowledge_risk`: bool:
            pass True when you know your SQL statement contains
            a potentially dangerous or data-altering operation
            and still want to run it against your DataWarehouse
        """
        response = check_response_type(response)
        if is_scary_sql(sql) and not acknowledge_risk:
            msg = (
                'It looks like your SQL statement contains a '
                'potentially dangerous or data-altering operation.'
                'If you are positive you want to run this, '
                'pass in acknowledge_risk=True and run this function again.'
            )
            raise SQLWarning(msg)
        verbose_message(
            f"Executing query: {sql}",
            logger,
        )
        if response == 'DICT':
            return self._query_into_dict(sql)
        if response == 'DF':
            return self._query_into_pandas(sql)
        return self._execute_string(sql, ignore_results=(response == 'NONE'))

    def get_ddl(
        self,
        fqtn: str,
    ) -> str:
        """
        Returns the create statement for a table or view

        `fqtn`: str:
            Fully-qualified Table Name (database.schema.table)
        """
        fqtn = self.magic_fqtn_handler(fqtn, self.default_namespace)
        proj, ds, tbl = self.parse_fqtn(fqtn)
        sql = f"SELECT DDL FROM {proj}.{ds}.INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME='{tbl}'"
        query_response = self.execute_query(sql)
        return query_response[0]

    def get_object_details(
        self,
        fqtn: str,
    ) -> tuple:
        """
        Return details of a table or view in BigQuery

        Params:
        `fqtn`: str:
            Fully-qualified table name (database.schema.table)

        Response:
            object exists: bool
            is rasgo object: bool
            object type: [table|view|unknown]
        """
        fqtn = self.magic_fqtn_handler(fqtn, self.default_namespace)
        obj_exists = False
        is_rasgo_obj = False
        obj_type = 'unknown'
        try:
            table = self.connection.get_table(fqtn)
            obj_exists = True
            obj_type = table.table_type
            if table.labels:
                for _label, value in table.labels.items():
                    if value == 'rasgoql':
                        is_rasgo_obj = True
            return obj_exists, is_rasgo_obj, obj_type
        except gcp_exc.NotFound:
            return obj_exists, is_rasgo_obj, obj_type

    def get_schema(
        self,
        fqtn_or_sql: str,
    ) -> List[Tuple[str, str]]:
        """
        Return the schema of a table, view, or select statement

        Params:
        `fqtn_or_sql`: str:
            Either a Fully-qualified table name (database.schema.table)
            or a SQL select statement that will create a view.
        """
        # Check for SQL
        if self._is_select_statement(fqtn_or_sql):
            query_job = self.connection.query(
                query=fqtn_or_sql,
                job_config=bq.QueryJobConfig(dry_run=True),
            )
            schema = query_job._properties['statistics']['query']['schema']['fields']
            return [(schema_field['name'], schema_field['type']) for schema_field in schema]
        # Otherwise assume fqtn:
        fqtn = self.magic_fqtn_handler(fqtn_or_sql, self.default_namespace)
        if self._table_exists(fqtn):
            table = self.connection.get_table(fqtn)
            schema = table.schema
            return [(schema_field.name, schema_field.field_type) for schema_field in schema]
        else:
            raise TableAccessError(f'Table {fqtn} does not exist or cannot be accessed.')

    def list_tables(
        self,
        database: str = None,
        schema: str = None,
    ) -> pd.DataFrame:
        """
        List all tables and views available in default namespace

        Params:
        `database`: str:
            override database
        `schema`: str:
            override schema
        """
        project = database or self.default_project
        dataset = schema or self.default_dataset
        namespace = f'{project}.{dataset}'
        try:
            tables = self.connection.list_tables(namespace)
            records = []
            columns = [
                'TABLE_NAME',
                'FQTN',
                'TABLE_TYPE',
                'ROW_COUNT',
                'CREATED',
                'LAST_ALTERED',
            ]
            for tbl in tables:
                table = self.connection.get_table(tbl)
                records.append(
                    (
                        table.table_id,
                        f'{table.project}.{table.dataset_id}.{table.table_id}',
                        table.table_type,
                        table.num_rows,
                        table.created,
                        table.modified,
                    )
                )
            return pd.DataFrame(
                records,
                columns=columns,
            )
        except Exception as e:
            self._error_handler(e)

    def preview(
        self,
        sql: str,
        limit: int = 10,
    ) -> pd.DataFrame:
        """
        Returns 10 records into a pandas DataFrame

        Params:
        `sql`: str:
            SQL statment to run
        `limit`: int:
            Records to return
        """
        return self.execute_query(
            f'{sql} LIMIT {limit}',
            response='df',
            acknowledge_risk=True,
        )

    def save_df(
        self,
        df: pd.DataFrame,
        fqtn: str,
        method: str = None,
    ) -> str:
        """
        Creates a table in BigQuery from a pandas Dataframe

        Params:
        `df`: pandas DataFrame:
            DataFrame to upload
        `fqtn`: str:
            Fully-qualied table name (database.schema.table)
            Name for the new table
        `overwrite`: bool
            pass True when this table name already exists in your DataWarehouse
            and you know you want to overwrite it
            WARNING: This will completely overwrite data in the existing table
        """
        if method:
            method = check_write_method(method)
        fqtn = self.magic_fqtn_handler(fqtn, self.default_namespace)
        table_exists = self._table_exists(fqtn)
        if table_exists and not method:
            msg = (
                f"A table named {fqtn} already exists. "
                "If you are sure you want to write over it, pass in "
                "method='append' or method='replace' and run this function again"
            )
            raise TableConflictException(msg)
        try:
            cleanse_sql_dataframe(df)
            job_config = bq.LoadJobConfig(write_disposition='WRITE_TRUNCATE' if method == 'REPLACE' else None)
            _job = self.connection.load_table_from_dataframe(
                df,
                fqtn,
                job_config=job_config,
            )
            # Wait for the job to complete
            _job.result()
            return fqtn
        except Exception as e:
            self._error_handler(e)

    # ---------------------------
    # Core Data Warehouse helpers
    # ---------------------------
    def _table_exists(
        self,
        fqtn: str,
    ) -> bool:
        """
        Check for existence of fqtn in the Data Warehouse and return a boolean

        Params:
        `fqtn`: str:
            Fully-qualified table name (database.schema.table)
        """
        fqtn = self.magic_fqtn_handler(fqtn, self.default_namespace)
        try:
            self.connection.get_table(fqtn)
            return True
        except gcp_exc.NotFound:
            return False

    # --------------------------
    # BigQuery specific helpers
    # --------------------------
    @property
    def _default_job_config(self) -> 'bq.QueryJobConfig':
        return bq.QueryJobConfig(default_dataset=self.default_namespace)

    def _error_handler(
        self,
        exception: Exception,
        query: str = None,
    ) -> None:
        """
        Handle Snowflake exceptions that need additional info
        """
        verbose_message(f'Exception occurred while running query: {query}', logger),
        if exception is None:
            return
        if isinstance(exception, gcp_exc.NotFound):
            raise TableAccessError('Table does not exist') from exception
        if isinstance(exception, gcp_exc.BadRequest):
            if {"reason": "invalidQuery"} in exception.errors:
                raise DWQueryError(exception.errors) from exception
        if isinstance(exception, gcp_exc.ServiceUnavailable):
            raise DWConnectionError(
                'BigQuery is unavailable. Please check that your are using '
                'valid credentials, that you have internet access, and '
                'that http connections to GCP are whitelisted in your env. '
                'Finally check https://status.cloud.google.com/incident/bigquery '
                'for outage status.'
            ) from exception
        raise exception

    def _execute_string(
        self,
        query: str,
        ignore_results: bool = False,
    ) -> Optional[List[tuple]]:
        """
        Execute a query string against the DataWarehouse connection and fetch all results
        """
        try:
            query_job = self.connection.query(
                query,
                job_config=self._default_job_config,
            )
            if ignore_results:
                return
            return list(query_job.result())
        except Exception as e:
            self._error_handler(e)

    def _get_appflow_credentials(
        self,
        filepath: str,
    ) -> 'bq.Credentials':
        try:
            appflow = gcp_flow.InstalledAppFlow.from_client_secrets_file(
                filepath,
                scopes=["https://www.googleapis.com/auth/bigquery"],
            )
            appflow.run_local_server()
            # appflow.run_console()
            return appflow.credentials
        except Exception as e:
            self._error_handler(e)

    def _get_credentials(
        self,
        filepath: str,
    ) -> 'bq.Credentials':
        try:
            return self._get_appflow_credentials(filepath)
        except ValueError:
            return self._get_service_credentials(filepath)

    def _get_service_credentials(
        self,
        filepath: str,
    ) -> 'bq.Credentials':
        try:
            credentials = gcp_svc.Credentials.from_service_account_file(
                filepath,
                scopes=["https://www.googleapis.com/auth/cloud-platform"],
            )
            return credentials
        except Exception as e:
            self._error_handler(e)

    def _query_into_dict(
        self,
        query: str,
    ) -> dict:
        """
        Return results of query in a pandas DataFrame
        """
        try:
            return (
                self.connection.query(
                    query,
                    job_config=self._default_job_config,
                )
                .result()
                .to_dataframe()
                .to_dict()
            )
        except Exception as e:
            self._error_handler(e)

    def _query_into_pandas(
        self,
        query: str,
    ) -> pd.DataFrame:
        """
        Return results of query in a pandas DataFrame
        """
        try:
            return (
                self.connection.query(
                    query,
                    job_config=self._default_job_config,
                )
                .result()
                .to_dataframe()
            )
        except Exception as e:
            self._error_handler(e)
