#include <mongoc.h>
#include <mongoc-uri-private.h>

#include "TestSuite.h"
#include "mock_server/future.h"
#include "mock_server/future-functions.h"
#include "mock_server/mock-server.h"
#include "test-conveniences.h"


static bool
_can_be_command (const char *query) {
   return (!bson_empty (tmp_bson (query)));
}

static void
_test_op_query (const mongoc_uri_t *uri,
                mock_server_t *server,
                const char *query_in,
                mongoc_read_prefs_t *read_prefs,
                mongoc_query_flags_t expected_query_flags,
                const char          *expected_query)
{
   mongoc_client_t *client;
   mongoc_collection_t *collection;
   mongoc_cursor_t *cursor;
   const bson_t *doc;
   bson_t b = BSON_INITIALIZER;
   future_t *future;
   request_t *request;

   client = mongoc_client_new_from_uri (uri);
   collection = mongoc_client_get_collection (client, "test", "test");
   mongoc_collection_set_read_prefs (collection, read_prefs);

   cursor = mongoc_collection_find (collection,
                                    MONGOC_QUERY_NONE,
                                    0,
                                    1,
                                    0,
                                    tmp_bson (query_in),
                                    NULL,
                                    read_prefs);

   future = future_cursor_next (cursor, &doc);

   request = mock_server_receives_query (
      server,
      "test.test",
      expected_query_flags,
      0,
      1,
      expected_query,
      NULL);

   mock_server_replies (request,
                        MONGOC_REPLY_NONE,    /* flags */
                        0,                    /* cursorId */
                        0,                    /* startingFrom */
                        1,                    /* numberReturned */
                        "{'a': 1}");

   /* mongoc_cursor_next returned true */
   assert (future_get_bool (future));

   request_destroy (request);
   future_destroy (future);
   mongoc_cursor_destroy (cursor);
   mongoc_collection_destroy (collection);
   mongoc_client_destroy (client);
   bson_destroy (&b);
}


static void
_test_find_command (const mongoc_uri_t  *uri,
                    mock_server_t       *server,
                    const char          *query_in,
                    mongoc_read_prefs_t *read_prefs,
                    mongoc_query_flags_t expected_find_cmd_query_flags,
                    const char          *expected_find_cmd)
{
   mongoc_client_t *client;
   mongoc_collection_t *collection;
   mongoc_cursor_t *cursor;
   const bson_t *doc;
   bson_t b = BSON_INITIALIZER;
   future_t *future;
   request_t *request;

   client = mongoc_client_new_from_uri (uri);
   collection = mongoc_client_get_collection (client, "test", "test");
   mongoc_collection_set_read_prefs (collection, read_prefs);

   cursor = mongoc_collection_find (collection,
                                    MONGOC_QUERY_NONE,
                                    0,
                                    1,
                                    0,
                                    tmp_bson (query_in),
                                    NULL,
                                    read_prefs);

   future = future_cursor_next (cursor, &doc);

   request = mock_server_receives_command (
      server,
      "test",
      expected_find_cmd_query_flags,
      expected_find_cmd);

   mock_server_replies (request,
                        MONGOC_REPLY_NONE,    /* flags */
                        0,                    /* cursorId */
                        0,                    /* startingFrom */
                        1,                    /* numberReturned */
                        "{'ok': 1,"
                        " 'cursor': {"
                        "    'id': 0,"
                        "    'ns': 'db.collection',"
                        "    'firstBatch': [{'a': 1}]}}");

   /* mongoc_cursor_next returned true */
   assert (future_get_bool (future));

   request_destroy (request);
   future_destroy (future);
   mongoc_cursor_destroy (cursor);
   mongoc_collection_destroy (collection);
   mongoc_client_destroy (client);
   bson_destroy (&b);
}

static void
_test_command (const mongoc_uri_t *uri,
               mock_server_t *server,
               const char *command,
               mongoc_read_prefs_t *read_prefs,
               mongoc_query_flags_t expected_query_flags,
               const char *expected_query)
{
   mongoc_client_t *client;
   mongoc_collection_t *collection;
   mongoc_cursor_t *cursor;
   const bson_t *doc;
   bson_t b = BSON_INITIALIZER;
   future_t *future;
   request_t *request;

   client = mongoc_client_new_from_uri (uri);
   collection = mongoc_client_get_collection (client, "test", "test");
   mongoc_collection_set_read_prefs (collection, read_prefs);

   cursor = mongoc_collection_command (collection,
                                       MONGOC_QUERY_NONE,
                                       0,
                                       1,
                                       0,
                                       tmp_bson (command),
                                       NULL,
                                       read_prefs);

   future = future_cursor_next (cursor, &doc);

   request = mock_server_receives_command (
      server,
      "test",
      expected_query_flags,
      expected_query);

   mock_server_replies (request,
                        MONGOC_REPLY_NONE,    /* flags */
                        0,                    /* cursorId */
                        0,                    /* startingFrom */
                        1,                    /* numberReturned */
                        "{'a': 1}");

   /* mongoc_cursor_next returned true */
   assert (future_get_bool (future));

   request_destroy (request);
   future_destroy (future);
   mongoc_cursor_destroy (cursor);
   mongoc_collection_destroy (collection);
   mongoc_client_destroy (client);
   bson_destroy (&b);
}

