#include <chrono>
#include <iostream>
#include <optional>
#include <stdint.h>
#include <string>
#include <vector>

#include <sodium.h>

// From https://github.com/tkislan/base64
#include "cpp-base64/base64.h"
#include "json.hpp"

#include "string_helpers.h"

using namespace std;
using chrono::duration_cast;
using chrono::seconds;
using chrono::system_clock;
using nlohmann::json;

using namespace BareMinimumCrypto;

const int64_t about_1_week = (int64_t)7 * 86400;
const int64_t about_3_months = (int64_t)105 * 86400;
const int64_t about_1_year = (int64_t)365 * 86400;

struct ExpiringSignature {
	string cert_s;
	vector <uint8_t> sig;
	
	// C++ nonsense
	bool operator == (const ExpiringSignature & o) const {
		return
			cert_s == o.cert_s &&
			sig == o.sig
		;
	}
	
	bool operator != (const ExpiringSignature & o) const {
		return ! (*this == o);
	}
};

int64_t get_seconds_since_epoch () {
	const auto utc_now = system_clock::now ();
	return duration_cast <seconds> (utc_now.time_since_epoch ()).count ();
}

void try_sodium_init () {
	if (sodium_init () < 0) {
		throw std::runtime_error ("Can't initialize libsodium");
	}
}

struct VerifiedData {
	vector <uint8_t> payload;
	string purpose;
};

optional <VerifiedData> try_verify_signed_data (
	const ExpiringSignature & sig,
	int64_t now,
	const vector <uint8_t> & pubkey
) {
	try_sodium_init ();
	
	if (pubkey.size () != crypto_sign_PUBLICKEYBYTES) {
		return nullopt;
	}
	
	if (crypto_sign_verify_detached (
		sig.sig.data (),
		(const uint8_t *)sig.cert_s.data (),
		sig.cert_s.size (),
		pubkey.data ()
	) != 0) {
		return nullopt;
	}
	
	const json j = json::parse (sig.cert_s);
	
	const int64_t not_before = j ["not_before"];
	const int64_t not_after = j ["not_after"];
	
	if (now < not_before) {
		return nullopt;
	}
	if (now > not_after) {
		return nullopt;
	}
	
	const string purpose = j ["purpose"];
	const string payload_b64 = j ["payload_b64"];
	const auto payload = std::move (*b64_decode (payload_b64));
	
	return VerifiedData {
		payload,
		purpose
	};
}

optional <VerifiedData> verify_signed_data (
	const ExpiringSignature & sig,
	int64_t now,
	const vector <uint8_t> & pubkey
) {
	try {
		return try_verify_signed_data (sig, now, pubkey);
	}
	catch (json::exception &) {
		return nullopt;
	}
}

class SigningKey {
	vector <uint8_t> pk;
	vector <uint8_t> sk;
	
public:
	SigningKey () {
		try_sodium_init ();
		
		pk.resize (crypto_sign_PUBLICKEYBYTES);
		sk.resize (crypto_sign_SECRETKEYBYTES);
		
		crypto_sign_keypair (pk.data (), sk.data ());
	}
	
	vector <uint8_t> pubkey () const {
		return pk;
	}
	
	string pub_to_base64 () const {
		return base64_encode (pk);
	}
	
	optional <ExpiringSignature> sign_base64 (
		const string & payload_b64, 
		string purpose,
		int64_t now,
		int64_t duration
	) const {
		try_sodium_init ();
		
		if (duration > about_1_year) {
			return nullopt;
		}
		
		const auto not_after = now + duration;
		
		const json j {
			{"not_before", now},
			{"not_after", not_after},
			{"purpose", purpose},
			{"payload_b64", payload_b64},
		};
		
		const auto cert_s = j.dump ();
		
		vector <uint8_t> sig;
		sig.resize (crypto_sign_BYTES);
		
		crypto_sign_detached (sig.data (), nullptr, (const uint8_t *)cert_s.data (), cert_s.size (), sk.data ());
		
		return ExpiringSignature {
			cert_s,
			sig,
		};
	}
	
	optional <ExpiringSignature> sign_key (const SigningKey & k, int64_t now) const 
	{
		return sign_base64 (k.pub_to_base64 (), "4QHAB7O5 trusted public key", now, about_3_months);
	}
	
	optional <ExpiringSignature> sign_data (const vector <uint8_t> & v, int64_t now) const 
	{
		return sign_base64 (base64_encode (v), "MS7WL26L signed data", now, about_1_week);
	}
};

