"""Snowflake helper methods."""
from connectors.snowflake import snow_session_wrap

from constants import db
from constants.fields import (
    DO_NOT_EXECUTE_STATUSES,
    FINISHED_STATUSES,
    IN_PROGRESS_STATUSES,
    ORPHANED_STATUSES
)

from sql.copy_table_to_json import COPY_TABLE_TO_JSON, COPY_TABLE_TO_JSON_WITH_TOKEN
from sql.create_asset_sub_table_queue_table import CREATE_ASSET_SUB_TABLE_QUEUE_TABLE  # noqa
from sql.create_wakeup_sub_table_queue_table import CREATE_WAKEUP_SUB_TABLE_QUEUE_TABLE  # noqa
from sql.create_bulk_asset_sfn_queue_table import CREATE_BULK_ASSET_SFN_QUEUE_TABLE  # noqa
from sql.create_bulk_asset_table_queue_table import CREATE_BULK_ASSET_TABLE_QUEUE_TABLE  # noqa
from sql.delete_table_queue_item import DELETE_TABLE_QUEUE_ITEM
from sql.get_table_row_count import GET_TABLE_ROW_COUNT
from sql.get_release_sizes import GET_RELEASE_SIZES
from sql.insert_sfn_queue_item import INSERT_SFN_QUEUE_ITEM
from sql.insert_table_queue_item import INSERT_TABLE_QUEUE_ITEM
from sql.select_all_tables_to_ingest import SELECT_ALL_TABLES_TO_INGEST
from sql.select_in_progress_ingestions import SELECT_IN_PROGRESS_INGESTIONS
from sql.select_queued_ingestions import SELECT_QUEUED_INGESTIONS


from sql.select_table_by_table_name import SELECT_TABLE_BY_TABLE_NAME
from sql.select_sfn_by_table_name import SELECT_SFN_BY_TABLE_NAME
from sql.select_sfn_by_table_name_fuzzy import SELECT_SFN_BY_TABLE_NAME_FUZZY
from sql.select_table_to_ingest import SELECT_TABLE_TO_INGEST
from sql.select_todo_sfn_queue_items import SELECT_TODO_SFN_QUEUE_ITEMS
from sql.select_orphaned_sfn_queue_items import SELECT_ORPHANED_SFN_QUEUE_ITEMS
from sql.update_queue_field_by_exec_name import UPDATE_QUEUE_FIELD_BY_EXEC_NAME
from sql.update_queue_field_by_table_name import UPDATE_QUEUE_FIELD_BY_TABLE_NAME  # noqa

import config

DEFAULT_PARAMS = {
    'catalog': config.SNOWFLAKE_CATALOG,
    'schema': config.SNOWFLAKE_SCHEMA,
    'env': config.ENVIRONMENT
}


@snow_session_wrap
def create_sfn_queue(cursor):
    """Create sfn queue table."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(CREATE_BULK_ASSET_SFN_QUEUE_TABLE.format(**params))


@snow_session_wrap
def create_table_queue(cursor):
    """Create table queue... table."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(CREATE_BULK_ASSET_TABLE_QUEUE_TABLE.format(**params))


@snow_session_wrap
def create_asset_sub_table(
        cursor, upc_list, source_table, sub_table, wakeup=False):
    """Create bulk asset queue sub table."""
    params = {
        **DEFAULT_PARAMS,
        'source_table': source_table,
        'sub_table': sub_table
    }

    if wakeup:
        sql = CREATE_WAKEUP_SUB_TABLE_QUEUE_TABLE
    else:
        sql = CREATE_ASSET_SUB_TABLE_QUEUE_TABLE

    cursor.execute(sql.format(**params), {'upc_list': upc_list})


@snow_session_wrap
def add_table_queue_item(cursor, table_name):
    """Add item to the table queue."""
    params = {
        **DEFAULT_PARAMS,
        'source_table': table_name
    }

    cursor.execute(INSERT_TABLE_QUEUE_ITEM.format(**params))


@snow_session_wrap
def add_sfn_queue_item(cursor,
                       sfn_name,
                       source_table,
                       execution_name=None,
                       trigger_json_key=None,
                       status=None):
    """Add item to the queue."""
    params = {
        **DEFAULT_PARAMS,
        'sfn_name': f"'{sfn_name}'",
        'source_table': f"'{source_table}'",
        'execution_name': execution_name or 'NULL',
        'trigger_json_key': trigger_json_key or 'NULL',
        'status': f"'{status}'" if status else "'RUNNING'"
    }

    cursor.execute(INSERT_SFN_QUEUE_ITEM.format(**params))


@snow_session_wrap
def get_queued_ingestions(cursor):
    """Get all queued ingestions."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(SELECT_QUEUED_INGESTIONS.format(**params))

    return cursor.fetchall()


@snow_session_wrap
def get_todo_ingestions(cursor):
    """Get todo queued ingestions."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(
        SELECT_TODO_SFN_QUEUE_ITEMS.format(**params),
        {'bad_status_list': tuple(DO_NOT_EXECUTE_STATUSES)}
        )

    return cursor.fetchall()


