"""
Tests encoding functionality during parsing
for all of the parsers defined in parsers.py
"""

from io import (
    BytesIO,
    StringIO,
    TextIOWrapper,
)
import sys
import tempfile

import numpy as np
import pytest

from pandas.errors import (
    EmptyDataError,
    ParserWarning,
)

import pandas as pd
import pandas._testing as tm

skip_pyarrow = pytest.mark.usefixtures("pyarrow_skip")


def test_bytes_io_input(all_parsers):
    encoding = "cp1255"
    parser = all_parsers

    data = BytesIO("שלום:1234\n562:123".encode(encoding))
    result = parser.read_csv(data, sep=":", encoding=encoding)

    expected = pd.DataFrame([[562, 123]], columns=["שלום", "1234"])
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize(
    "encoding_errors, expected", [("replace", "x?"), ("ignore", "x")]
)
def test_text_buffer_encoding_errors(all_parsers, encoding_errors, expected):
    # GH#70405 the C and pyarrow engines encode a text buffer to bytes and
    # must apply encoding_errors when doing so
    parser = all_parsers
    if parser.engine == "python":
        pytest.skip("python engine reads text buffers without encoding them")
    data = StringIO("a\nx\udcef\n")
    result = parser.read_csv(data, encoding_errors=encoding_errors)
    tm.assert_frame_equal(result, pd.DataFrame({"a": [expected]}))


@skip_pyarrow  # CSV parse error: Empty CSV file or block
def test_read_csv_unicode(all_parsers):
    parser = all_parsers
    data = BytesIO("\u0141aski, Jan;1".encode())

    result = parser.read_csv(data, sep=";", encoding="utf-8", header=None)
    expected = pd.DataFrame([["\u0141aski, Jan", 1]])
    tm.assert_frame_equal(result, expected)


@skip_pyarrow
@pytest.mark.parametrize("sep", [",", "\t"])
@pytest.mark.parametrize("encoding", ["utf-16", "utf-16le", "utf-16be"])
def test_utf16_bom_skiprows(all_parsers, sep, encoding, temp_file):
    # see gh-2298
    parser = all_parsers
    data = """skip this
skip this too
A,B,C
1,2,3
4,5,6""".replace(",", sep)
    kwargs = {"sep": sep, "skiprows": 2}
    utf8 = "utf-8"

    bytes_data = data.encode(encoding)

    with open(temp_file, "wb") as f:
        f.write(bytes_data)

    with TextIOWrapper(BytesIO(data.encode(utf8)), encoding=utf8) as bytes_buffer:
        result = parser.read_csv(temp_file, encoding=encoding, **kwargs)
        expected = parser.read_csv(bytes_buffer, encoding=utf8, **kwargs)
    tm.assert_frame_equal(result, expected)


def test_utf16_example(all_parsers, datapath):
    path = datapath("io", "parser", "data", "utf16_ex.txt")
    parser = all_parsers
    result = parser.read_csv(path, encoding="utf-16", sep="\t")
    assert len(result) == 50


def test_unicode_encoding(all_parsers, datapath):
    path = datapath("io", "parser", "data", "unicode_series.csv")
    parser = all_parsers

    result = parser.read_csv(path, header=None, encoding="latin-1")
    result = result.set_index(0)
    got = result[1][1632]

    expected = "\xc1 k\xf6ldum klaka (Cold Fever) (1994)"
    assert got == expected


@pytest.mark.parametrize(
    "data,kwargs,expected",
    [
        # Basic test
        ("a\n1", {}, [1]),
        # "Regular" quoting
        ('"a"\n1', {"quotechar": '"'}, [1]),
        # Test in a data row instead of header
        ("b\n1", {"names": ["a"]}, ["b", "1"]),
        # Test in empty data row with skipping
        ("\n1", {"names": ["a"], "skip_blank_lines": True}, [1]),
        # Test in empty data row without skipping
        (
            "\n1",
            {"names": ["a"], "skip_blank_lines": False},
            [np.nan, 1],
        ),
    ],
)
@pytest.mark.parametrize("encoding", ["utf-8", "utf-8-sig"])
def test_utf8_bom_removal(all_parsers, data, kwargs, expected, encoding):
    # see gh-4793
    parser = all_parsers
    bom = "\ufeff"

    def _encode_data_with_bom(_data):
        bom_data = (bom + _data).encode("utf-8")
        return BytesIO(bom_data)

    if (
        parser.engine == "pyarrow"
        and data == "\n1"
        and kwargs.get("skip_blank_lines", True)
    ):
        # CSV parse error: Empty CSV file or block: cannot infer number of columns
        pytest.skip(reason="https://github.com/apache/arrow/issues/38676")

    result = parser.read_csv(
        _encode_data_with_bom(data),
        encoding=encoding,
        **kwargs,
    )
    expected = pd.DataFrame({"a": expected})
    tm.assert_frame_equal(result, expected)


