"""Application index."""

import os
import re
import sys
import time

import backoff
from lambdacommon import util
from neo4j import GraphDatabase
from neo4j.exceptions import Neo4jError
from owslogger import logger
import sentry_sdk

import config


# Global logger
log = logger.setup(
    config.ENVIRONMENT, config.LOGGER_NAME,
    config.LOGGER_LEVEL, config.SERVICE_NAME, '1.1.1', dsn=config.LOGGER_DSN)

# Sentry handler
if not config.DISABLE_SENTRY:
    sentry_sdk.init(config.SENTRY_DSN)


def _send_datadog_metric(metric):
    """Send metric to datadog.

    Args:
        metric (dict): metric object to send.
    """
    with util.datadog_connection(
        api_key=config.DATADOG_API_KEY,
        app_key=config.DATADOG_APP_KEY
    ) as datadog:
        datadog.Metric.send([metric])


def call_datadog_with_metric(
    metric, metric_type, query_file=None, value=1, tags=[]
):
    """
    Wrap the DataDog metric call.

    Args:
        metric (str): custom metric name to send to DataDog.
        metric_type (str): info or count metric type.
        query_file (str): Query filename.
        value (int): metric value.
        tags (list): list of additional string tags to be added.
    """
    metric_name = f'{config.SERVICE_NAME}.{metric}.{metric_type}'
    log.info(f'Sending {metric_name} metric to datadog.')
    now = time.time()
    metric_obj = {
        'metric': metric_name,
        'type': 'count',
        'interval': 60,
        'points': (now, value),
        'tags': [
            f'environment:{config.ENVIRONMENT}',
            f'cypher_source:{config.CYPHER_SOURCE}'
        ]
    }
    if query_file:
        metric_obj['tags'].append(f'query_file:{query_file}')
    if tags:
        metric_obj['tags'].extend(tags)

    _send_datadog_metric(metric_obj)


def capture_exception(exception):
    """Capture exception in Sentry if Sentry is configured."""
    if config.SENTRY_DSN and not config.DISABLE_SENTRY:
        sentry_sdk.capture_exception(exception)
    else:
        log.info('Sentry is not configured for this job.')


def format_query(source_name):
    """
    Read Neo4j query from a file.

    Args:
        source_name (str): Name of cypher source. Should match the name
        of the file, minus cypher extension, in the queries/ directory
        e.g. mysql-pricing, snowflake-chartmetric

    Returns:
        str: Query to run
    """
    query_file = open(
        'queries/{}.cypher'.format(source_name), 'r', encoding='utf-8')
    cypher = query_file.read()
    query_file.close()

    return cypher


def get_queries_from_folder(source_folder_name):
    """
    Return the list of files containing the queries to run ordered by filename.

    Args:
        source_folder_name (str): The name of the folder containing the query
            files.

    Returns:
        [str]: The list of queries to run ordered by filename.
    """
    source_folder_path = 'queries/{}'.format(source_folder_name)
    queries = []

    for item in os.listdir(source_folder_path):
        is_file = os.path.isfile(os.path.join(source_folder_path, item))
        is_cypher = os.path.splitext(item)[1] == '.cypher'

        if is_file and is_cypher:
            queries.append('{}/{}'.format(
                source_folder_name, item.replace('.cypher', '')))

    queries.sort()
    return queries


def get_query_list_from_cypher(cypher):
    """
    Return the list of valid queries contained in the cypher file.

    Args:
        cypher (str): The content of the cypher file.

    Returns:
        List[str]: The list of valid queries contained in the cypher file.
    """
    return [query.strip() for query in cypher.split(';') if query.strip()]


def execute_query(tx, query, params):
    """
    Execute a query in a transaction.

    Args:
        tx (neo4j.Transaction): The neo4j transaction to execute the query in.
        query (str): The formatted query ready to be executed.
        params (dict): The parameters to use to execute the query.

    Returns:
        dict: The result of the query.
    """
    return tx.run(query, params).data()


def backoff_handler(details):
    """Log a warning everytime the backoff mechanism is triggered."""
    log.warning('Backing off {wait:0.1f} seconds after {tries} tries '
                'calling function {target} with args {args} and kwargs '
                '{kwargs}'.format(**details))


