"""Helper functions and classes for Python based integration tests."""
import os
from pathlib import Path

import boto3
from garcon_contrib.aws.utils import garcon_s3

import feed_ingestion.conf.config as global_config
from feed_ingestion.flows.helpers import upload_raw_file_to_s3
from feed_ingestion.util.aws import s3 as s3utils
from integration_tests.python_integration_tests.snowflake_executor import \
    IntegrationTestSFExecutor

from feed_ingestion.flows.base import get_aws_config


FIXTURE_PATH = Path(__file__).parents[1]


def get_list_of_files(path, s3_full_path):
    """Put sample files on S3.

    Args:
        path (str): A local path with files.
        s3_full_path (str): S3 path for coping files.

    Yields:
        (tuple): Full local path and full s3 path.
    """
    for root, _, files in os.walk(path):
        for filename in files:
            yield (
                os.path.join(path, filename),
                os.path.join(s3_full_path, filename))


def put_sample_files_on_s3(feed_name, s3_path):
    """Put sample files on S3.

    Args:
        feed_name (str): A feed name
        s3_path (str): S3 path for coping files.
    """
    path = FIXTURE_PATH / 'data' / feed_name

    for local_path, s3_full_path in get_list_of_files(path, s3_path):
            upload_raw_file_to_s3(local_path, s3_full_path)


def seed_snowflake_tables(feed_name, s3_path):
    """Seed tables in snowflake with csv files.

    Copy csv files with data on s3 first, truncate the table and seed it.
    The table name should be the same as filename.

    Args:
        feed_name (str): A feed name
        s3_path (str): S3 path for coping files.
    """
    with IntegrationTestSFExecutor(global_config.SF_CONFIG) as executor:
        path = FIXTURE_PATH / 'snowflake_seed' / feed_name

        for local_path, s3_full_path in get_list_of_files(path, s3_path):
            upload_raw_file_to_s3(local_path, s3_full_path)
            table_name = Path(local_path).stem
            executor.truncate_table(table_name)
            executor.load_table(
                table_name, get_aws_config(), s3_full_path)


def truncate_tables(tables_list):
    """Truncate destination tables in the integration_tests db schema.

    Args:
        tables_list (list): The list of table names.
    """
    with IntegrationTestSFExecutor(global_config.SF_CONFIG) as executor:
        for table_name in tables_list:
            executor.truncate_table(table_name)


def get_number_of_files(s3_directory):
    """Return number of files in s3_directory.

    Args:
        s3_directory (str): S3 directory to check.

    Returns:
        int: Number of keys for the specified S3 path prefix.
    """
    return len(s3utils.get_list_of_files_and_directories(s3_directory))


def get_number_of_rows_in_table(table_name, **kwargs):
    """Return number of rows in a table in snowflake.

    Args:
        table_name (str): A table name.
        kwargs (dict): Dictionary with additional filtering params.
    """
    with IntegrationTestSFExecutor(global_config.SF_CONFIG) as executor:
        return executor.get_number_of_rows(table_name, **kwargs)[0][0]


def remove_test_files_from_s3(s3_directory,
                              expected_bucket_owner='437795906767'):
    """Delete sample files and directories from S3.

    Args:
        s3_directory (str): S3 directory to check.
        expected_bucket_owner (str): Expected bucket owner.
    """
    s3_client = boto3.client('s3')

    bucket, bucket_path = garcon_s3.extract_bucket_path(s3_directory)
    existing_keys = s3_client.list_objects(
        Bucket=bucket,
        Prefix=bucket_path,
        ExpectedBucketOwner=expected_bucket_owner)['Contents']
    for key in existing_keys:
        s3_client.delete_object(
            Bucket=bucket,
            Key=key['Key'],
            ExpectedBucketOwner=expected_bucket_owner
        )