import io

import numpy as np
import pytest

from pandas._config import using_string_dtype

from pandas.compat import HAS_PYARROW

import pandas as pd
import pandas._testing as tm
from pandas.util import _test_decorators as td
from pandas.util.version import Version


@pytest.fixture
def fsspectest():
    pytest.importorskip("fsspec")
    from fsspec import register_implementation
    from fsspec.implementations.memory import MemoryFileSystem
    from fsspec.registry import _registry as registry

    class TestMemoryFS(MemoryFileSystem):
        protocol = "testmem"
        test = [None]

        def __init__(self, **kwargs) -> None:
            self.test[0] = kwargs.pop("test", None)
            super().__init__(**kwargs)

    register_implementation("testmem", TestMemoryFS, clobber=True)
    yield TestMemoryFS()
    registry.pop("testmem", None)
    TestMemoryFS.test[0] = None
    TestMemoryFS.store.clear()


@pytest.fixture
def df1():
    return pd.DataFrame(
        {
            "int": [1, 3],
            "float": [2.0, np.nan],
            "str": ["t", "s"],
            "dt": pd.date_range("2018-06-18", periods=2),
        }
    )


@pytest.fixture
def cleared_fs():
    fsspec = pytest.importorskip("fsspec")

    memfs = fsspec.filesystem("memory")
    yield memfs
    memfs.store.clear()


def test_read_csv(cleared_fs, df1):
    text = str(df1.to_csv(index=False)).encode()
    with cleared_fs.open("test/test.csv", "wb") as w:
        w.write(text)
    df2 = pd.read_csv("memory://test/test.csv", parse_dates=["dt"])

    expected = df1.copy()
    expected["dt"] = expected["dt"].astype("M8[us]")
    tm.assert_frame_equal(df2, expected)


def test_reasonable_error(monkeypatch, cleared_fs):
    from fsspec.registry import known_implementations

    with pytest.raises(ValueError, match="nosuchprotocol"):
        pd.read_csv("nosuchprotocol://test/test.csv")
    err_msg = "test error message"
    monkeypatch.setitem(
        known_implementations,
        "couldexist",
        {"class": "unimportable.CouldExist", "err": err_msg},
    )
    with pytest.raises(ImportError, match=err_msg):
        pd.read_csv("couldexist://test/test.csv")


def test_to_csv(cleared_fs, df1):
    df1.to_csv("memory://test/test.csv", index=True)

    df2 = pd.read_csv("memory://test/test.csv", parse_dates=["dt"], index_col=0)

    expected = df1.copy()
    expected["dt"] = expected["dt"].astype("M8[us]")
    tm.assert_frame_equal(df2, expected)


@pytest.mark.filterwarnings(
    "ignore:The default engine for reading:pandas.errors.Pandas4Warning"
)
def test_to_excel(cleared_fs, df1):
    pytest.importorskip("openpyxl")
    ext = "xlsx"
    path = f"memory://test/test.{ext}"
    df1.to_excel(path, index=True)

    df2 = pd.read_excel(path, parse_dates=["dt"], index_col=0)

    expected = df1.copy()
    expected["dt"] = expected["dt"].astype("M8[us]")
    tm.assert_frame_equal(df2, expected)


@pytest.mark.parametrize("binary_mode", [False, True])
def test_to_csv_fsspec_object(cleared_fs, binary_mode, df1):
    fsspec = pytest.importorskip("fsspec")

    path = "memory://test/test.csv"
    mode = "wb" if binary_mode else "w"
    with fsspec.open(path, mode=mode).open() as fsspec_object:
        df1.to_csv(fsspec_object, index=True)
        assert not fsspec_object.closed

    mode = mode.replace("w", "r")
    with fsspec.open(path, mode=mode) as fsspec_object:
        df2 = pd.read_csv(
            fsspec_object,
            parse_dates=["dt"],
            index_col=0,
        )
        assert not fsspec_object.closed

    expected = df1.copy()
    expected["dt"] = expected["dt"].astype("M8[us]")
    tm.assert_frame_equal(df2, expected)


