import gzip
import bz2
import lzma

from autogluon.common.utils import compression_utils


def test_get_validated_path_no_compression_fn():
    no_extension_filepath = 'dummy_file'
    assert(compression_utils.get_validated_path(no_extension_filepath) == no_extension_filepath)

    single_extension_filepath = 'dummy_file.pkl'
    assert(compression_utils.get_validated_path(single_extension_filepath) == single_extension_filepath)

    multiple_extension_filepath = 'dummy_file.fake-foo.zip.pkl'
    assert(compression_utils.get_validated_path(multiple_extension_filepath) == multiple_extension_filepath)


def test_get_validated_path_with_compression_fn():
    compression_fns = ['gzip', 'lzma', 'bz2']

    no_extension_filepath = 'dummy_file'

    expected_no_ext_gzip_filepath = "dummy_file.gz"
    expected_no_ext_lzma_filepath = "dummy_file.lzma"
    expected_no_ext_bz2_filepath = "dummy_file.bz2"
    assert(compression_utils.get_validated_path(no_extension_filepath, compression_fns[0])
           == expected_no_ext_gzip_filepath)
    assert(compression_utils.get_validated_path(no_extension_filepath, compression_fns[1])
           == expected_no_ext_lzma_filepath)
    assert(compression_utils.get_validated_path(no_extension_filepath, compression_fns[2])
           == expected_no_ext_bz2_filepath)

    single_extension_filepath = 'dummy_file.pkl'

    expected_pkl_gzip_filepath = "dummy_file.pkl.gz"
    expected_pkl_lzma_filepath = "dummy_file.pkl.lzma"
    expected_pkl_bz2_filepath = "dummy_file.pkl.bz2"
    assert(compression_utils.get_validated_path(single_extension_filepath, compression_fns[0])
           == expected_pkl_gzip_filepath)
    assert(compression_utils.get_validated_path(single_extension_filepath, compression_fns[1])
           == expected_pkl_lzma_filepath)
    assert(compression_utils.get_validated_path(single_extension_filepath, compression_fns[2])
           == expected_pkl_bz2_filepath)

    multiple_extension_filepath = 'dummy_file.fake-foo.zip.pkl'

    expected_multi_gzip_filepath = "dummy_file.fake-foo.zip.pkl.gz"
    expected_multi_lzma_filepath = "dummy_file.fake-foo.zip.pkl.lzma"
    expected_multi_bz2_filepath = "dummy_file.fake-foo.zip.pkl.bz2"
    assert (compression_utils.get_validated_path(multiple_extension_filepath, compression_fns[0])
            == expected_multi_gzip_filepath)
    assert (compression_utils.get_validated_path(multiple_extension_filepath, compression_fns[1])
            == expected_multi_lzma_filepath)
    assert (compression_utils.get_validated_path(multiple_extension_filepath, compression_fns[2])
            == expected_multi_bz2_filepath)


def test_get_compression_map():
    expected_compression_fn_map = {
        None: {
            'open': open,
            'extension': '',
        },
        'gzip': {
            'open': compression_utils._gzip_open,
            'extension': 'gz',
        },
        'bz2': {
            'open': compression_utils._bz2_open,
            'extension': 'bz2',
        },
        'lzma': {
            'open': compression_utils._lzma_open,
            'extension': 'lzma',
        },
    }
    assert compression_utils.get_compression_map() == expected_compression_fn_map
