import os
import platform
import shutil

import pytest

import py7zr
import py7zr.helpers

testdata_path = os.path.join(os.path.dirname(__file__), 'data')
targets = [("zstd", [{"id": py7zr.FILTER_ZSTD}]),
           ("bzip2", [{"id": py7zr.FILTER_BZIP2}]),
           ("lzma+bcj", [{"id": py7zr.FILTER_X86}, {"id": py7zr.FILTER_LZMA, "preset": 7}]),
           ("lzma2+bcj", [{"id": py7zr.FILTER_X86}, {"id": py7zr.FILTER_LZMA2, "preset": 7}]),
           ("bzip2+aes", [{"id": py7zr.FILTER_BZIP2}, {"id": py7zr.FILTER_CRYPTO_AES256_SHA256}]),
           ("lzma2+aes", [{"id": py7zr.FILTER_LZMA2, "preset": 7}, {"id": py7zr.FILTER_CRYPTO_AES256_SHA256}])
           ]


@pytest.mark.benchmark(group='compress')
@pytest.mark.parametrize("name, filters", targets)
def test_benchmark_filters_compress(tmp_path, benchmark, name, filters):

    def compressor(filters, password):
        with py7zr.SevenZipFile(tmp_path.joinpath('target.7z'), 'w', filters=filters, password=password) as szf:
            szf.writeall(tmp_path.joinpath('src'), 'src')

    def setup():
        if tmp_path.joinpath('target.7z').exists():
            tmp_path.joinpath('target.7z').unlink()

    with py7zr.SevenZipFile(os.path.join(testdata_path, 'mblock_1.7z'), 'r') as szf:
        szf.extractall(path=tmp_path.joinpath('src'))
    with py7zr.SevenZipFile(os.path.join(testdata_path, 'mblock_1.7z'), 'r') as szf:
        archive_info = szf.archiveinfo()
        source_size = archive_info.uncompressed
    if name.endswith('aes'):
        password = 'secret'
    else:
        password = None
    benchmark.extra_info['data_size'] = source_size
    benchmark.pedantic(compressor, setup=setup, args=[filters, password], iterations=1, rounds=3)
    benchmark.extra_info['ratio'] = str(tmp_path.joinpath('target.7z').stat().st_size / source_size)


@pytest.mark.benchmark(group='decompress')
@pytest.mark.parametrize("name, filters", targets)
def test_benchmark_filters_decompress(tmp_path, benchmark, name, filters):

    def decompressor(secret):
        with py7zr.SevenZipFile(tmp_path.joinpath('target.7z'), 'r', password=secret) as szf:
            szf.extractall(tmp_path.joinpath('tgt'))

    def setup():
        shutil.rmtree(tmp_path.joinpath('tgt'), ignore_errors=True)

    with py7zr.SevenZipFile(os.path.join(testdata_path, 'mblock_1.7z'), 'r') as szf:
        szf.extractall(path=tmp_path.joinpath('src'))
    with py7zr.SevenZipFile(os.path.join(testdata_path, 'mblock_1.7z'), 'r') as szf:
        archive_info = szf.archiveinfo()
        source_size = archive_info.uncompressed

    if name.endswith('aes'):
        password = 'secret'
    else:
        password = None
    with py7zr.SevenZipFile(tmp_path.joinpath('target.7z'), 'w', filters=filters, password=password) as szf:
        szf.writeall(tmp_path.joinpath('src'), 'src')
    benchmark.extra_info['data_size'] = source_size
    benchmark.extra_info['ratio'] = str(tmp_path.joinpath('target.7z').stat().st_size / source_size)
    benchmark.pedantic(decompressor, setup=setup, args=[password], iterations=1, rounds=3)


@pytest.mark.benchmark(group='calculate_key')
def test_benchmark_calculate_key1(benchmark):
    password = 'secret'.encode('utf-16LE')
    cycles = 19
    salt = b''
    expected = b'e\x11\xf1Pz<*\x98*\xe6\xde\xf4\xf6X\x18\xedl\xf2Be\x1a\xca\x19\xd1\\\xeb\xc6\xa6z\xe2\x89\x1d'
    key = benchmark(py7zr.helpers._calculate_key1, password, cycles, salt, 'sha256')
    assert key == expected


@pytest.mark.benchmark(group='calculate_key')
@pytest.mark.skipif(platform.python_implementation() == "PyPy", reason="Pypy has a bug around ctypes")
def test_benchmark_calculate_key2(benchmark):
    password = 'secret'.encode('utf-16LE')
    cycles = 19
    salt = b''
    expected = b'e\x11\xf1Pz<*\x98*\xe6\xde\xf4\xf6X\x18\xedl\xf2Be\x1a\xca\x19\xd1\\\xeb\xc6\xa6z\xe2\x89\x1d'
    key = benchmark(py7zr.helpers._calculate_key2, password, cycles, salt, 'sha256')
    assert key == expected


@pytest.mark.benchmark(group='calculate_key')
def test_benchmark_calculate_key3(benchmark):
    password = 'secret'.encode('utf-16LE')
    cycles = 19
    salt = b''
    expected = b'e\x11\xf1Pz<*\x98*\xe6\xde\xf4\xf6X\x18\xedl\xf2Be\x1a\xca\x19\xd1\\\xeb\xc6\xa6z\xe2\x89\x1d'
    key = benchmark(py7zr.helpers._calculate_key3, password, cycles, salt, 'sha256')
    assert key == expected