def test_latin1_bom_preserved(all_parsers):
    parser = all_parsers
    if parser.engine == "pyarrow":
        pytest.skip(reason="pyarrow does not support latin-1 encoding")

    bom_bytes = b"\xef\xbb\xbf"
    bom_text = bom_bytes.decode("latin1")
    data = b"a\n1"
    bio = BytesIO(bom_bytes + data)

    result = parser.read_csv(bio, encoding="latin1")
    expected = pd.DataFrame({f"{bom_text}a": [1]})
    tm.assert_frame_equal(result, expected)


def test_read_csv_utf_aliases(all_parsers, utf_value, encoding_fmt):
    # see gh-13549
    expected = pd.DataFrame({"mb_num": [4.8], "multibyte": ["test"]})
    parser = all_parsers

    encoding = encoding_fmt.format(utf_value)
    data = "mb_num,multibyte\n4.8,test".encode(encoding)

    result = parser.read_csv(BytesIO(data), encoding=encoding)
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize(
    "file_path,encoding",
    [
        (("io", "data", "csv", "test1.csv"), "utf-8"),
        (("io", "parser", "data", "unicode_series.csv"), "latin-1"),
        (("io", "parser", "data", "sauron.SHIFT_JIS.csv"), "shiftjis"),
    ],
)
def test_binary_mode_file_buffers(all_parsers, file_path, encoding, datapath):
    # gh-23779: Python csv engine shouldn't error on files opened in binary.
    # gh-31575: Python csv engine shouldn't error on files opened in raw binary.
    parser = all_parsers

    fpath = datapath(*file_path)
    expected = parser.read_csv(fpath, encoding=encoding)

    with open(fpath, encoding=encoding) as fa:
        result = parser.read_csv(fa)
        assert not fa.closed
    tm.assert_frame_equal(expected, result)

    with open(fpath, mode="rb") as fb:
        result = parser.read_csv(fb, encoding=encoding)
        assert not fb.closed
    tm.assert_frame_equal(expected, result)

    with open(fpath, mode="rb", buffering=0) as fb:
        result = parser.read_csv(fb, encoding=encoding)
        assert not fb.closed
    tm.assert_frame_equal(expected, result)


@pytest.mark.parametrize("pass_encoding", [True, False])
def test_encoding_temp_file(
    all_parsers, utf_value, encoding_fmt, pass_encoding, temp_file
):
    # see gh-24130
    parser = all_parsers
    encoding = encoding_fmt.format(utf_value)

    if parser.engine == "pyarrow" and pass_encoding is True and utf_value in [16, 32]:
        # GH#24130: pyarrow's CSV reader splits the input into byte-blocks
        # before decoding, which corrupts multi-byte utf-16/utf-32 sequences
        # (raises "straddling object straddles two block boundaries"). The
        # error path is slow (~9s per case), so skip rather than xfail.
        pytest.skip("pyarrow utf-16/32 block-boundary error is slow to surface")

    expected = pd.DataFrame({"foo": ["bar"]})

    with temp_file.open(mode="w+", encoding=encoding) as f:
        f.write("foo\nbar")
        f.seek(0)

        result = parser.read_csv(f, encoding=encoding if pass_encoding else None)
        tm.assert_frame_equal(result, expected)


def test_encoding_named_temp_file(all_parsers):
    # see gh-31819
    parser = all_parsers
    encoding = "shift-jis"

    title = "てすと"
    data = "こむ"

    expected = pd.DataFrame({title: [data]})

    with tempfile.NamedTemporaryFile() as f:
        f.write(f"{title}\n{data}".encode(encoding))

        f.seek(0)

        result = parser.read_csv(f, encoding=encoding)
        tm.assert_frame_equal(result, expected)
        assert not f.closed


@pytest.mark.parametrize(
    "encoding", ["utf-8", "utf-16", "utf-16-be", "utf-16-le", "utf-32"]
)
def test_parse_encoded_special_characters(encoding):
    # GH16218 Verify parsing of data with encoded special characters
    # Data contains a Unicode 'FULLWIDTH COLON' (U+FF1A) at position (0,"a")
    data = "a\tb\n：foo\t0\nbar\t1\nbaz\t2"  # noqa: RUF001
    encoded_data = BytesIO(data.encode(encoding))
    result = pd.read_csv(encoded_data, delimiter="\t", encoding=encoding)

    expected = pd.DataFrame(
        data=[["：foo", 0], ["bar", 1], ["baz", 2]],  # noqa: RUF001
        columns=["a", "b"],
    )
    tm.assert_frame_equal(result, expected)


@pytest.mark.parametrize("encoding", ["utf-8", None, "utf-16", "cp1255", "latin-1"])
def test_encoding_memory_map(all_parsers, encoding, temp_file):
    # GH40986
    parser = all_parsers
    expected = pd.DataFrame(
        {
            "name": ["Raphael", "Donatello", "Miguel Angel", "Leonardo"],
            "mask": ["red", "purple", "orange", "blue"],
            "weapon": ["sai", "bo staff", "nunchunk", "katana"],
        }
    )
    expected.to_csv(temp_file, index=False, encoding=encoding)

    if parser.engine == "pyarrow":
        msg = "The 'memory_map' option is not supported with the 'pyarrow' engine"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(temp_file, encoding=encoding, memory_map=True)
        return

    df = parser.read_csv(temp_file, encoding=encoding, memory_map=True)
    tm.assert_frame_equal(df, expected)


