import operator

import numpy as np
import pytest

import pandas.util._test_decorators as td

from pandas.core.dtypes.common import is_dtype_equal

import pandas as pd
import pandas._testing as tm
from pandas.core.arrays.string_arrow import ArrowStringArray, ArrowStringDtype

skip_if_no_pyarrow = td.skip_if_no("pyarrow", min_version="1.0.0")


@pytest.fixture(
    params=[
        # pandas\tests\arrays\string_\test_string.py:16: error: List item 1 has
        # incompatible type "ParameterSet"; expected
        # "Sequence[Collection[object]]"  [list-item]
        "string",
        pytest.param(
            "arrow_string", marks=skip_if_no_pyarrow
        ),  # type:ignore[list-item]
    ]
)
def dtype(request):
    return request.param


@pytest.fixture
def dtype_object(dtype):
    if dtype == "string":
        return pd.StringDtype
    else:
        return ArrowStringDtype


@pytest.fixture(
    params=[
        pd.arrays.StringArray,
        pytest.param(ArrowStringArray, marks=skip_if_no_pyarrow),
    ]
)
def cls(request):
    return request.param


def test_repr(dtype, request):
    if dtype == "arrow_string":
        reason = (
            "AssertionError: assert '      A\n0     a\n1  None\n2     b' "
            "== '      A\n0     a\n1  <NA>\n2     b'"
        )
        mark = pytest.mark.xfail(reason=reason)
        request.node.add_marker(mark)

    df = pd.DataFrame({"A": pd.array(["a", pd.NA, "b"], dtype=dtype)})
    expected = "      A\n0     a\n1  <NA>\n2     b"
    assert repr(df) == expected

    expected = "0       a\n1    <NA>\n2       b\nName: A, dtype: string"
    assert repr(df.A) == expected

    expected = "<StringArray>\n['a', <NA>, 'b']\nLength: 3, dtype: string"
    assert repr(df.A.array) == expected


def test_none_to_nan(cls):
    a = cls._from_sequence(["a", None, "b"])
    assert a[1] is not None
    assert a[1] is pd.NA


def test_setitem_validates(cls):
    arr = cls._from_sequence(["a", "b"])

    if cls is pd.arrays.StringArray:
        msg = "Cannot set non-string value '10' into a StringArray."
    else:
        msg = "Scalar must be NA or str"
    with pytest.raises(ValueError, match=msg):
        arr[0] = 10

    if cls is pd.arrays.StringArray:
        msg = "Must provide strings."
    else:
        msg = "Scalar must be NA or str"
    with pytest.raises(ValueError, match=msg):
        arr[:] = np.array([1, 2])


def test_setitem_with_scalar_string(dtype):
    # is_float_dtype considers some strings, like 'd', to be floats
    # which can cause issues.
    arr = pd.array(["a", "c"], dtype=dtype)
    arr[0] = "d"
    expected = pd.array(["d", "c"], dtype=dtype)
    tm.assert_extension_array_equal(arr, expected)


@pytest.mark.parametrize(
    "input, method",
    [
        (["a", "b", "c"], operator.methodcaller("capitalize")),
        (["a", "b", "c"], operator.methodcaller("capitalize")),
        (["a b", "a bc. de"], operator.methodcaller("capitalize")),
    ],
)
def test_string_methods(input, method, dtype, request):
    if dtype == "arrow_string":
        reason = "AttributeError: 'ArrowStringDtype' object has no attribute 'base'"
        mark = pytest.mark.xfail(reason=reason)
        request.node.add_marker(mark)

    a = pd.Series(input, dtype=dtype)
    b = pd.Series(input, dtype="object")
    result = method(a.str)
    expected = method(b.str)

    assert result.dtype.name == dtype
    tm.assert_series_equal(result.astype(object), expected)


def test_astype_roundtrip(dtype, request):
    if dtype == "arrow_string":
        reason = "ValueError: Could not convert object to NumPy datetime"
        mark = pytest.mark.xfail(reason=reason, raises=ValueError)
        request.node.add_marker(mark)
    else:
        mark = pytest.mark.xfail(
            reason="GH#36153 casting from StringArray to dt64 fails", raises=ValueError
        )
        request.node.add_marker(mark)

    ser = pd.Series(pd.date_range("2000", periods=12))
    ser[0] = None

    casted = ser.astype(dtype)
    assert is_dtype_equal(casted.dtype, dtype)

    result = casted.astype("datetime64[ns]")
    tm.assert_series_equal(result, ser)


def test_add(dtype, request):
    if dtype == "arrow_string":
        reason = (
            "TypeError: unsupported operand type(s) for +: 'ArrowStringArray' and "
            "'ArrowStringArray'"
        )
        mark = pytest.mark.xfail(reason=reason)
        request.node.add_marker(mark)

    a = pd.Series(["a", "b", "c", None, None], dtype=dtype)
    b = pd.Series(["x", "y", None, "z", None], dtype=dtype)

    result = a + b
    expected = pd.Series(["ax", "by", None, None, None], dtype=dtype)
    tm.assert_series_equal(result, expected)

    result = a.add(b)
    tm.assert_series_equal(result, expected)

    result = a.radd(b)
    expected = pd.Series(["xa", "yb", None, None, None], dtype=dtype)
    tm.assert_series_equal(result, expected)

    result = a.add(b, fill_value="-")
    expected = pd.Series(["ax", "by", "c-", "-z", None], dtype=dtype)
    tm.assert_series_equal(result, expected)


def test_add_2d(dtype, request):
    if dtype == "arrow_string":
        reason = "Failed: DID NOT RAISE <class 'ValueError'>"
        mark = pytest.mark.xfail(reason=reason)
        request.node.add_marker(mark)

    a = pd.array(["a", "b", "c"], dtype=dtype)
    b = np.array([["a", "b", "c"]], dtype=object)
    with pytest.raises(ValueError, match="3 != 1"):
        a + b

    s = pd.Series(a)
    with pytest.raises(ValueError, match="3 != 1"):
        s + b


def test_add_sequence(dtype, request):
    if dtype == "arrow_string":
        reason = (
            "TypeError: unsupported operand type(s) for +: 'ArrowStringArray' "
            "and 'list'"
        )
        mark = pytest.mark.xfail(reason=reason)
        request.node.add_marker(mark)

    a = pd.array(["a", "b", None, None], dtype=dtype)
    other = ["x", None, "y", None]

    result = a + other
    expected = pd.array(["ax", None, None, None], dtype=dtype)
    tm.assert_extension_array_equal(result, expected)

    result = other + a
    expected = pd.array(["xa", None, None, None], dtype=dtype)