def test_csv_options(fsspectest):
    df = pd.DataFrame({"a": [0]})
    df.to_csv(
        "testmem://test/test.csv", storage_options={"test": "csv_write"}, index=False
    )
    assert fsspectest.test[0] == "csv_write"
    pd.read_csv("testmem://test/test.csv", storage_options={"test": "csv_read"})
    assert fsspectest.test[0] == "csv_read"


def test_read_table_options(fsspectest):
    # GH #39167
    df = pd.DataFrame({"a": [0]})
    df.to_csv(
        "testmem://test/test.csv", storage_options={"test": "csv_write"}, index=False
    )
    assert fsspectest.test[0] == "csv_write"
    pd.read_table("testmem://test/test.csv", storage_options={"test": "csv_read"})
    assert fsspectest.test[0] == "csv_read"


@pytest.mark.filterwarnings(
    "ignore:The default engine for reading:pandas.errors.Pandas4Warning"
)
def test_excel_options(fsspectest):
    pytest.importorskip("openpyxl")
    extension = "xlsx"

    df = pd.DataFrame({"a": [0]})

    path = f"testmem://test/test.{extension}"

    df.to_excel(path, storage_options={"test": "write"}, index=False)
    assert fsspectest.test[0] == "write"
    pd.read_excel(path, storage_options={"test": "read"})
    assert fsspectest.test[0] == "read"


@pytest.mark.filterwarnings(
    "ignore:The 'fastparquet' engine is deprecated:DeprecationWarning"
)
def test_to_parquet_new_file(cleared_fs, df1, request):
    """Regression test for writing to a not-yet-existent GCS Parquet file."""
    fp = pytest.importorskip("fastparquet")

    request.applymarker(
        pytest.mark.xfail(
            using_string_dtype()
            and HAS_PYARROW
            and Version(fp.__version__) < Version("2026.5.0"),
            reason="TODO(infer_string) fastparquet",
        )
    )

    df1.to_parquet(
        "memory://test/test.csv", index=True, engine="fastparquet", compression=None
    )


def test_arrowparquet_options(fsspectest):
    """Regression test for writing to a not-yet-existent GCS Parquet file."""
    pytest.importorskip("pyarrow")
    df = pd.DataFrame({"a": [0]})
    df.to_parquet(
        "testmem://test/test.csv",
        engine="pyarrow",
        compression=None,
        storage_options={"test": "parquet_write"},
    )
    assert fsspectest.test[0] == "parquet_write"
    pd.read_parquet(
        "testmem://test/test.csv",
        engine="pyarrow",
        storage_options={"test": "parquet_read"},
    )
    assert fsspectest.test[0] == "parquet_read"


def test_arrowparquet_falls_back_to_fsspec(cleared_fs, df1, monkeypatch):
    # GH#58078 pyarrow cannot build a filesystem for "hdfs:///" but fsspec can
    pytest.importorskip("pyarrow")
    from fsspec.implementations.memory import MemoryFileSystem
    from fsspec.registry import _registry as registry

    monkeypatch.setitem(registry, "hdfs", MemoryFileSystem)
    path = "hdfs:///test/test.parquet"
    df1.to_parquet(path, engine="pyarrow")
    result = pd.read_parquet(path, engine="pyarrow")
    tm.assert_frame_equal(result, df1)


def test_arrowparquet_falls_back_to_fsspec_on_arrow_exception(
    cleared_fs, df1, monkeypatch
):
    # GH#58078 the fallback also covers non-ArrowInvalid subclasses of ArrowException,
    # e.g. a pyarrow build without HDFS support raises ArrowNotImplementedError
    pa = pytest.importorskip("pyarrow")
    pa_fs = pytest.importorskip("pyarrow.fs")
    from fsspec.implementations.memory import MemoryFileSystem
    from fsspec.registry import _registry as registry

    real_file_system = pa_fs.FileSystem

    class FailingFileSystem:
        @staticmethod
        def from_uri(uri):
            # restore the real FileSystem before pyarrow's own isinstance
            # checks run on the fsspec-fallback filesystem later in this call
            pa_fs.FileSystem = real_file_system
            raise pa.ArrowNotImplementedError("HDFS support not built")

    monkeypatch.setitem(registry, "hdfs", MemoryFileSystem)
    path = "hdfs:///test/test.parquet"

    monkeypatch.setattr(pa_fs, "FileSystem", FailingFileSystem)
    df1.to_parquet(path, engine="pyarrow")

    monkeypatch.setattr(pa_fs, "FileSystem", FailingFileSystem)
    result = pd.read_parquet(path, engine="pyarrow")
    tm.assert_frame_equal(result, df1)