@snow_session_wrap
def get_orphaned_ingestions(cursor):
    """Get orphaned queued ingestions."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(
        SELECT_ORPHANED_SFN_QUEUE_ITEMS.format(**params),
        {'orphaned_status_list': tuple(ORPHANED_STATUSES)})

    return cursor.fetchall()


@snow_session_wrap
def get_release_sizes(cursor, table_name):
    """Get release sizes."""
    params = {
        **DEFAULT_PARAMS,
        'source_table': table_name
    }

    cursor.execute(GET_RELEASE_SIZES.format(**params))

    return cursor.fetchall()


@snow_session_wrap
def delete_table_queue_item(cursor, table_name):
    """Delete a table queue item by name."""
    params = {
        **DEFAULT_PARAMS,
        'source_table': table_name
    }

    cursor.execute(DELETE_TABLE_QUEUE_ITEM.format(**params))


@snow_session_wrap
def get_in_progress_ingestions(cursor):
    """Get in-progress ingestions."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(
        SELECT_IN_PROGRESS_INGESTIONS.format(**params),
        {'finished_status_list': tuple(FINISHED_STATUSES),
         'in_progress_status_list': tuple(IN_PROGRESS_STATUSES)}
    )

    return cursor.fetchall()


@snow_session_wrap
def get_table_to_ingest(cursor):
    """Get a single table from the table queue."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(SELECT_TABLE_TO_INGEST.format(**params))

    return cursor.fetchone()[db.SOURCE_TABLE]


@snow_session_wrap
def get_all_tables_to_ingest(cursor):
    """Get all tables in the table queue."""
    params = {
        **DEFAULT_PARAMS
    }

    cursor.execute(SELECT_ALL_TABLES_TO_INGEST.format(**params))

    return cursor.fetchall()


@snow_session_wrap
def get_rows_in_sfn_queue_by_table(cursor, table_name, status=None):
    """Get all rows in the SFN queue by table name.

    Args:
        cursor: Snowflake cursor
        table_name: str, the name of the table by which to filter rows
        status: str or list, the status by which to filter rows

    Returns:
        list, all rows in the SFN queue that match the table name and status
        (if provided)
    """
    params = {
        **DEFAULT_PARAMS,
        'source_table': table_name
    }

    cursor.execute(SELECT_SFN_BY_TABLE_NAME.format(**params))

    result = cursor.fetchall()

    # Filter by status
    if status and isinstance(status, str):  # Single status
        result = [row for row in result if row[db.STATUS] == status]
    elif status and isinstance(status, list):  # Multiple statuses
        result = [row for row in result if row[db.STATUS] in status]

    return result


@snow_session_wrap
def get_rows_in_sfn_queue_by_fuzzy_table(cursor, table_name_prefix, status=None):  # noqa
    """Get all rows in the SFN queue by using a table name prefix.

    Args:
        cursor: Snowflake cursor
        table_name_prefix: str, the prefix of the table by which to filter rows
        status: str or list, the status by which to filter rows

    Returns:
        list, all rows in the SFN queue that match the table name and status
        (if provided)
    """
    params = {
        **DEFAULT_PARAMS,
        'source_table_prefix': table_name_prefix
    }

    cursor.execute(SELECT_SFN_BY_TABLE_NAME_FUZZY.format(**params))

    result = cursor.fetchall()

    # Filter by status
    if status and isinstance(status, str):  # Single status
        result = [row for row in result if row[db.STATUS] == status]
    elif status and isinstance(status, list):  # Multiple statuses
        result = [row for row in result if row[db.STATUS] in status]

    return result


@snow_session_wrap
def is_table_in_table_queue(cursor, table_name):
    """Check if a table is in the table queue.

    Args:
        cursor: Snowflake cursor
        table_name: str, the name of the table to check

    Returns:
        bool, True if the table is in the queue, False otherwise
    """
    params = {
        **DEFAULT_PARAMS,
        'source_table': table_name
    }

    cursor.execute(SELECT_TABLE_BY_TABLE_NAME.format(**params))

    return cursor.fetchone() is not None


@snow_session_wrap
def update_queue_field_by_exec(cursor, field_name, value, execution_name):
    """Update a field in the queue using the execution name."""
    params = {
        **DEFAULT_PARAMS,
        'field_name': field_name,
        'value': value,
        'execution_name': execution_name
    }

    cursor.execute(UPDATE_QUEUE_FIELD_BY_EXEC_NAME.format(**params))


@snow_session_wrap
def update_queue_field_by_table(cursor, field_name, value, table_name):
    """Update a field in the queue using the table name."""
    params = {
        **DEFAULT_PARAMS,
        'field_name': field_name,
        'value': value,
        'source_table': table_name
    }

    cursor.execute(UPDATE_QUEUE_FIELD_BY_TABLE_NAME.format(**params))


@snow_session_wrap
def drop_json_trigger(cursor, bucket, prefix, file_name, table_name):
    """Drop a JSON file to trigger a SFN."""
    params = {
        **DEFAULT_PARAMS,
        'bucket': bucket,
        'prefix': prefix,
        'key': file_name,
        'table_name': table_name,
        'aws_key_id': config.AWS_ACCESS_KEY_ID,
        'aws_secret_key': config.AWS_SECRET_ACCESS_KEY
    }

    if config.AWS_SESSION_TOKEN:
        params['aws_token'] = config.AWS_SESSION_TOKEN
        cursor.execute(COPY_TABLE_TO_JSON_WITH_TOKEN.format(**params))
    else:
        cursor.execute(COPY_TABLE_TO_JSON.format(**params))


@snow_session_wrap
def get_table_row_count(cursor, table_name):
    """Get the number of rows in a table."""
    params = {
        **DEFAULT_PARAMS,
        'table_name': table_name
    }

    cursor.execute(GET_TABLE_ROW_COUNT.format(**params))

    return cursor.fetchone()['TOTAL_ROWS']