static void
_test_command_simple (const mongoc_uri_t *uri,
                      mock_server_t *server,
                      const char *command,
                      mongoc_read_prefs_t *read_prefs,
                      mongoc_query_flags_t expected_query_flags,
                      const char *expected_query)
{
   mongoc_client_t *client;
   mongoc_collection_t *collection;
   bson_t b = BSON_INITIALIZER;
   future_t *future;
   request_t *request;

   client = mongoc_client_new_from_uri (uri);
   collection = mongoc_client_get_collection (client, "test", "test");
   mongoc_collection_set_read_prefs (collection, read_prefs);

   future = future_client_command_simple (client,
                                          "test",
                                          tmp_bson (command),
                                          read_prefs,
                                          NULL,
                                          NULL);

   request = mock_server_receives_command (
      server,
      "test",
      expected_query_flags,
      expected_query);

   mock_server_replies (request,
                        MONGOC_REPLY_NONE,    /* flags */
                        0,                    /* cursorId */
                        0,                    /* startingFrom */
                        1,                    /* numberReturned */
                        "{'ok': 1}");

   /* mongoc_cursor_next returned true */
   assert (future_get_bool (future));

   request_destroy (request);
   future_destroy (future);
   mongoc_collection_destroy (collection);
   mongoc_client_destroy (client);
   bson_destroy (&b);
}


typedef enum {
    READ_PREF_TEST_STANDALONE,
    READ_PREF_TEST_MONGOS,
    READ_PREF_TEST_PRIMARY,
    READ_PREF_TEST_SECONDARY,
} read_pref_test_type_t;


static mock_server_t *
_run_server (read_pref_test_type_t test_type,
             int32_t               max_wire_version)
{
   mock_server_t *server;

   server = mock_server_new ();
   mock_server_run (server);

   switch (test_type) {
   case READ_PREF_TEST_STANDALONE:
      mock_server_auto_ismaster (server,
                                 "{'ok': 1,"
                                 " 'maxWireVersion': %d,"
                                 " 'ismaster': true}",
                                 max_wire_version);
      break;
   case READ_PREF_TEST_MONGOS:
      mock_server_auto_ismaster (server,
                                 "{'ok': 1,"
                                 " 'maxWireVersion': %d,"
                                 " 'ismaster': true,"
                                 " 'msg': 'isdbgrid'}",
                                 max_wire_version);
      break;
   case READ_PREF_TEST_PRIMARY:
      mock_server_auto_ismaster (server,
                                 "{'ok': 1,"
                                 " 'maxWireVersion': %d,"
                                 " 'ismaster': true,"
                                 " 'setName': 'rs',"
                                 " 'hosts': ['%s']}",
                                 max_wire_version,
                                 mock_server_get_host_and_port (server));
      break;
   case READ_PREF_TEST_SECONDARY:
      mock_server_auto_ismaster (server,
                                 "{'ok': 1,"
                                 " 'maxWireVersion': %d,"
                                 " 'ismaster': false,"
                                 " 'secondary': true,"
                                 " 'setName': 'rs',"
                                 " 'hosts': ['%s']}",
                                 max_wire_version,
                                 mock_server_get_host_and_port (server));
      break;
   default:
      fprintf (stderr, "Invalid test_type: : %d\n", test_type);
      abort ();
   }

   return server;
}


static mongoc_uri_t *
_get_uri (mock_server_t        *server,
          read_pref_test_type_t test_type)
{
   mongoc_uri_t *uri;

   uri = mongoc_uri_copy (mock_server_get_uri (server));

   switch (test_type) {
   case READ_PREF_TEST_PRIMARY:
   case READ_PREF_TEST_SECONDARY:
      mongoc_uri_set_option_as_utf8 (uri, "replicaSet", "rs");
      break;
   case READ_PREF_TEST_STANDALONE:
   case READ_PREF_TEST_MONGOS:
   default:
      break;
   }

   return uri;
}