def test_chunk_splits_multibyte_char(all_parsers, temp_file):
    """
    Chunk splits a multibyte character with memory_map=True

    GH 43540
    """
    parser = all_parsers
    # DEFAULT_CHUNKSIZE = 262144, defined in parsers.pyx
    df = pd.DataFrame(data=["a" * 127] * 2048)

    # Put two-bytes utf-8 encoded character "ą" at the end of chunk
    # utf-8 encoding of "ą" is b'\xc4\x85'
    df.iloc[2047] = "a" * 127 + "ą"
    df.to_csv(temp_file, index=False, header=False, encoding="utf-8")

    if parser.engine == "pyarrow":
        msg = "The 'memory_map' option is not supported with the 'pyarrow' engine"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(temp_file, header=None, memory_map=True)
        return

    dfr = parser.read_csv(temp_file, header=None, memory_map=True)
    tm.assert_frame_equal(dfr, df)


def test_readcsv_memmap_utf8(all_parsers, temp_file):
    """
    GH 43787

    Test correct handling of UTF-8 chars when memory_map=True and encoding is UTF-8
    """
    lines = []
    line_length = 128
    start_char = " "
    end_char = "\U00010080"
    # This for loop creates a list of 128-char strings
    # consisting of consecutive Unicode chars
    for lnum in range(ord(start_char), ord(end_char), line_length):
        line = "".join([chr(c) for c in range(lnum, lnum + 0x80)]) + "\n"
        try:
            line.encode("utf-8")
        except UnicodeEncodeError:
            continue
        lines.append(line)
    parser = all_parsers
    df = pd.DataFrame(lines)
    df.to_csv(temp_file, index=False, header=False, encoding="utf-8")

    if parser.engine == "pyarrow":
        msg = "The 'memory_map' option is not supported with the 'pyarrow' engine"
        with pytest.raises(ValueError, match=msg):
            parser.read_csv(temp_file, header=None, memory_map=True, encoding="utf-8")
        return

    dfr = parser.read_csv(temp_file, header=None, memory_map=True, encoding="utf-8")
    tm.assert_frame_equal(df, dfr)


@pytest.mark.parametrize("mode", ["w+b", "w+t"])
def test_not_readable(all_parsers, mode):
    # GH43439
    parser = all_parsers
    content = b"abcd"
    if "t" in mode:
        content = "abcd"
    with tempfile.SpooledTemporaryFile(mode=mode, encoding="utf-8") as handle:
        handle.write(content)
        handle.seek(0)
        if parser.engine == "pyarrow":
            # pyarrow's CSV reader cannot read from a SpooledTemporaryFile
            if "t" in mode:
                msg = "The 'pyarrow' engine can only read from a binary file object"
                with pytest.raises(TypeError, match=msg):
                    parser.read_csv(handle)
            else:
                msg = "No columns to parse from file"
                with pytest.raises(EmptyDataError, match=msg):
                    parser.read_csv(handle)
            return
        df = parser.read_csv(handle)
    expected = pd.DataFrame([], columns=["abcd"])
    tm.assert_frame_equal(df, expected)


def test_sep_encodeable_check_ignores_filesystem_encoding(monkeypatch, temp_file):
    # GH#46456 whether the "c" engine can handle the separator depends on utf-8,
    # not on the platform's filesystem encoding
    monkeypatch.setattr(sys, "getfilesystemencoding", lambda: "latin-1")
    temp_file.write_text(
        "key\u00a5value\ntables\u00a5rectangular\n", encoding="latin-1"
    )

    with tm.assert_produces_warning(
        ParserWarning, match="encoded in utf-8", check_stacklevel=False
    ):
        result = pd.read_csv(temp_file, sep="\u00a5", encoding="latin-1")
    expected = pd.DataFrame([["tables", "rectangular"]], columns=["key", "value"])
    tm.assert_frame_equal(result, expected)

    with pytest.raises(ValueError, match="encoded in utf-8"):
        pd.read_csv(temp_file, sep="\u00a5", encoding="latin-1", engine="c")


def test_sep_unencodeable_falls_back(temp_file):
    # GH#46456 a separator that cannot be encoded at all falls back instead of raising
    temp_file.write_text("a,b\n1,2\n", encoding="utf-8")

    with tm.assert_produces_warning(
        ParserWarning, match="encoded in utf-8", check_stacklevel=False
    ):
        result = pd.read_csv(temp_file, sep="\udcff")
    expected = pd.DataFrame({"a,b": ["1,2"]})
    tm.assert_frame_equal(result, expected)