import os
from tempfile import NamedTemporaryFile

import boto3
from moto import mock_s3
import pandas as pd
import pandavro as pdx
import pickle
import pytest


@pytest.fixture(autouse=True, scope='session')
def aws_credentials():
    """
    Sets AWS credentials to invalid values. Applied to all test functions and
    scoped to the entire testing session, so there's no chance of interfering
    with production buckets.
    """
    os.environ['AWS_ACCESS_KEY_ID'] = 'testing'
    os.environ['AWS_SECRET_ACCESS_KEY'] = 'testing'
    os.environ['AWS_SECURITY_TOKEN'] = 'testing'
    os.environ['AWS_SESSION_TOKEN'] = 'testing'


@pytest.fixture
def test_bucket():
    """Universal bucket name for use throughout testing"""
    return 'test_bucket'


@pytest.fixture
def test_keys():
    """List of keys to be used for populating a bucket with empty objects"""
    return sorted([
        'test_key_0.csv',
        'folder0/test_key_1.pq',
        'folder1/test_key_2.pkl',
        'folder1/subfolder0/test_key_3.pkl',
        'folder2/'
    ])


@pytest.fixture
def test_df_keys():
    """List of keys to be used for populating a bucket with DataFrames"""
    return {
        'avro': ['df.avro'],
        'csv': ['df.csv'],
        'csv.gz': ['df.csv.gz'],
        'csv.zip': ['df.csv.zip'],
        'csv.bz2': ['df.csv.bz2'],
        'csv.xz': ['df.csv.xz'],
        'psv': ['df.psv'],
        'psv.gz': ['df.psv.gz'],
        'psv.zip': ['df.psv.zip'],
        'psv.bz2': ['df.psv.bz2'],
        'psv.xz': ['df.psv.xz'],
        'feather': ['df.feather'],
        'json': ['df.json'],
        'pkl': ['df.pkl', 'df.pickle'],
        'pq': ['df.pq', 'df.parquet']
    }


@pytest.fixture
def test_df():
    """
    Universal dataframe for use throughout testing. Multiple data types
    used to test for proper encoding/decoding.
    """
    return pd.DataFrame({
        'intcol': [1, 2, 3],
        'strcol': ['four', 'five', 'six'],
        'floatcol': [7.0, 8.5, 9.0]
    })


@pytest.fixture
def mock_s3_client():
    """Mocks all s3 connections in any test or fixture that includes it"""
    with mock_s3():
        yield


@pytest.fixture
def setup_bucket_w_contents(mock_s3_client, test_bucket, test_keys):
    """
    Sets up a bucket with objects containing the empty string, based off
    keys in 'test_keys'
    """
    s3 = boto3.client('s3')
    s3.create_bucket(Bucket=test_bucket)

    for key in test_keys:
        s3.put_object(Bucket=test_bucket, Key=key, Body='')
    yield


@pytest.fixture
def setup_bucket_wo_contents(mock_s3_client, test_bucket):
    """Sets up a bucket with no contents."""
    s3 = boto3.client('s3')
    s3.create_bucket(Bucket=test_bucket)

    yield