@pytest.mark.parametrize(
    "path", ["s3://missing/test.parquet", "hdfs://namenode:8020/test.parquet"]
)
def test_arrowparquet_oserror_raises(monkeypatch, path):
    # GH#58078 only "hdfs:///" falls back to fsspec on OSError
    pa_fs = pytest.importorskip("pyarrow.fs")

    class FailingFileSystem:
        @staticmethod
        def from_uri(uri):
            raise OSError("pyarrow error")

    monkeypatch.setattr(pa_fs, "FileSystem", FailingFileSystem)
    with pytest.raises(OSError, match="pyarrow error"):
        pd.read_parquet(path, engine="pyarrow")


def test_arrowparquet_fsspec_fallback_fails_keeps_pyarrow_error(monkeypatch):
    # GH#58078 if the fsspec fallback also fails, chain pyarrow's error
    pa_fs = pytest.importorskip("pyarrow.fs")
    pytest.importorskip("fsspec")
    from fsspec.implementations.memory import MemoryFileSystem
    from fsspec.registry import _registry as registry

    class FailingFileSystem:
        @staticmethod
        def from_uri(uri):
            raise OSError("Unable to load libjvm")

    class FailingHadoopFileSystem(MemoryFileSystem):
        def __init__(self, *args, **kwargs) -> None:
            raise OSError("Prior attempt to load libhdfs failed")

    monkeypatch.setattr(pa_fs, "FileSystem", FailingFileSystem)
    monkeypatch.setitem(registry, "hdfs", FailingHadoopFileSystem)
    with pytest.raises(OSError, match="Prior attempt") as excinfo:
        pd.read_parquet("hdfs:///test.parquet", engine="pyarrow")
    assert "Unable to load libjvm" in str(excinfo.value.__cause__)


@pytest.mark.filterwarnings(
    "ignore:The 'fastparquet' engine is deprecated:DeprecationWarning"
)
def test_fastparquet_options(fsspectest):
    """Regression test for writing to a not-yet-existent GCS Parquet file."""
    pytest.importorskip("fastparquet")

    df = pd.DataFrame({"a": [0]})
    df.to_parquet(
        "testmem://test/test.csv",
        engine="fastparquet",
        compression=None,
        storage_options={"test": "parquet_write"},
    )
    assert fsspectest.test[0] == "parquet_write"
    pd.read_parquet(
        "testmem://test/test.csv",
        engine="fastparquet",
        storage_options={"test": "parquet_read"},
    )
    assert fsspectest.test[0] == "parquet_read"


@pytest.mark.single_cpu
@pytest.mark.parametrize("compression_suffix", ["", ".gz", ".bz2"])
def test_from_s3_csv(s3_bucket_public_with_data, s3so, tips_file, compression_suffix):
    pytest.importorskip("s3fs")
    df_from_s3 = pd.read_csv(
        f"s3://{s3_bucket_public_with_data.name}/tips.csv{compression_suffix}",
        storage_options=s3so,
    )
    df_from_local = pd.read_csv(tips_file)
    tm.assert_equal(df_from_s3, df_from_local)


@pytest.mark.single_cpu
@pytest.mark.parametrize("protocol", ["s3", "s3a", "s3n"])
def test_s3_protocols(s3_bucket_public_with_data, s3so, tips_file, protocol):
    pytest.importorskip("s3fs")
    df_from_s3 = pd.read_csv(
        f"{protocol}://{s3_bucket_public_with_data.name}/tips.csv",
        storage_options=s3so,
    )
    df_from_local = pd.read_csv(tips_file)
    tm.assert_equal(df_from_s3, df_from_local)


