from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
from airflow.providers.google.cloud.operators.bigquery import BigQueryExecuteQueryOperator
from airflow.exceptions import AirflowException
from airflow.models.baseoperator import BaseOperator
from airflow.models.skipmixin import SkipMixin
from airflow.utils.decorators import apply_defaults

from awehflow.operators.flow import EventEmittingOperator
from awehflow.utils import utc_now
from typing import Iterable
import time

class BigQueryJobOperator(BigQueryExecuteQueryOperator):
    """Exactly like BigQueryOperator, except that it pushes the job_id to xcom"""
        
    template_fields = ('cleanup_query', 'sql', 'destination_dataset_table', 'labels', "query_params", "impersonation_chain")
    def __init__(self,
                 cleanup_query: str = None,
                 *args,
                 **kwargs):
        super(BigQueryJobOperator, self).__init__(*args, **kwargs)
        self.cleanup_query = cleanup_query 
                

    def execute(self, context):
        if self.hook is None:
            self.hook = BigQueryHook(
                gcp_conn_id=self.gcp_conn_id,
                use_legacy_sql=self.use_legacy_sql,
                delegate_to=self.delegate_to,
                location=self.location,
                impersonation_chain=self.impersonation_chain,
            )
            credential_email = self.hook._get_credentials_email()
            self.log.info(f"Created BigQueryHook with account: {credential_email}")

        if self.cleanup_query:
            self.log.info('Executing cleanup query: {}'.format(self.cleanup_query))
            self.hook.run(sql=self.cleanup_query)
            self.log.info('Completed execution of cleanup query')

        if isinstance(self.sql, str):
            self.sql = [self.sql]
        
        job_id = []
        if isinstance(self.sql, Iterable):
            for s in self.sql:
                self.log.info('Executing: %s', self.sql)
                query_job_id = self.hook.run_query(
                    sql=s,
                    destination_dataset_table=self.destination_dataset_table,
                    write_disposition=self.write_disposition,
                    allow_large_results=self.allow_large_results,
                    flatten_results=self.flatten_results,
                    udf_config=self.udf_config,
                    maximum_billing_tier=self.maximum_billing_tier,
                    maximum_bytes_billed=self.maximum_bytes_billed,
                    create_disposition=self.create_disposition,
                    query_params=self.query_params,
                    labels=self.labels,
                    schema_update_options=self.schema_update_options,
                    priority=self.priority,
                    time_partitioning=self.time_partitioning,
                    api_resource_configs=self.api_resource_configs,
                    cluster_fields=self.cluster_fields,
                    encryption_configuration=self.encryption_configuration,
                )
                job_id.append(query_job_id)
                qj = self.hook.get_job(job_id=query_job_id)
                while not qj.done():
                    self.log.info(f'Waiting for job [{query_job_id}] to complete...')
                    time.sleep(1)
        else:
            raise AirflowException(f"argument 'sql' of type {type(self.sql)} is neither a string nor an iterable")
        
        self.log.info('Completed BigQuery jobs: {}'.format(job_id))
        return job_id


class BigQueryJobTaskMetricOperator(EventEmittingOperator):
    """
    Neeeds help
    """

    @apply_defaults
    def __init__(
            self,
            task_ids: list=[],
            xcom_key: str='return_value',
            bigquery_conn_id: str='bigquery_default',
            *args, **kwargs):
        """
        :param task_ids: List of task_ids saying which tasks to sink their job_metrics for
        :param xcom_key: XCOM key used to pull the bigquery job id from the specified tasks
        """
        self.task_ids = task_ids
        self.xcom_key = xcom_key
        self.gcp_conn_id = bigquery_conn_id
        if 'gcp_conn_id' in kwargs:
            self.gcp_conn_id = kwargs['gcp_conn_id']
            del kwargs['gcp_conn_id']

        super(BigQueryJobTaskMetricOperator, self).__init__(*args, **kwargs)
        

    def execute(self, context):
        hook = BigQueryHook(
            gcp_conn_id=self.gcp_conn_id
        )
        credential_email = hook._get_credentials_email()
        self.log.info(f"Created BigQueryHook with account: {credential_email}")

        jobs = hook.get_service().jobs()

        for task_id in self.task_ids:
            job_id = context['task_instance'].xcom_pull(key=self.xcom_key, task_ids=task_id)

            if isinstance(job_id, str):
                job_id = [job_id]

            if job_id:
                for jid in job_id:
                    job = jobs.get(
                        projectId=hook.project_id,
                        jobId=jid
                    ).execute()
                    self.emit_event('task_metric', {
                        'run_id': context['dag_run'].run_id,
                        'dag_id': self.dag.dag_id,
                        'job_name': context['task'].params.get('job_name', ''),
                        'task_id': task_id,
                        'value': job,
                        'created_time': utc_now(),
                        'reference_time': context['next_execution_date']
                    })


class BigQueryShortCircuitOperator(BaseOperator, SkipMixin):
    """
    A "short circuit" operator that can be used a a "pre check" system.  The supplied sql statement should turn a single BOOL column.

    If the BOOL value is TRUE then downstream processors will execute as normal.
    If the BOOL value is FALSE then any downstream processors will be skipped.

    """

    template_fields = ('sql',)
    template_ext = ('.sql',)

    @apply_defaults
    def __init__(
            self,
            sql: str,
            bigquery_conn_id: str='bigquery_default',
            use_legacy_sql: bool=True,
            *args, **kwargs):
        
        self.sql = sql
        self.gcp_conn_id = bigquery_conn_id
        if 'gcp_conn_id' in kwargs:
            self.gcp_conn_id = kwargs['gcp_conn_id']
            del kwargs['gcp_conn_id']

        self.use_legacy_sql = use_legacy_sql
        super(BigQueryShortCircuitOperator, self).__init__(*args, **kwargs)


    def execute(self, context):
        records = self.db_hook.get_first(self.sql)
        success = records and all([bool(r) for r in records])

        if success:
            return

        self.log.info('Skipping downstream tasks...')

        downstream_tasks = context['task'].get_flat_relatives(upstream=False)
        self.log.debug("Downstream task_ids %s", downstream_tasks)

        if downstream_tasks:
            self.skip(context['dag_run'], context['ti'].execution_date, downstream_tasks)

        self.log.info("Done.")

    @property
    def db_hook(self):
        hook = BigQueryHook(gcp_conn_id=self.gcp_conn_id, use_legacy_sql=self.use_legacy_sql)
        credential_email = hook._get_credentials_email()
        self.log.info(f"Created BigQueryHook with account: {credential_email}")
        return hook
