import shutil
import subprocess
import tempfile
import time

import pandas as pd
import pyarrow as pa
import pytest

import ray
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.mock_http_server import *  # noqa
from ray.tests.conftest import *  # noqa

# To run tests locally, make sure you install mongodb-org and have mongod
# available on your PATH. Started directly since mongodb-org has no SysV init
# script. See https://hub.docker.com/_/mongo


@pytest.fixture
def start_mongo():
    import pymongo
    import pymongo.errors

    dbpath = tempfile.mkdtemp(prefix="mongod_test_")
    proc = subprocess.Popen(
        ["mongod", "--dbpath", dbpath, "--bind_ip", "127.0.0.1"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    # Wait for mongod to accept connections.
    mongo_url = "mongodb://localhost:27017"
    for _ in range(30):
        if proc.poll() is not None:
            raise RuntimeError(
                f"mongod exited unexpectedly (returncode={proc.returncode})"
            )
        try:
            client = pymongo.MongoClient(mongo_url, serverSelectionTimeoutMS=1000)
            client.admin.command("ping")
            break
        except pymongo.errors.PyMongoError:
            time.sleep(0.5)
    else:
        proc.kill()
        raise RuntimeError("mongod failed to start")

    # Make sure a clean slate for each test by dropping
    # previously created ones (if any).
    for db in client.list_database_names():
        # Keep the MongoDB default databases.
        if db not in ("admin", "local", "config"):
            client.drop_database(db)
    yield client, mongo_url

    proc.terminate()
    proc.wait(timeout=10)
    shutil.rmtree(dbpath)


def test_read_write_mongo(ray_start_regular_shared, start_mongo):
    from pymongo.errors import ServerSelectionTimeoutError
    from pymongoarrow.api import Schema

    client, mongo_url = start_mongo
    foo_db = "foo-db"
    foo_collection = "foo-collection"
    foo = client[foo_db][foo_collection]
    foo.delete_many({})

    # Read nonexistent URI.
    with pytest.raises(ServerSelectionTimeoutError):
        ds = ray.data.read_mongo(
            uri="nonexistent-uri",
            database=foo_db,
            collection=foo_collection,
        )
    # Read nonexistent database.
    with pytest.raises(ValueError):
        ds = ray.data.read_mongo(
            uri=mongo_url,
            database="nonexistent-db",
            collection=foo_collection,
        )
    # Read nonexistent collection.
    with pytest.raises(ValueError):
        ds = ray.data.read_mongo(
            uri=mongo_url,
            database=foo_db,
            collection="nonexistent-collection",
        )

    # Inject 5 test docs.
    docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5)]
    df = pd.DataFrame(docs).astype({"int_field": "int32"})
    foo.insert_many(docs)

    # Read a non-empty database, with schema specified.
    schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        schema=schema,
        override_num_blocks=2,
    )
    assert ds._block_num_rows() == [3, 2]
    ds_schema = ds.schema()
    assert ds_schema.names == ["float_field", "int_field"]
    assert ds_schema.types == [pa.float64(), pa.int32()]
    result = ds.to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Read with schema inference, which will read all columns (including the auto
    # generated internal column "_id").
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        override_num_blocks=2,
    )
    assert ds._block_num_rows() == [3, 2]
    assert ds.count() == 5
    assert ds.schema().names == ["_id", "float_field", "int_field"]
    # We are not testing the datatype of _id here, because it varies per platform
    assert ds.schema().types[1:] == [
        pa.float64(),
        pa.int32(),
    ]
    result = ds.drop_columns(["_id"]).to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Read a subset of the collection.
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
        override_num_blocks=2,
    )
    assert ds._block_num_rows() == [2, 1]
    assert ds.count() == 3
    assert ds.schema().names == ["_id", "float_field", "int_field"]
    df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())

    # Read with auto-tuned parallelism.
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
    )

    assert ds.count() == 5
    assert ds.schema().names == ["_id", "float_field", "int_field"]
    # We are not testing the datatype of _id here, because it varies per platform
    assert ds.schema().types[1:] == [
        pa.float64(),
        pa.int32(),
    ]
    result = ds.drop_columns(["_id"]).to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Read with a parallelism larger than number of rows.
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        override_num_blocks=1000,
    )

    assert ds.count() == 5
    assert ds.schema().names == ["_id", "float_field", "int_field"]
    # We are not testing the datatype of _id here, because it varies per platform
    assert ds.schema().types[1:] == [
        pa.float64(),
        pa.int32(),
    ]
    result = ds.drop_columns(["_id"]).to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Add a column and then write back to MongoDB.
    # Inject 2 more test docs.
    new_docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5, 7)]
    new_df = pd.DataFrame(new_docs).astype({"int_field": "int32"})
    ds2 = ray.data.from_pandas(new_df)
    ds2.write_mongo(uri=mongo_url, database=foo_db, collection=foo_collection)

    # Read again to verify the content.
    expected_ds = ds.drop_columns(["_id"]).union(ds2)
    ds3 = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
    )
    ds3.drop_columns(["_id"]).to_pandas().equals(expected_ds.to_pandas())

    # Destination database doesn't exist.
    with pytest.raises(ValueError):
        ray.data.range(10).write_mongo(
            uri=mongo_url, database="nonexistent-db", collection=foo_collection
        )
    # Destination collection doesn't exist.
    with pytest.raises(ValueError):
        ray.data.range(10).write_mongo(
            uri=mongo_url, database=foo_db, collection="nonexistent-collection"
        )