static void
_test_read_prefs (
   read_pref_test_type_t test_type,
   mongoc_read_prefs_t  *read_prefs,
   const char           *query_in,
   const char           *expected_query,
   mongoc_query_flags_t  expected_query_flags,
   const char           *expected_find_cmd,
   mongoc_query_flags_t  expected_find_cmd_query_flags)
{
   mock_server_t *server;
   mongoc_uri_t *uri;

   server = _run_server (test_type, 3);
   uri = _get_uri (server, test_type);

   _test_op_query (uri,
                   server,
                   query_in,
                   read_prefs,
                   expected_query_flags,
                   expected_query);

   if (_can_be_command (query_in)) {
      _test_command (uri,
                     server,
                     query_in,
                     read_prefs,
                     expected_query_flags,
                     expected_query);

      _test_command_simple (uri,
                            server,
                            query_in,
                            read_prefs,
                            expected_query_flags,
                            expected_query);
   }

   mock_server_destroy (server);
   mongoc_uri_destroy (uri);

   server = _run_server (test_type, 4);
   uri = _get_uri (server, test_type);

   _test_find_command (uri,
                       server,
                       query_in,
                       read_prefs,
                       expected_find_cmd_query_flags,
                       expected_find_cmd);

   mock_server_destroy (server);
   mongoc_uri_destroy (uri);
}


/* test that a NULL read pref is the same as PRIMARY */
static void
test_read_prefs_standalone_null (void)
{
   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, NULL, "{}",
      "{}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, NULL, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_SLAVE_OK);
}

static void
test_read_prefs_standalone_primary (void)
{
   mongoc_read_prefs_t *read_prefs;

   /* Server Selection Spec: for topology type single and server types other
    * than mongos, "clients MUST always set the slaveOK wire protocol flag on
    * reads to ensure that any server type can handle the request."
    * */
   read_prefs = mongoc_read_prefs_new (MONGOC_READ_PRIMARY);

   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, read_prefs, "{}",
      "{}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, read_prefs, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_SLAVE_OK);

   mongoc_read_prefs_destroy (read_prefs);
}


static void
test_read_prefs_standalone_secondary (void)
{
   mongoc_read_prefs_t *read_prefs;

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_SECONDARY);

   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, read_prefs, "{}",
      "{}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, read_prefs, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_SLAVE_OK);

   mongoc_read_prefs_destroy (read_prefs);
}


static void
test_read_prefs_standalone_tags (void)
{
   bson_t b = BSON_INITIALIZER;
   mongoc_read_prefs_t *read_prefs;

   bson_append_utf8 (&b, "dc", 2, "ny", 2);

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_SECONDARY_PREFERRED);
   mongoc_read_prefs_add_tag (read_prefs, &b);
   mongoc_read_prefs_add_tag (read_prefs, NULL);

   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, read_prefs, "{}",
      "{}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_STANDALONE, read_prefs, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_SLAVE_OK);

   mongoc_read_prefs_destroy (read_prefs);
}


static void
test_read_prefs_primary_rsprimary (void)
{
   mongoc_read_prefs_t *read_prefs;

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_PRIMARY);

   _test_read_prefs (
      READ_PREF_TEST_PRIMARY, read_prefs, "{}",
      "{}", MONGOC_QUERY_NONE,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_NONE);

   _test_read_prefs (
      READ_PREF_TEST_PRIMARY, read_prefs, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_NONE,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_NONE);

   mongoc_read_prefs_destroy (read_prefs);
}


static void
test_read_prefs_secondary_rssecondary (void)
{
   mongoc_read_prefs_t *read_prefs;

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_SECONDARY);

   _test_read_prefs (
      READ_PREF_TEST_SECONDARY, read_prefs, "{}",
      "{}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_SECONDARY, read_prefs, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_SLAVE_OK);

   mongoc_read_prefs_destroy (read_prefs);
}



/* test that a NULL read pref is the same as PRIMARY */
static void
test_read_prefs_mongos_null (void)
{
   _test_read_prefs (
      READ_PREF_TEST_MONGOS, NULL, "{}",
      "{}", MONGOC_QUERY_NONE,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_NONE);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, NULL, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_NONE,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_NONE);
}


static void
test_read_prefs_mongos_primary (void)
{
   mongoc_read_prefs_t *read_prefs;

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_PRIMARY);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{}",
      "{}", MONGOC_QUERY_NONE,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_NONE);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_NONE,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_NONE);

   mongoc_read_prefs_destroy (read_prefs);
}