@pytest.mark.filterwarnings(
    "ignore:The 'fastparquet' engine is deprecated:DeprecationWarning"
)
@pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string) fastparquet")
@pytest.mark.single_cpu
def test_s3_parquet(s3_bucket_public, s3so, df1):
    pytest.importorskip("fastparquet")
    pytest.importorskip("s3fs")

    fn = f"s3://{s3_bucket_public.name}/test.parquet"
    df1.to_parquet(
        fn, index=False, engine="fastparquet", compression=None, storage_options=s3so
    )
    df2 = pd.read_parquet(fn, engine="fastparquet", storage_options=s3so)
    tm.assert_equal(df1, df2)


@td.skip_if_installed("fsspec")
def test_not_present_exception():
    msg = (
        "`Import fsspec` failed.  Use pip, conda, or your preferred package "
        "management tool to install the fsspec package."
    )
    with pytest.raises(ImportError, match=msg):
        pd.read_csv("memory://test/test.csv")


def test_feather_options(fsspectest):
    pytest.importorskip("pyarrow")
    df = pd.DataFrame({"a": [0]})
    df.to_feather("testmem://mockfile", storage_options={"test": "feather_write"})
    assert fsspectest.test[0] == "feather_write"
    out = pd.read_feather(
        "testmem://mockfile", storage_options={"test": "feather_read"}
    )
    assert fsspectest.test[0] == "feather_read"
    tm.assert_frame_equal(df, out)


def test_pickle_options(fsspectest):
    df = pd.DataFrame({"a": [0]})
    df.to_pickle("testmem://mockfile", storage_options={"test": "pickle_write"})
    assert fsspectest.test[0] == "pickle_write"
    out = pd.read_pickle("testmem://mockfile", storage_options={"test": "pickle_read"})
    assert fsspectest.test[0] == "pickle_read"
    tm.assert_frame_equal(df, out)


def test_json_options(fsspectest, compression):
    df = pd.DataFrame({"a": [0]})
    df.to_json(
        "testmem://mockfile",
        compression=compression,
        storage_options={"test": "json_write"},
    )
    assert fsspectest.test[0] == "json_write"
    out = pd.read_json(
        "testmem://mockfile",
        compression=compression,
        storage_options={"test": "json_read"},
    )
    assert fsspectest.test[0] == "json_read"
    tm.assert_frame_equal(df, out)


def test_stata_options(fsspectest):
    df = pd.DataFrame({"a": [0]})
    df.to_stata(
        "testmem://mockfile", storage_options={"test": "stata_write"}, write_index=False
    )
    assert fsspectest.test[0] == "stata_write"
    out = pd.read_stata("testmem://mockfile", storage_options={"test": "stata_read"})
    assert fsspectest.test[0] == "stata_read"
    tm.assert_frame_equal(df, out.astype("int64"))


def test_markdown_options(fsspectest):
    pytest.importorskip("tabulate")
    df = pd.DataFrame({"a": [0]})
    df.to_markdown("testmem://mockfile", storage_options={"test": "md_write"})
    assert fsspectest.test[0] == "md_write"
    assert fsspectest.cat("testmem://mockfile")


def test_non_fsspec_options():
    pytest.importorskip("pyarrow")
    with pytest.raises(ValueError, match="storage_options"):
        pd.read_csv("localfile", storage_options={"a": True})
    with pytest.raises(ValueError, match="storage_options"):
        # separate test for parquet, which has a different code path
        pd.read_parquet("localfile", storage_options={"a": True})
    by = io.BytesIO()

    with pytest.raises(ValueError, match="storage_options"):
        pd.read_csv(by, storage_options={"a": True})

    df = pd.DataFrame({"a": [0]})
    with pytest.raises(ValueError, match="storage_options"):
        df.to_parquet("nonfsspecpath", storage_options={"a": True})