#include <atomic>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <random>
#include <thread>
#include <vector>

#include "dpf/launch.hpp"
#include "dpf/net/client_link.hpp"
#include "dpf/run_log.hpp"

// A client splits a secret into two additive shares and sends share i to party
// i over a client link; the two parties then open the sum over their own
// party link. Both links are TLS 1.3; each end logs how it authenticated the
// other.
//
//   c++ -std=c++17 -march=native -pthread -I include -I thirdparty \
//       examples/protocol/client_shares.cpp -lsctp -lssl -lcrypto -o client_shares
//   ./client_shares                          # development certificate (logged)
//   ./dpf_keygen srv.key                     # prints srv's public key
//   ./client_shares --server_identity=srv.key --client_pin=<srv public key>
//   ./client_shares --client_verify=off      # accepts any server (logged)
//
// With --server_identity and no --client_pin the client refuses the server:
// clients always verify unless told not to.

int main(int argc, char ** argv)
{
    try
    {
        auto cfg = dpf::app::run_config::from_env();
        for (const auto & extra : cfg.apply_args(argc, argv))
            throw std::invalid_argument("unknown argument " + extra);
        if (cfg.kind != dpf::net::transport::mux && cfg.kind != dpf::net::transport::parallel)
            cfg.kind = dpf::net::transport::mux;
        dpf::app::start_logging(cfg);

        // Each party accepts one client and keeps the share it sends.
        std::atomic<unsigned short> ports[2] = {{0}, {0}};
        std::uint64_t shares[2] = {0, 0};
        std::string errors[2];
        std::vector<std::thread> parties;
        for (int p = 0; p < 2; ++p)
            parties.emplace_back([&, p] {
                try
                {
                    asio::io_context io;
                    dpf::net::client_listener l(io, cfg.server, 1, cfg.policy, cfg.limits);
                    ports[p].store(l.listen());
                    auto c = l.accept();
                    bool done = false;
                    std::error_code ec;
                    c.link->async_read(0, &shares[p], 8, [&](const std::error_code & e) {
                        ec = e;
                        done = true;
                    });
                    while (!done)
                    {
                        if (io.stopped())
                            io.restart();
                        io.run_one();
                    }
                    if (ec)
                        throw std::system_error(ec, "reading the client's share");
                }
                catch (const std::exception & e)
                {
                    errors[p] = e.what();
                    ports[p].store(1);
                }
            });

        // The client: one fresh share per party, each on its own verified link.
        const std::uint64_t secret = 42;
        std::random_device rd;
        const std::uint64_t r = (static_cast<std::uint64_t>(rd()) << 32) | rd();
        const std::uint64_t mine[2] = {r, secret - r};
        std::string client_error;
        for (int p = 0; p < 2 && client_error.empty(); ++p)
        {
            while (ports[p].load() == 0)
                std::this_thread::yield();
            try
            {
                asio::io_context io;
                auto c = dpf::net::connect_server(io, cfg.host, ports[p].load(), cfg.client,
                    1, cfg.policy, cfg.limits);
                std::cout << "client -> party " << p << ": " << c.security.protocol << " "
                          << c.security.cipher << ", server auth=" << c.security.peer_auth
                          << "\n";
                bool done = false;
                c.link->async_write(0, &mine[p], 8, [&](const std::error_code &) {
                    done = true;
                });
                while (!done)
                {
                    if (io.stopped())
                        io.restart();
                    io.run_one();
                }
                c.link.reset();
                io.poll();
            }
            catch (const std::exception & e)
            {
                client_error = e.what();
            }
        }
        if (!client_error.empty())
        {
            // Unblock the listeners that are still waiting for this client.
            for (int p = 0; p < 2; ++p)
            {
                std::error_code ec;
                asio::io_context io;
                asio::ip::tcp::socket s(io);
                if (ports[p].load() > 1)
                    s.connect({asio::ip::make_address("127.0.0.1"), ports[p].load()}, ec);
            }
        }
        for (auto & t : parties)
            t.join();
        if (!client_error.empty())
            throw std::runtime_error("client: " + client_error);
        for (const auto & e : errors)
            if (!e.empty())
                throw std::runtime_error("party: " + e);

        // The parties open the sum over their party link.
        dpf::protocol::composer c0(0), c1(1);
        auto x0 = c0.input(dpf::protocol::domain::a, 8);
        auto x1 = c1.input(dpf::protocol::domain::a, 8);
        auto o0 = c0.exchange(x0);
        (void)c1.exchange(x1);
        auto p0 = c0.schedule();
        auto p1 = c1.schedule();
        dpf::app::party_values v0(p0.nodes().size()), v1(p1.nodes().size());
        v0[x0.id].assign(8, 0);
        v1[x1.id].assign(8, 0);
        std::memcpy(v0[x0.id].data(), &shares[0], 8);
        std::memcpy(v1[x1.id].data(), &shares[1], 8);
        (void)dpf::run_two_party(p0, p1, v0, v1, {}, cfg);
        std::uint64_t open = 0;
        std::memcpy(&open, v0[o0.id].data(), 8);
        std::cout << "parties opened " << open << " (share 0 = " << shares[0] << ")\n";
        return open == secret ? 0 : 1;
    }
    catch (const std::exception & e)
    {
        std::cerr << "client_shares: " << e.what() << "\n";
        return 1;
    }
}