static void
test_read_prefs_mongos_secondary (void)
{
   mongoc_read_prefs_t *read_prefs;

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_SECONDARY);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{}",
      "{'$query': {}, '$readPreference': {'mode': 'secondary'}}",
      MONGOC_QUERY_SLAVE_OK,
      "{'$query': {'find': 'test', 'filter':  {}},"
      " '$readPreference': {'mode': 'secondary'}}",
      MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{'a': 1}",
      "{'$query': {'a': 1}, '$readPreference': {'mode': 'secondary'}}",
      MONGOC_QUERY_SLAVE_OK,
      "{'$query': {'find': 'test', 'filter':  {'a': 1}},"
      " '$readPreference': {'mode': 'secondary'}}",
      MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{'$query': {'a': 1}}",
      "{'$query': {'a': 1}, '$readPreference': {'mode': 'secondary'}}",
      MONGOC_QUERY_SLAVE_OK,
      "{'$query': {'find': 'test', 'filter':  {'a': 1}},"
      " '$readPreference': {'mode': 'secondary'}}",
      MONGOC_QUERY_SLAVE_OK);

   mongoc_read_prefs_destroy (read_prefs);
}


static void
test_read_prefs_mongos_secondary_preferred (void)
{
   mongoc_read_prefs_t *read_prefs;

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_SECONDARY_PREFERRED);

   /* $readPreference not sent, only slaveOk */
   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{}",
      "{}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {}}", MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{'a': 1}",
      "{'a': 1}", MONGOC_QUERY_SLAVE_OK,
      "{'find': 'test', 'filter':  {'a': 1}}", MONGOC_QUERY_SLAVE_OK);

   mongoc_read_prefs_destroy (read_prefs);
}


static void
test_read_prefs_mongos_tags (void)
{
   bson_t b = BSON_INITIALIZER;
   mongoc_read_prefs_t *read_prefs;

   bson_append_utf8 (&b, "dc", 2, "ny", 2);

   read_prefs = mongoc_read_prefs_new (MONGOC_READ_SECONDARY_PREFERRED);
   mongoc_read_prefs_add_tag (read_prefs, &b);
   mongoc_read_prefs_add_tag (read_prefs, NULL);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{}",
      "{'$query': {}, '$readPreference': {'mode': 'secondaryPreferred',"
      "                                   'tags': [{'dc': 'ny'}, {}]}}",
      MONGOC_QUERY_SLAVE_OK,
      "{'$query': {'find': 'test', 'filter':  {}},"
      " '$readPreference': {'mode': 'secondaryPreferred',"
      "                             'tags': [{'dc': 'ny'}, {}]}}",
      MONGOC_QUERY_SLAVE_OK);

   _test_read_prefs (
      READ_PREF_TEST_MONGOS, read_prefs, "{'a': 1}",
      "{'$query': {'a': 1},"
      " '$readPreference': {'mode': 'secondaryPreferred',"
      "                     'tags': [{'dc': 'ny'}, {}]}}",
      MONGOC_QUERY_SLAVE_OK,
      "{'$query': {'find': 'test', 'filter':  {}},"
      " '$readPreference': {'mode': 'secondaryPreferred',"
      "                             'tags': [{'dc': 'ny'}, {}]}}",
      MONGOC_QUERY_SLAVE_OK);

   mongoc_read_prefs_destroy (read_prefs);
}


void
test_read_prefs_install (TestSuite *suite)
{
   TestSuite_Add (suite, "/ReadPrefs/standalone/null",
                  test_read_prefs_standalone_null);
   TestSuite_Add (suite, "/ReadPrefs/standalone/primary",
                  test_read_prefs_standalone_primary);
   TestSuite_Add (suite, "/ReadPrefs/standalone/secondary",
                  test_read_prefs_standalone_secondary);
   TestSuite_Add (suite, "/ReadPrefs/standalone/tags",
                  test_read_prefs_standalone_tags);
   TestSuite_Add (suite, "/ReadPrefs/rsprimary/primary",
                  test_read_prefs_primary_rsprimary);
   TestSuite_Add (suite, "/ReadPrefs/rssecondary/secondary",
                  test_read_prefs_secondary_rssecondary);
   TestSuite_Add (suite, "/ReadPrefs/mongos/null",
                  test_read_prefs_mongos_null);
   TestSuite_Add (suite, "/ReadPrefs/mongos/primary",
                  test_read_prefs_mongos_primary);
   TestSuite_Add (suite, "/ReadPrefs/mongos/secondary",
                  test_read_prefs_mongos_secondary);
   TestSuite_Add (suite, "/ReadPrefs/mongos/secondaryPreferred",
                  test_read_prefs_mongos_secondary_preferred);
   TestSuite_Add (suite, "/ReadPrefs/mongos/tags",
                  test_read_prefs_mongos_tags);
}
