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, 'pytest.fixture', 'pytest.fixture', ([], {'autouse': '(True)', 'scope': '"""session"""'}), False, 'import pytest\n'), (71, 'pandas.DataFrame', 'pd.DataFrame', (["{'intcol': [1, 2, 3], 'strcol': ['four', 'five', 'six'], 'floatcol': [7.0, \n    8.5, 9.0]}"], {}), True, 'import pandas as pd\n'), (91, 'boto3.client', 'boto3.client', (['"""s3"""'], {}), False, 'import boto3\n'), (102, 'boto3.client', 'boto3.client', (['"""s3"""'], {}), False, 'import boto3\n'), (115, 'boto3.client', 'boto3.client', (['"""s3"""'], {}), False, 'import boto3\n'), (81, 'moto.mock_s3', 'mock_s3', ([], {}), False, 'from moto import mock_s3\n'), (119, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {}), False, 'from tempfile import NamedTemporaryFile\n'), (120, 'pandavro.to_avro', 'pdx.to_avro', (['tmpfile', 'test_df'], {}), True, 'import pandavro as pdx\n'), (124, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {}), False, 'from tempfile import NamedTemporaryFile\n'), (129, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".csv.gz"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (134, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".csv.zip"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (139, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".csv.bz2"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (144, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".csv.xz"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (149, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {}), False, 'from tempfile import NamedTemporaryFile\n'), (154, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".psv.gz"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (159, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".psv.zip"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (164, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".psv.bz2"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (169, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {'suffix': '""".psv.xz"""'}), False, 'from tempfile import NamedTemporaryFile\n'), (174, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {}), False, 'from tempfile import NamedTemporaryFile\n'), (179, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {}), False, 'from tempfile import NamedTemporaryFile\n'), (184, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {}), False, 'from tempfile import NamedTemporaryFile\n'), (185, 'pickle.dump', 'pickle.dump', (['test_df', 'tmpfile'], {'protocol': 'pickle.HIGHEST_PROTOCOL'}), False, 'import pickle\n'), (190, 'tempfile.NamedTemporaryFile', 'NamedTemporaryFile', ([], {}), False, 'from tempfile import NamedTemporaryFile\n')]