// Most tests will use a virtual clock. But just as a smoke test,
// make sure real time is realistic.

int check_real_time () {
	const auto seconds_since_epoch = get_seconds_since_epoch ();
	
	const auto time_of_writing = 1610844872;
	if (seconds_since_epoch < time_of_writing) {
		cerr << "Error: Real time is in the past." << endl;
		return 1;
	}
	
	const int64_t about_100_years = (int64_t)100 * 365 * 86400;
	if (seconds_since_epoch > time_of_writing + about_100_years) {
		cerr << "Error: Real time is in the far future." << endl;
		return 1;
	}
	
	return 0;
}

int check_base64 () {
	// I assume that char is 8 bits
	// char is C++ nonsense inherited from C nonsense
	if (sizeof (char) != sizeof (uint8_t)) {
		cerr << "char is not the same size as uint8_t" << endl;
		return 1;
	}
	
	vector <uint8_t> v {1, 2, 3, 4, 5, 6};
	const auto s = base64_encode (v);
	
	if (s != "AQIDBAUG") {
		cerr << "Base64 encoding failed" << endl;
		return 1;
	}
	
	// Trivial decode
	const auto v2 = std::move (*b64_decode (s));
	
	if (v2 != v) {
		cerr << "Base64 trivial decode failed" << endl;
		return 1;
	}
	
	// Decode should fail
	const auto v3 = b64_decode ("AQIDBAUG.");
	
	if (v3 != nullopt) {
		cerr << "Base64 decode should have failed" << endl;
		return 1;
	}
	
	return 0;
}

int happy_path () {
	// We generate a root key and keep it somewhere safe
	// (offline, hopefully)
	SigningKey root_key;
	cerr << "Root pub key " << root_key.pub_to_base64 () << endl;
	
	if (check_real_time () != 0) {
		return 1;
	}
	
	// The server generates a signing key
	SigningKey signing_key;
	cerr << "Signing key " << signing_key.pub_to_base64 () << endl;
	
	const auto now = get_seconds_since_epoch ();
	
	// That signing key signs some important data
	const auto important_data = copy_to_bytes ("Nikolai, Anna, Ivan, Mikhail, Ivan, Nikolai, Anna. 7 4 1 4 3 5 7 4");
	const ExpiringSignature signed_data = std::move (*signing_key.sign_data (important_data, now));
	
	// The server signs our temporary signing key
	const ExpiringSignature cert = std::move (*root_key.sign_key (signing_key, now));
	
	{
		// Check that a different time results in a different cert
		const auto cert_2 = std::move (*root_key.sign_key (signing_key, now));
		const auto cert_3 = std::move (*root_key.sign_key (signing_key, now + 1));
		
		if (cert != cert_2) {
			cerr << "Certs should have been identical" << endl;
			return 1;
		}
		
		if (cert == cert_3) {
			cerr << "Certs should have been different" << endl;
			return 1;
		}
		
		if (cert == cert_3) {
			cerr << "Signatures should have been different" << endl;
			return 1;
		}
	}
	
	{
		cerr << "Cert:" << endl;
		cerr << cert.cert_s << endl;
		cerr << base64_encode (cert.sig) << endl;
	}
	
	// The client knows our root public key
	const auto root_pubkey = root_key.pubkey ();
	
	{
		auto verified_opt = verify_signed_data (cert, now, root_pubkey);
		if (! verified_opt) {
			cerr << "Couldn't verify cert" << endl;
			return 1;
		}
		const auto verified = std::move (*verified_opt);
		
		if (verified.purpose != "4QHAB7O5 trusted public key") {
			cerr << "Purpose did not match" << endl;
			return 1;
		}
		if (verified.payload != signing_key.pubkey ()) {
			cerr << "Pubkey payload did not match" << endl;
			return 1;
		}
	}
	
	{
		auto verified_opt = verify_signed_data (signed_data, now, signing_key.pubkey ());
		if (! verified_opt) {
			cerr << "Couldn't verify cert" << endl;
			return 1;
		}
		const auto verified = std::move (*verified_opt);
		
		if (verified.purpose != "MS7WL26L signed data") {
			cerr << "Purpose did not match" << endl;
			return 1;
		}
		if (verified.payload != important_data) {
			cerr << "Pubkey payload did not match" << endl;
			return 1;
		}
	}
	
	return 0;
}

int main () {
	if (check_base64 () != 0) {
		return 1;
	}
	
	if (happy_path () != 0) {
		return 1;
	}
	
	cerr << "All good." << endl;
	return 0;
}