import base64
import json
import logging
from google.cloud.bigtable import Client
from bt_df_data_retention_manager.config import Config
from bt_df_data_retention_manager.jobs.engines.dataflow import DataflowJob, DataflowJobRuntimeParams, running_jobs_bt_instance
from bt_df_data_retention_manager.jobs.job_status import JobStatusEnum
from bt_df_data_retention_manager.jobs.job_processor import JobProcessor
from bt_df_data_retention_manager.retention_dates import RetentionDateParams
from bt_df_data_retention_manager.state.record.retention import TableRetentionState
from bt_df_data_retention_manager.state.storage.bigtable import StateStorageManagerBigtable
from bt_df_data_retention_manager.state.storage.interface import NonexistentStateId
from bt_df_data_retention_manager.slack_client import SlackClient
from typing import List


def main(event, context):
    """Background Cloud Function to be triggered by Pub/Sub.
    Args:
         event (dict):  The dictionary with data specific to this type of
         event. The `data` field contains the PubsubMessage message. The
         `attributes` field will contain custom attributes if there are any.
         context (google.cloud.functions.Context): The Cloud Functions event
         metadata. The `event_id` field contains the Pub/Sub message ID. The
         `timestamp` field contains the publish time.
    """
    print("""This Function was triggered by messageId {} published at {}
        """.format(context.event_id, context.timestamp))

    if 'data' in event:
        data = json.loads(base64.b64decode(event['data']).decode('utf-8'))
        if 'tables' in data:
            delete_rows = data.get('delete_rows', False)
            if not isinstance(delete_rows, bool):
                delete_rows = False
            handler(data['tables'], delete_rows)
        else:
            raise Exception('expected list of tables not found in data')
    else:
        raise Exception('data is missing')


def handler(tables: List[str], delete_rows: bool):
    slack_client = SlackClient(slack_token=Config.SLACK_TOKEN)
    # initialize BT client
    client = Client(project=Config.PROJECT_ID, admin=False)
    instance = client.instance(Config.BT_INSTANCE_ID)
    state_table = instance.table(Config.BT_STATE_TABLE)

    # get dates
    date_params = RetentionDateParams(start_year=Config.START_YEAR,
                                      retention_months=Config.RETENTION_MONTHS,
                                      key_prefix=Config.KEY_PREFIX)
    date_params.generate()

    state_storage_manager = StateStorageManagerBigtable(
        state_table, Config.BT_STATE_TABLE_CF)

    # Run state machine with runnable=False
    # to update stored state with new statuses of the jobs and/or new job parameters if needed.
    # Count currently running jobs to verify how many jobs could be started to fit Config.SIMULTANEOUS_JOBS_LIMIT
    report = {}
    currently_running_jobs = running_jobs_bt_instance(
        project_id=Config.PROJECT_ID,
        instance_id=Config.BT_INSTANCE_ID,
        location=Config.REGION_ID)

    # fall back on env variable if no tables have been set up
    if len(tables) == 0:
        tables = Config.TABLES

    for table in tables:
        try:
            # Only launch runnable jobs if there are less than limit running
            if currently_running_jobs < Config.SIMULTANEOUS_JOBS_LIMIT:
                runnable = True
            else:
                runnable = False
            job_processor = get_job_processor(table,
                                              state_storage_manager,
                                              date_params,
                                              runnable=runnable,
                                              delete_rows=delete_rows)
            job_processor.to_next_state()
            if job_processor.is_in_progress():
                currently_running_jobs += 1

            state_storage_manager.put(job_processor.get_latest_state())
            report[table] = str(job_processor.get_latest_state().job_status)
        except Exception as e:
            error_msg = getattr(e, 'message',
                                str(e)).replace('\n',
                                                '\\n').replace('\r',
                                                               '\\r')[:989]
            notification = f"Error during cleaning up {table} " + error_msg
            slack_client.post_text_message(Config.SLACK_CHANNEL_ID,
                                           notification)
            logging.error(notification)
            raise

    # send report to slack:
    block = slack_client.get_section_block("report",
                                           ["*Table*", "*Job status*"], report)
    slack_client.post_block_message(Config.SLACK_CHANNEL_ID, [block])


def get_job_processor(table: str,
                      state_storage_manager: StateStorageManagerBigtable,
                      date_params: RetentionDateParams, delete_rows: bool,
                      runnable: bool) -> JobProcessor:
    """Generate job processor based on Config and given values"""
    try:
        stored_state = state_storage_manager.read(table)
    except NonexistentStateId:
        # initialize new state
        stored_state = TableRetentionState(
            table_id=table,
            job_status=JobStatusEnum.ready,
            latest_job_id='',
            current_end_date=date_params.end_date,
            retries=0)

    runtime_params = DataflowJobRuntimeParams(
        bigtable_instance_id=Config.BT_INSTANCE_ID,
        bigtable_table_id=table,
        max_num_workers=10,
        delete_rows=delete_rows,
        project_id=Config.PROJECT_ID,
        region=Config.REGION_ID,
        row_key_regex=date_params.filter_regex,
        end_retention_date=date_params.end_date,
        timeout_hours=Config.JOB_TIMEOUT_HOURS,
        staging_location=Config.DF_STAGING_LOCATION,
        template_location=Config.DF_TEMPLATE)

    job = DataflowJob(job_params=runtime_params,
                      job_id=stored_state.latest_job_id)

    return JobProcessor(job=job,
                        max_retries=Config.JOB_MAX_RETRIES,
                        runnable=runnable,
                        table_state=stored_state)