@pytest.fixture
def setup_bucket_w_dfs(mock_s3_client, test_bucket, test_df, test_df_keys):
    """
    Sets up a bucket populated with dataframes that contain the data as
    defined in 'test_df', at the keys and storage formats defined in
    'test_df_keys'
    """
    s3 = boto3.client('s3')
    s3.create_bucket(Bucket=test_bucket)

    for key in test_df_keys['avro']:
        with NamedTemporaryFile() as tmpfile:
            pdx.to_avro(tmpfile, test_df)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['csv']:
        with NamedTemporaryFile() as tmpfile:
            test_df.to_csv(tmpfile.name, index=False)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['csv.gz']:
        with NamedTemporaryFile(suffix='.csv.gz') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['csv.zip']:
        with NamedTemporaryFile(suffix='.csv.zip') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['csv.bz2']:
        with NamedTemporaryFile(suffix='.csv.bz2') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['csv.xz']:
        with NamedTemporaryFile(suffix='.csv.xz') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['psv']:
        with NamedTemporaryFile() as tmpfile:
            test_df.to_csv(tmpfile.name, index=False, sep='|')
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['psv.gz']:
        with NamedTemporaryFile(suffix='.psv.gz') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False, sep='|')
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['psv.zip']:
        with NamedTemporaryFile(suffix='.psv.zip') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False, sep='|')
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['psv.bz2']:
        with NamedTemporaryFile(suffix='.psv.bz2') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False, sep='|')
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['psv.xz']:
        with NamedTemporaryFile(suffix='.psv.xz') as tmpfile:
            test_df.to_csv(tmpfile.name, index=False, sep='|')
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['feather']:
        with NamedTemporaryFile() as tmpfile:
            test_df.to_feather(tmpfile.name)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['json']:
        with NamedTemporaryFile() as tmpfile:
            test_df.to_json(tmpfile.name)
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['pkl']:
        with NamedTemporaryFile() as tmpfile:
            pickle.dump(test_df, tmpfile, protocol=pickle.HIGHEST_PROTOCOL)
            tmpfile.flush()
            s3.upload_file(tmpfile.name, test_bucket, key)

    for key in test_df_keys['pq']:
        with NamedTemporaryFile() as tmpfile:
            test_df.to_parquet(tmpfile.name, index=False)
            s3.upload_file(tmpfile.name, test_bucket, key)

    yield

[((12, 1, 12, 46), 'pytest.fixture', 'pytest.fixture', (), '', False, 'import pytest\n'), ((71, 11, 75, 6), 'pandas.DataFrame', 'pd.DataFrame', ({(71, 24, 75, 5): "{'intcol': [1, 2, 3], 'strcol': ['four', 'five', 'six'], 'floatcol': [7.0, \n    8.5, 9.0]}"}, {}), "({'intcol': [1, 2, 3], 'strcol': ['four', 'five', 'six'],\n    'floatcol': [7.0, 8.5, 9.0]})", True, 'import pandas as pd\n'), ((91, 9, 91, 27), 'boto3.client', 'boto3.client', ({(91, 22, 91, 26): '"""s3"""'}, {}), "('s3')", False, 'import boto3\n'), ((102, 9, 102, 27), 'boto3.client', 'boto3.client', ({(102, 22, 102, 26): '"""s3"""'}, {}), "('s3')", False, 'import boto3\n'), ((115, 9, 115, 27), 'boto3.client', 'boto3.client', ({(115, 22, 115, 26): '"""s3"""'}, {}), "('s3')", False, 'import boto3\n'), ((81, 9, 81, 18), 'moto.mock_s3', 'mock_s3', ({}, {}), '()', False, 'from moto import mock_s3\n'), ((119, 13, 119, 33), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ({}, {}), '()', False, 'from tempfile import NamedTemporaryFile\n'), ((120, 12, 120, 41), 'pandavro.to_avro', 'pdx.to_avro', ({(120, 24, 120, 31): 'tmpfile', (120, 33, 120, 40): 'test_df'}, {}), '(tmpfile, test_df)', True, 'import pandavro as pdx\n'), ((124, 13, 124, 33), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ({}, {}), '()', False, 'from tempfile import NamedTemporaryFile\n'), ((129, 13, 129, 49), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((134, 13, 134, 50), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((139, 13, 139, 50), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((144, 13, 144, 49), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((149, 13, 149, 33), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ({}, {}), '()', False, 'from tempfile import NamedTemporaryFile\n'), ((154, 13, 154, 49), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((159, 13, 159, 50), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((164, 13, 164, 50), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((169, 13, 169, 49), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', (), '', False, 'from tempfile import NamedTemporaryFile\n'), ((174, 13, 174, 33), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ({}, {}), '()', False, 'from tempfile import NamedTemporaryFile\n'), ((179, 13, 179, 33), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ({}, {}), '()', False, 'from tempfile import NamedTemporaryFile\n'), ((184, 13, 184, 33), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ({}, {}), '()', False, 'from tempfile import NamedTemporaryFile\n'), ((185, 12, 185, 75), 'pickle.dump', 'pickle.dump', (), '', False, 'import pickle\n'), ((190, 13, 190, 33), 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ({}, {}), '()', False, 'from tempfile import NamedTemporaryFile\n')]