// #include <boost/program_options.hpp>
// namespace po = boost::program_options;

#include <iostream>
#include <algorithm>
#include <random>
#include <chrono>

#include "BGN.hpp"
#include "client.hpp"
#include "server.hpp"
#include "serverEntity.hpp"

using namespace std;

const int MAX_ALLOWED_VOTE = 2;

// int argparse(
//     int argc,
//     char *argv[],
//     size_t &numServers,
//     size_t &numClients,
//     size_t &numRounds,
//     size_t &numVotes,
//     bool &maliciousServers,
//     bool &maliciousUsers,
//     string &seedStr) 
// {
//     string config_file;

//     // Declare a group of options that will be 
//     // allowed only on command line
//     po::options_description generic("General options");
//     generic.add_options()
//         ("help", "produce this help message")
//         ("config,c", po::value<string>(&config_file)->default_value(""),
//               "name of a configuration file")
//         ;

//     // Declare a group of options that will be 
//     // allowed both on command line and in
//     // config file
//     po::options_description config("Configuration");
//     config.add_options()
//         ("malicious-servers,M",
//                "presence of this flag indicates servers will operate in malicious model")
//         ("malicious-users,U",
//                "presence of this flag indicates users will operate in malicious model")
//         ("seed", po::value<string>(&seedStr)->default_value("default"),
//             "the random seed to use for this test")
//         ;

//     // Hidden options, will be allowed both on command line and
//     // in config file, but will not be shown to the user.
//     po::options_description hidden("Hidden options");
//     hidden.add_options()
//         ("number-servers,S", po::value<size_t>(&numServers)->default_value(2), 
//               "number of servers in test")
//         ("number-users,N", po::value<size_t>(&numClients)->default_value(5), 
//               "number of users in test")
//         ("number-rounds,R", po::value<size_t>(&numRounds)->default_value(3), 
//               "number of rounds to perform in test")
//         ("number-votes,V", po::value<size_t>(&numVotes)->default_value(3), 
//               "number of votes each user makes per round during test")
//         ;
    
//     po::options_description cmdline_options;
//     cmdline_options.add(generic).add(config).add(hidden);

//     po::options_description config_file_options;
//     config_file_options.add(config).add(hidden);

//     po::options_description visible("Allowed options");
//     visible.add(generic).add(config);
    
//     po::positional_options_description p;
//     p.add("number-servers", 1);
//     p.add("number-users", 1);
//     p.add("number-rounds", 1);
//     p.add("number-votes", 1);
    
//     po::variables_map vm;
//     store(po::command_line_parser(argc, argv).
//           options(cmdline_options).positional(p).run(), vm);
//     notify(vm);

//     if (!config_file.empty())
//     {
//         ifstream config(config_file.c_str());
//         if (!config)
//         {
//             cerr << "Cannot open config file: " << config_file << "\n";
//             return 2;
//         }
//         else
//         {
//             store(parse_config_file(config, config_file_options), vm);
//             notify(vm);
//         }
//     }

//     if (vm.count("help"))
//     {
//         cout << visible << endl;
//         return 1;
//     }

//     maliciousServers = vm.count("malicious-servers");
//     maliciousUsers = vm.count("malicious-users");

//     return 0;
// }

// Initialize the classes we use
void initialize_prsona_classes()
{
    Scalar::init();
    Curvepoint elGamalBlindGenerator = PrsonaServer::init();
    PrsonaClient::init(elGamalBlindGenerator);
}

// Do an epoch (including votes, etc.), and return the timing to print out
double epoch(
    default_random_engine& generator,
    PrsonaServerEntity& servers,
    vector<PrsonaClient>& users,
    size_t numVotes)
{
    Proof unused;
    uniform_int_distribution<int> voteDistribution(0, MAX_ALLOWED_VOTE);
    size_t numUsers = users.size();

    for (size_t i = 0; i < numUsers; i++)
    {
        // Make the correct number of new votes, but shuffle where they go
        vector<Scalar> votes;
        vector<bool> replace;
        for (size_t j = 0; j < numUsers; j++)
        {
            votes.push_back(Scalar(voteDistribution(generator)));
            replace.push_back(j < numVotes);
        }
        shuffle(replace.begin(), replace.end(), generator);

        // Make the actual votes to give to the servers
        Proof pi;
        Curvepoint shortTermPublicKey = users[i].get_short_term_public_key(pi);
        vector<CurveBipoint> encryptedVotes = users[i].make_votes(
                                                pi, votes, replace);

        // Give the servers these new votes
        servers.receive_vote(pi, encryptedVotes, shortTermPublicKey);
    }

    // Do the epoch server calculations
    chrono::high_resolution_clock::time_point t0 =
        chrono::high_resolution_clock::now();
    servers.epoch(unused);
    chrono::high_resolution_clock::time_point t1 =
        chrono::high_resolution_clock::now();
    
    // Transmit the results of the epoch to each user
    for (size_t i = 0; i < numUsers; i++)
        servers.transmit_updates(users[i]);

    // Return the timing of the epoch server calculations
    chrono::duration<double> time_span =
        chrono::duration_cast<chrono::duration<double>>(t1 - t0);
    return time_span.count();
}

int main(int argc, char *argv[])
{
    initialize_prsona_classes();

    // Defaults
    size_t numServers = 2;
    size_t numUsers = 5;
    size_t numRounds = 3;
    size_t numVotesPerRound = 3;
    bool maliciousServers = false;
    bool maliciousUsers = false;
    string seedStr = "seed";

    // Potentially accept command line inputs
    if (argc > 1)
        numServers = atoi(argv[1]);
    if (argc > 2)
        numUsers = atoi(argv[2]);
    if (argc > 3)
        numRounds = atoi(argv[3]);
    if (argc > 4)
        numVotesPerRound = atoi(argv[4]);

    cout << "Running the protocol with the following parameters: " << endl;
    cout << numServers << " PRSONA servers" << endl;
    cout << numUsers << " participants (voters/votees)" << endl;
    cout << numRounds << " epochs" << endl;
    cout << numVotesPerRound << " new (random) votes by each user per epoch"
        << endl;

    // Set malicious flags where necessary
    if (maliciousServers)
    {
        PrsonaServer::set_server_malicious();
        PrsonaClient::set_server_malicious();
    }
    if (maliciousUsers)
    {
        PrsonaServer::set_client_malicious();
        PrsonaClient::set_client_malicious();
    }

    // Entities we operate with
    PrsonaServerEntity servers(numServers);
    BGNPublicKey bgnPublicKey = servers.get_bgn_public_key();

    cout << "Initialization: adding users to system" << endl;
    vector<PrsonaClient> users;
    for (size_t i = 0; i < numUsers; i++)
    {
        PrsonaClient currUser(bgnPublicKey, &servers);
        servers.add_new_client(currUser);
        users.push_back(currUser);
    }

    // Seeded randomness for random votes used in epoch
    seed_seq seed(seedStr.begin(), seedStr.end());
    default_random_engine generator(seed);

    for (size_t i = 0; i < numRounds; i++)
    {
        cout << "Round " << i+1 << " of " << numRounds << ": " << endl;
        double timing = epoch(generator, servers, users, numVotesPerRound);
        cout << "Server computation: " << timing << " seconds" << endl;
    }

    return 0;
}