def test_mongo_datasource(ray_start_regular_shared, start_mongo):
    from pymongoarrow.api import Schema

    client, mongo_url = start_mongo
    foo_db = "foo-db"
    foo_collection = "foo-collection"
    foo = client[foo_db][foo_collection]
    foo.delete_many({})

    # Inject 5 test docs.
    docs = [{"float_field": 2.0 * key, "int_field": key} for key in range(5)]
    df = pd.DataFrame(docs).astype({"int_field": "int32"})
    foo.insert_many(docs)

    # Read non-empty datasource with a specified schema.
    schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        schema=schema,
        override_num_blocks=2,
    ).materialize()
    assert ds._block_num_rows() == [3, 2]
    assert ds.num_blocks() == 2
    assert ds.count() == 5
    ds_schema = ds.schema()
    assert ds_schema.names == ["float_field", "int_field"]
    assert ds_schema.types == [pa.float64(), pa.int32()]
    result = ds.to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Read with schema inference, which will read all columns (including the auto
    # generated internal column "_id").
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        override_num_blocks=2,
    ).materialize()
    assert ds._block_num_rows() == [3, 2]
    assert ds.num_blocks() == 2
    assert ds.count() == 5
    ds_schema = ds.schema()
    assert ds_schema.names == ["_id", "float_field", "int_field"]
    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
    result = ds.drop_columns(["_id"]).to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Read with auto-tuned parallelism.
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
    ).materialize()
    assert ds.num_blocks() == 2
    assert ds.count() == 5
    ds_schema = ds.schema()
    assert ds_schema.names == ["_id", "float_field", "int_field"]
    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
    result = ds.drop_columns(["_id"]).to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Read with a parallelism larger than number of rows.
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        override_num_blocks=1000,
    )
    assert ds.schema(fetch_if_missing=False) is None
    result = ds.drop_columns(["_id"]).to_pandas()
    pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)

    # Read a subset of the collection.
    ds = ray.data.read_mongo(
        uri=mongo_url,
        database=foo_db,
        collection=foo_collection,
        pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
        override_num_blocks=2,
    )
    assert ds._block_num_rows() == [2, 1]
    ds_schema = ds.schema()
    assert ds_schema.names == ["_id", "float_field", "int_field"]
    assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
    df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))