#include <iostream>
#include <stdlib.h>
#include <sys/random.h>
#include <sys/time.h>
#include <unistd.h>
#include "spir.hpp"

using std::cout;
using std::cerr;

static inline size_t elapsed_us(const struct timeval *start)
{
    struct timeval end;
    gettimeofday(&end, NULL);
    return (end.tv_sec-start->tv_sec)*1000000 + end.tv_usec - start->tv_usec;
}

int main(int argc, char **argv)
{
    if (argc < 2 || argc > 5) {
        cerr << "Usage: " << argv[0] << " r [num_threads [num_preproc [num_pirs]]]\n";
        cerr << "r = log_2(num_records)\n";
        exit(1);
    }
    uint32_t r, num_threads = 1, num_preproc = 1, num_pirs = 1;
    r = strtoul(argv[1], NULL, 10);
    size_t num_records = ((size_t) 1)<<r;
    size_t num_records_mask = num_records - 1;
    if (argc > 2) {
        num_threads = strtoul(argv[2], NULL, 10);
    }
    if (argc > 3) {
        num_preproc = strtoul(argv[3], NULL, 10);
    }
    if (argc > 4) {
        num_pirs = strtoul(argv[4], NULL, 10);
    } else {
        num_pirs = num_preproc;
    }

    cout << "===== ONE-TIME SETUP =====\n\n";

    struct timeval otsetup_start;
    gettimeofday(&otsetup_start, NULL);

    SPIR::init(num_threads);
    string pub_params;
    SPIR_Client client(r, pub_params);
    SPIR_Server server(r, pub_params);

    size_t otsetup_us = elapsed_us(&otsetup_start);
    cout << "One-time setup: " << otsetup_us << " µs\n";
    cout << "pub_params len = " << pub_params.length() << "\n";

    cout << "\n===== PREPROCESSING =====\n\n";

    cout << "num_preproc = " << num_preproc << "\n";

    struct timeval preproc_client_start;
    gettimeofday(&preproc_client_start, NULL);

    string preproc_msg = client.preproc(num_preproc);
    size_t preproc_client_us = elapsed_us(&preproc_client_start);
    cout << "Preprocessing client: " << preproc_client_us << " µs\n";
    cout << "preproc_msg len = " << preproc_msg.length() << "\n";

    struct timeval preproc_server_start;
    gettimeofday(&preproc_server_start, NULL);

    string preproc_resp = server.preproc_process(preproc_msg);
    size_t preproc_server_us = elapsed_us(&preproc_server_start);
    cout << "Preprocessing server: " << preproc_server_us << " µs\n";
    cout << "preproc_resp len = " << preproc_resp.length() << "\n";

    struct timeval preproc_finish_start;
    gettimeofday(&preproc_finish_start, NULL);

    client.preproc_finish(preproc_resp);
    size_t preproc_finish_us = elapsed_us(&preproc_finish_start);
    cout << "Preprocessing client finish: " << preproc_finish_us << " µs\n";

    // Create the database
    SPIR::DBEntry *db = new SPIR::DBEntry[num_records];
    for (size_t i=0; i<num_records; ++i) {
        db[i] = i * 10000001;
    }

    for (size_t i=0; i<num_pirs; ++i) {
        cout << "\n===== SPIR QUERY " << i+1 << " =====\n\n";

        size_t idx;
        if (getrandom(&idx, sizeof(idx), 0) != sizeof(idx)) {
            cerr << "Failure in getrandom\n";
            exit(1);
        }
        idx &= num_records_mask;

        struct timeval query_client_start;
        gettimeofday(&query_client_start, NULL);

        string query_msg = client.query(idx);
        size_t query_client_us = elapsed_us(&query_client_start);
        cout << "Query client: " << query_client_us << " µs\n";
        cout << "query_msg len = " << query_msg.length() << "\n";

        struct timeval query_server_start;
        gettimeofday(&query_server_start, NULL);

        string query_resp = server.query_process(query_msg, db, 100, 20);
        size_t query_server_us = elapsed_us(&query_server_start);
        cout << "Query server: " << query_server_us << " µs\n";
        cout << "query_resp len = " << query_resp.length() << "\n";

        struct timeval query_finish_start;
        gettimeofday(&query_finish_start, NULL);

        SPIR::DBEntry entry = client.query_finish(query_resp);
        size_t query_finish_us = elapsed_us(&query_finish_start);
        cout << "Query client finish: " << query_finish_us << " µs\n";
        cout << "idx = " << idx << "; entry = " << entry << "\n";
    }

    delete[] db;

    return 0;
}