import io
import pytest

from alchemtest.gmx import load_expanded_ensemble_case_1
from alchemlyb.parsing.util import anyopen


def test_gzip():
    """Test that gzip reads .gz files in the correct (text) mode.

    """
    dataset = load_expanded_ensemble_case_1()

    for leg in dataset['data']:
        for filename in dataset['data'][leg]:
            with anyopen(filename, 'r') as f:
                assert type(f.readline()) is str


def test_gzip_stream():
    """Test that `anyopen` reads streams with specified compression.

    """
    dataset = load_expanded_ensemble_case_1()

    for leg in dataset['data']:
        for filename in dataset['data'][leg]:
            with open(filename, 'rb') as f:
                with anyopen(f, mode='r', compression='gzip') as f_uc:
                    assert type(f_uc.readline()) is str


def test_gzip_stream_wrong():
    """Test that `anyopen` gives failure for attempting to decompress gzip
    stream with bz2.

    """
    dataset = load_expanded_ensemble_case_1()

    for leg in dataset['data']:
        for filename in dataset['data'][leg]:
            with open(filename, 'rb') as f:
                with anyopen(f, mode='r', compression='bzip2') as f_uc:
                    with pytest.raises(OSError, match='Invalid data stream'):
                        assert type(f_uc.readline()) is str


def test_gzip_stream_wrong_no_compression():
    """Test that `anyopen` gives passthrough when no compression specified on a
    stream.

    """
    dataset = load_expanded_ensemble_case_1()

    for leg in dataset['data']:
        for filename in dataset['data'][leg]:
            with open(filename, 'rb') as f:
                with anyopen(f, mode='r') as f_uc:
                    assert type(f_uc.readline()) is bytes


@pytest.mark.parametrize('extension', ['bz2', 'gz'])
def test_file_roundtrip(extension, tmp_path):
    """Test that roundtripping write/read to a file works with `anyopen`.

    """

    data = "my momma told me to pick the very best one and you are not it"

    filepath = tmp_path / f'testfile.txt.{extension}'
    with anyopen(filepath, mode='w') as f:
        f.write(data)

    with anyopen(filepath, 'r') as f:
        data_out = f.read()

    assert data_out == data


@pytest.mark.parametrize('extension,compression',
        [('bz2', 'gzip'), ('gz', 'bzip2')])
def test_file_roundtrip_force_compression(extension, compression, tmp_path):
    """Test that roundtripping write/read to a file works with `anyopen`,
    in which we force compression despite different extension.

    """

    data = "my momma told me to pick the very best one and you are not it"

    filepath = tmp_path / f'testfile.txt.{extension}'
    with anyopen(filepath, mode='w', compression=compression) as f:
        f.write(data)

    with anyopen(filepath, 'r', compression=compression) as f:
        data_out = f.read()

    assert data_out == data


@pytest.mark.parametrize('compression', ['bzip2', 'gzip'])
def test_stream_roundtrip(compression):
    """Test that roundtripping write/read to a stream works with `anyopen`

    """

    data = "my momma told me to pick the very best one and you are not it"

    with io.BytesIO() as stream:

        # write to stream
        with anyopen(stream, mode='w', compression=compression) as f:
            f.write(data)

        # start at the beginning
        stream.seek(0)

        # read from stream
        with anyopen(stream, 'r', compression=compression) as f:
            data_out = f.read()

        assert data_out == data

def test_stream_unsupported_compression():
    """Test that we throw a ValueError when an unsupported compression is used.

    """

    compression="fakez"

    data = b"my momma told me to pick the very best one and you are not it"

    with io.BytesIO() as stream:

        # write to stream
        stream.write(data)

        # start at the beginning
        stream.seek(0)

        # read from stream
        with pytest.raises(ValueError):
            with anyopen(stream, 'r', compression=compression) as f:
                data_out = f.read()