def clean_string(input_str: str):
    """Remove any special characters from string."""
    output = re.sub(r'[\[(){}<>:\]]', '', input_str.lower())
    output = re.sub('-', '_', output).strip('_')
    return output


# If a Neo4j error occurs, retry once after a 60s delay
@backoff.on_exception(
    backoff.constant,
    Neo4jError,
    max_tries=3,
    on_backoff=backoff_handler,
    jitter=None,
    interval=60
)
def run_single_query(
        source_name, driver, query=None, query_index=1, params={}):
    """
    Run a single query.

    Args:
        source_name (str): The name of the file containing the query to run.
        driver (neo4j.Neo4jDriver): The neo4j driver.
        query (str): The formatted query ready to be executed.
        query_index (int): The index of the query in the case where
            a file contains multiple queries.
        params (dict): The parameters to use to run the query.

    Returns:
        [dict]: The result of the query.
    """
    log.info('Running single query #{} from {}'.format(
        query_index, source_name))

    if not query:
        query = format_query(source_name)

    try:
        with driver.session(database=config.NEO4J_DATABASE_NAME) as session:
            response = session.run(query, params)
            result = response.data()
    except Exception as error:
        log.warning(f'An {type(error).__name__} exception occurred: {error}')
        raise error

    log.info('Running single query #{} from {} was successful'.format(
        query_index, source_name))

    # Filename without extension
    filename = source_name.split('/')[-1].split('.')[0]
    if filename in config.SEND_QUERY_RESULTS_TO_DATADOG:
        # cypher needs to have a single dictionary as result with key:value,
        # so we can expose metric as: <folder>.query-count with tag:entity:<key>
        first_result = next(iter(result[0].values()))
        if isinstance(first_result, dict):
            for key, value in first_result.items():
                clean_tag = clean_string(key)
                call_datadog_with_metric(
                    config.CYPHER_SOURCE,
                    'query-count',
                    query_file=filename,
                    value=value,
                    tags=[f'entity:{clean_tag}']
                )
    return result


def run_query_in_batches(source_name, count_query, query, driver):
    """
    Run a query in batches.

    Args:
        source_name (str): The name of the file containing the query to run.
        count_query (str): The query to get the number of batches to run.
        query (str): The query to run in batches.
        driver (neo4j.Neo4jDriver): The neo4j driver.

    Returns:
        Boolean: True to notify that the batches have been run successfully.
    """
    log.info('Running query from {} in batches'.format(source_name))

    count_result = run_single_query(
        source_name, driver, query=count_query, query_index=1)

    count = count_result[0]['count']
    skip = 0

    while count > 0:
        log.info('Running batch, {} items remaining'.format(count))

        params = {'skip': skip, 'limit': config.BATCH_SIZE}
        run_single_query(
            source_name, driver, query=query, query_index=2, params=params)

        skip += config.BATCH_SIZE
        count -= config.BATCH_SIZE

    return True


def run_multiple_queries(source_folder_name, driver):
    """
    Run multiple queries.

    Args:
        source_folder_name (str): The name of the folder containing the query
            files.
        driver (neo4j.Neo4jDriver): The neo4j driver.

    Returns:
        [dict]: An array containing the result of each query.
    """
    log.info('Running multiple queries from {}'.format(source_folder_name))

    query_files = get_queries_from_folder(source_folder_name)
    results = []

    for query_file in query_files:
        query_file_name = query_file.split('/')[-1]
        cypher = format_query(query_file)

        # There might be multiple queries in a single file
        query_list = get_query_list_from_cypher(cypher)

        try:
            # Some queries need to be run in multiple batches,
            # the file should contain two cypher queries,
            # the first one to get the total count of nodes to process,
            # the second one to do the actual processing
            if config.RUN_IN_BATCHES and len(query_list) == 2:
                count_query = query_list[0]
                query = query_list[1]
                result = run_query_in_batches(
                    query_file, count_query, query, driver)
                results.append(
                    {
                        'query_name': query_file_name,
                        'query': query,
                        'result': result
                    }
                )
            else:
                for index, query in enumerate(query_list):
                    query_index = index + 1
                    result = run_single_query(
                        query_file,
                        driver,
                        query=query,
                        query_index=query_index)
                    results.append(
                        {
                            'query_name': query_file_name,
                            'query': query,
                            'result': result
                        }
                    )
        except Exception as error:
            # For some flows we want to keep running the remaining queries,
            # even if a previous one failed
            if config.CONTINUE_AFTER_ERRORS:
                log.error('Error running query from {}: {}'.format(
                    query_file, error))
                continue
            else:
                raise error

    log.info('Running multiple queries from {} was successful'.format(
        source_folder_name))

    return results


