// mtnet::SecureCipher round trip — exercises the full Metin2 KX handshake +
// stream cipher + AEAD session-token path offline (no server needed).
#include "../src/net/secure_cipher.h"
#include "../src/net/wire.h"

#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

using mtnet::SecureCipher;

static int g_fail = 0;
#define CHECK(c, msg)                                    \
	do {                                                 \
		if (!(c)) {                                       \
			std::fprintf(stderr, "FAIL: %s\n", msg);      \
			++g_fail;                                     \
		}                                                 \
	} while (0)

int main() {
	CHECK(SecureCipher::ensure_sodium_init(), "sodium_init");

	// --- key exchange: server sends KEY_CHALLENGE, client answers KEY_RESPONSE ---
	SecureCipher server, client;
	CHECK(server.initialize(), "server initialize");
	CHECK(client.initialize(), "client initialize");

	uint8_t server_pk[SecureCipher::PK_SIZE];
	server.get_public_key(server_pk);

	uint8_t challenge[SecureCipher::CHALLENGE_SIZE];
	randombytes_buf(challenge, sizeof(challenge));

	CHECK(client.compute_client_keys(server_pk), "client compute keys");

	uint8_t client_pk[SecureCipher::PK_SIZE];
	client.get_public_key(client_pk);
	CHECK(server.compute_server_keys(client_pk), "server compute keys");

	uint8_t response[crypto_auth_BYTES];
	client.compute_challenge_response(challenge, response);
	CHECK(server.verify_challenge_response(challenge, response), "server verifies challenge response");

	// tamper -> must fail
	uint8_t bad = response[0] ^ 0xFF;
	uint8_t bad_resp[crypto_auth_BYTES];
	std::memcpy(bad_resp, response, sizeof(response));
	bad_resp[0] = bad;
	CHECK(!server.verify_challenge_response(challenge, bad_resp), "tampered response rejected");

	// --- KEY_COMPLETE: server encrypts a session token, client decrypts it ---
	uint8_t token[SecureCipher::SESSION_TOKEN_SIZE];
	randombytes_buf(token, sizeof(token));
	uint8_t enc[SecureCipher::SESSION_TOKEN_SIZE + SecureCipher::TAG_SIZE];
	uint8_t nonce[SecureCipher::NONCE_SIZE];
	CHECK(server.encrypt_token(token, sizeof(token), enc, nonce), "server encrypt token");

	uint8_t dec[SecureCipher::SESSION_TOKEN_SIZE];
	CHECK(client.decrypt_token(enc, sizeof(enc), nonce, dec), "client decrypt token");
	CHECK(std::memcmp(token, dec, sizeof(token)) == 0, "session token round trips");

	server.set_activated(true);
	client.set_activated(true);

	// --- stream cipher: C->S traffic in arbitrary chunks, order-sensitive ---
	const std::vector<size_t> chunk_sizes = {1, 4, 60, 3, 5, 100, 7, 64, 200, 13};
	std::string acc_plain, acc_recovered;
	for (size_t n : chunk_sizes) {
		std::vector<uint8_t> buf(n);
		for (size_t i = 0; i < n; ++i) {
			buf[i] = static_cast<uint8_t>((acc_plain.size() + i) * 7 + 1);
		}
		acc_plain.append(reinterpret_cast<char *>(buf.data()), n);
		client.encrypt_in_place(buf.data(), n); // C->S encrypt
		// on the wire it looks like ciphertext; server decrypts the same bytes
		server.decrypt_in_place(buf.data(), n);
		acc_recovered.append(reinterpret_cast<char *>(buf.data()), n);
	}
	CHECK(acc_plain == acc_recovered, "C->S stream recovers across chunk boundaries");
	CHECK(client.tx_nonce() == acc_plain.size(), "client tx byte counter advanced");
	CHECK(server.rx_nonce() == acc_plain.size(), "server rx byte counter advanced");

	// --- S->C direction is independent ---
	{
		std::vector<uint8_t> buf(150);
		for (size_t i = 0; i < buf.size(); ++i) {
			buf[i] = static_cast<uint8_t>(i ^ 0x5A);
		}
		std::vector<uint8_t> orig = buf;
		server.encrypt_in_place(buf.data(), buf.size()); // S->C
		CHECK(buf != orig, "S->C ciphertext differs from plaintext");
		client.decrypt_in_place(buf.data(), buf.size());
		CHECK(buf == orig, "S->C stream round trips");
	}

	// --- wire.h struct sizes (must match the fork's #pragma pack(1) layout) ---
	CHECK(sizeof(mtnet::DynHeader) == 4, "DynHeader is 4 bytes");
	CHECK(sizeof(mtnet::GCKeyChallenge) == 4 + 32 + 32 + 4, "GCKeyChallenge 72 bytes");
	CHECK(sizeof(mtnet::CGKeyResponse) == 4 + 32 + 32, "CGKeyResponse 68 bytes");
	CHECK(sizeof(mtnet::GCKeyComplete) == 4 + 48 + 24, "GCKeyComplete 76 bytes");
	CHECK(sizeof(mtnet::CGLogin3) == 4 + 31 + 17, "CGLogin3 52 bytes");
	CHECK(sizeof(mtnet::GCAuthSuccess) == 4 + 4 + 1, "GCAuthSuccess 9 bytes");

	if (g_fail) {
		std::fprintf(stderr, "%d check(s) failed\n", g_fail);
		return 1;
	}
	std::printf("all checks passed\n");
	return 0;
}