"""
Tests compressed data parsing functionality for all
of the parsers defined in parsers.py
"""

from pathlib import Path
import tarfile
import zipfile

import pytest

import pandas as pd
import pandas._testing as tm


@pytest.fixture(params=[True, False])
def buffer(request):
    return request.param


@pytest.fixture
def parser_and_data(all_parsers, csv1):
    parser = all_parsers

    with open(csv1, "rb") as f:
        data = f.read()
    expected = parser.read_csv(csv1)

    return parser, data, expected


@pytest.mark.parametrize("compression", ["zip", "infer", "zip2"])
def test_zip(tmp_path, parser_and_data, compression):
    parser, data, expected = parser_and_data

    path = tmp_path / "test_file.zip"
    with zipfile.ZipFile(path, mode="w") as tmp:
        tmp.writestr("test_file", data)

    if compression == "zip2":
        with open(path, "rb") as f:
            result = parser.read_csv(f, compression="zip")
    else:
        result = parser.read_csv(path, compression=compression)

    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("compression", ["zip", "infer"])
def test_zip_error_multiple_files(tmp_path, parser_and_data, compression):
    parser, data, expected = parser_and_data

    path = tmp_path / "combined_zip.zip"
    inner_file_names = ["test_file", "second_file"]

    with zipfile.ZipFile(path, mode="w") as tmp:
        for file_name in inner_file_names:
            tmp.writestr(file_name, data)

    with pytest.raises(ValueError, match="Multiple files"):
        parser.read_csv(path, compression=compression)


def test_zip_error_no_files(tmp_path, parser_and_data):
    parser, _, _ = parser_and_data

    path = tmp_path / "test_file.zip"
    with zipfile.ZipFile(path, mode="w"):
        pass

    with pytest.raises(ValueError, match="Zero files"):
        parser.read_csv(path, compression="zip")


def test_zip_error_invalid_zip(tmp_path, parser_and_data):
    parser, _, _ = parser_and_data

    path = tmp_path / "invalid_file.zip"
    path.touch()
    with open(path, "rb") as f:
        with pytest.raises(zipfile.BadZipFile, match="File is not a zip file"):
            parser.read_csv(f, compression="zip")


@pytest.mark.parametrize("filename", [None, "test.{ext}"])
def test_compression(
    tmp_path,
    request,
    parser_and_data,
    compression_only,
    buffer,
    filename,
    compression_to_extension,
):
    parser, data, expected = parser_and_data
    compress_type = compression_only

    ext = compression_to_extension[compress_type]
    filename = filename if filename is None else filename.format(ext=ext)

    if filename and buffer:
        request.applymarker(
            pytest.mark.xfail(
                reason="Cannot deduce compression from buffer of compressed data."
            )
        )

    path = tmp_path / filename if filename else tmp_path / "test_file"
    tm.write_to_compressed(compress_type, path, data)
    compression = "infer" if filename else compress_type

    if buffer:
        with open(path, "rb") as f:
            result = parser.read_csv(f, compression=compression)
    else:
        result = parser.read_csv(path, compression=compression)

    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("ext", [None, "gz", "bz2"])
def test_infer_compression(all_parsers, csv1, buffer, ext):
    # see gh-9770
    parser = all_parsers
    kwargs = {"index_col": 0, "parse_dates": True}

    expected = parser.read_csv(csv1, **kwargs)
    kwargs["compression"] = "infer"

    if buffer:
        with open(csv1, encoding="utf-8") as f:
            result = parser.read_csv(f, **kwargs)
    else:
        ext = "." + ext if ext else ""
        result = parser.read_csv(csv1 + ext, **kwargs)

    tm.assert_frame_equal(result, expected)


def test_compression_utf_encoding(all_parsers, datapath, utf_value, encoding_fmt):
    # see gh-18071, gh-24130
    parser = all_parsers
    encoding = encoding_fmt.format(utf_value)
    path = datapath("io", "parser", "data", f"utf{utf_value}_ex_small.zip")

    result = parser.read_csv(path, encoding=encoding, compression="zip", sep="\t")
    expected = pd.DataFrame(
        {
            "Country": ["Venezuela", "Venezuela"],
            "Twitter": ["Hugo Chávez Frías", "Henrique Capriles R."],
        }
    )

    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("invalid_compression", ["sfark", "bz3", "zipper"])
def test_invalid_compression(all_parsers, invalid_compression):
    parser = all_parsers
    compress_kwargs = {"compression": invalid_compression}

    msg = f"Unrecognized compression type: {invalid_compression}"

    with pytest.raises(ValueError, match=msg):
        parser.read_csv("test_file.zip", **compress_kwargs)


def test_compression_tar_archive(all_parsers, datapath):
    parser = all_parsers
    path = datapath("io", "parser", "data", "tar_csv.tar.gz")
    df = parser.read_csv(path)
    assert list(df.columns) == ["a"]


def test_ignore_compression_extension(tmp_path, all_parsers):
    parser = all_parsers
    df = pd.DataFrame({"a": [0, 1]})

    path_csv = tmp_path / "test.csv"
    path_zip = tmp_path / "test.csv.zip"
    # make sure to create un-compressed file with zip extension
    df.to_csv(path_csv, index=False)
    Path(path_zip).write_text(
        Path(path_csv).read_text(encoding="utf-8"), encoding="utf-8"
    )

    tm.assert_frame_equal(parser.read_csv(path_zip, compression=None), df)


def test_writes_tar_gz(tmp_path, all_parsers):
    parser = all_parsers
    data = pd.DataFrame(
        {
            "Country": ["Venezuela", "Venezuela"],
            "Twitter": ["Hugo Chávez Frías", "Henrique Capriles R."],
        }
    )
    tar_path = tmp_path / "test.tar.gz"
    data.to_csv(tar_path, index=False)

    # test that read_csv infers .tar.gz to gzip:
    tm.assert_frame_equal(parser.read_csv(tar_path), data)

    # test that file is indeed gzipped:
    with tarfile.open(tar_path, "r:gz") as tar:
        result = parser.read_csv(
            tar.extractfile(tar.getnames()[0]), compression="infer"
        )
        tm.assert_frame_equal(result, data)