def run_snowflake_sync_checks(driver):
    """
    Run row-level Neo4j <-> Snowflake sync checks for all configured checks.

    Args:
        driver (neo4j.Neo4jDriver): The neo4j driver.

    Returns:
        [dict]: A list of result dicts, one per check, each containing
            'query_name', 'query', and 'result' keys.
    """
    from src.assert_snowflake_sync import snowflake_sync
    from src.assert_snowflake_sync import checks as snowflake_checks

    log.info('Running Snowflake sync checks from assert-snowflake-sync')
    window_start, window_end = snowflake_sync.build_window()
    results = []

    with config.create_snowflake_executor() as executor:
        for check in snowflake_checks.CHECKS:
            check_name = check['name']

            summary, comparison = snowflake_sync.run_row_level_check(
                check, driver, executor, window_start, window_end)

            results.append({
                'query_name': check_name,
                'query': 'row_level_compare',
                'result': [summary],
            })

            if summary.get('shouldBeTrue'):
                log.info(f'Row-level check {check_name} succeeded.')
                continue

            log.warning(f'Row-level check {check_name} mismatched.')
            for label, items in comparison.items():
                if items:
                    log.warning(f'{label} for {check_name}: {items}')

    log.info('Snowflake sync checks completed.')
    return results


def assert_results(results):
    """
    Assert that the result of each query is the expected one.

    Args:
        results ([dict]): The array containing the result of each query.
    """
    for item in results:
        query_name = item.get('query_name', config.CYPHER_SOURCE)
        query = item['query']
        result = item['result'][0]

        # send success and fail metrics for each query file.
        if 'shouldBeTrue' in result and not result['shouldBeTrue']:
            error = f'Assertion for {query_name} FAILED!\nQuery: {query}'
            call_datadog_with_metric(
                config.CYPHER_SOURCE,
                config.DATADOG_CYPHER_FAILURE_METRIC,
                query_name
            )
            log.error(error)
        elif result.get('shouldBeTrue', False):
            call_datadog_with_metric(
                config.CYPHER_SOURCE,
                config.DATADOG_CYPHER_SUCCESS_METRIC,
                query_name
            )
            log.info(f'Assertion for {query_name} was successful!!')

    log.info('Results asserted')


def main():
    """Entrypoint function."""
    log.info('Creating a connection to {}'.format(config.NEO4J_URL))

    driver = GraphDatabase.driver(config.NEO4J_URL, auth=(
        config.NEO4J_CONNECTION_USER, config.NEO4J_CONNECTION_PASSWORD))

    try:
        if config.ASSERT_SNOWFLAKE_SYNC:
            log.info('Running snowflake sync checks')
            results = run_snowflake_sync_checks(driver)
        elif config.NEO4J_SOURCE_FOLDER_NAME:
            results = run_multiple_queries(
                config.NEO4J_SOURCE_FOLDER_NAME, driver)
        else:
            result = run_single_query(config.NEO4J_SOURCE_NAME, driver)
            results = [{
                'query': config.NEO4J_SOURCE_NAME,
                'result': result
            }]
    except Exception as error:
        capture_exception(error)
        driver.close()
        sys.exit('Failed to run cypher from {} on {}: {}'.format(
            (config.NEO4J_SOURCE_FOLDER_NAME or config.NEO4J_SOURCE_NAME),
            config.NEO4J_URL,
            error))

    driver.close()

    log.info('Running cyphers from {} on {} was successful'.format(
        (config.NEO4J_SOURCE_FOLDER_NAME or config.NEO4J_SOURCE_NAME),
        config.NEO4J_URL))

    if not config.ASSERT_SNOWFLAKE_SYNC:
        if config.ASSERT_RESULTS:
            assert_results(results)
        else:
            # send a single success if we are not going to ASSERT individual RESULTS
            call_datadog_with_metric(
                config.CYPHER_SOURCE,
                config.DATADOG_CYPHER_SUCCESS_METRIC
            )

    return results


if __name__ == '__main__':
    main()