#include <sys/types.h>
#include <sys/stat.h>

#include <err.h>
#include <sha2.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "oqs/oqs.h"

static int
writefile(const char *base, const char *suff, uint8_t *contents, size_t len)
{
	size_t	 cc;
	size_t	 wc;
	FILE	*fp;
	char	 f[1024];
	int	 fd;

	snprintf(f, sizeof(f), "%s%s", base, suff);
	if ((fd = mkstemp(f)) < 0) {
		warn("can't open file '%s'", f);
		return 0;
	}
	if ((fp = fdopen(fd, "w+")) == NULL) {
		warn("can't open file '%s'", f);
		return 0;
	}
	for (cc = 0 ; cc < len ; cc += wc) {
		if ((wc = fwrite(&contents[cc], 1, len - cc, fp)) == 0) {
			break;
		}
	}
	fclose(fp);
	return (cc == len);
}

static int
readfile(const char *f, uint8_t **contents, size_t *len)
{
	struct stat	 st;
	size_t		 cc;
	size_t		 wc;
	FILE		*fp;

	if ((fp = fopen(f, "r")) == NULL) {
		warn("can't open file to read '%s'", f);
		return 0;
	}
	if ((*contents = calloc(1, *len = st.st_size + 1)) == NULL) {
		warn("can't allocate memory");
		fclose(fp);
		return 0;
	}
	/* agcXXX - sanity check sizes here */
	fprintf(stderr, "agcXXX - ADD SANITY CHECK ON SIZES HERE\n");
	for (cc = 0 ; cc < st.st_size ; cc += wc) {
		if ((wc = fwrite(&(*contents)[cc], 1, *len - cc, fp)) == 0) {
			break;
		}
	}
	fclose(fp);
	return (cc == *len);
}

/* agcXXX - mmap */
static int
sumfile(const char *f, uint8_t *sum, size_t sumsize)
{
	struct stat	 st;
	SHA512_CTX	 ctx;
	uint8_t		 buf[8192];
	size_t		 cc;
	size_t		 rc;
	FILE		*fp;

	if ((fp = fopen(f, "r")) == NULL) {
		warnx("can't read file '%s'", f);
		return 0;
	}
	fstat(fileno(fp), &st);
	SHA512_Init(&ctx);
	for (cc = 0 ; cc < st.st_size ; cc += rc) {
		if ((rc = fread(buf, 1, st.st_size - cc, fp)) == 0) {
			break;
		}
		SHA512_Update(&ctx, buf, rc);
	}
	fclose(fp);
	SHA512_Final(sum, &ctx);
	return 1;
}

/* generate public and secret keys */
static int
keygen(int argc, char **argv, char *key, int keyc)
{
	const char	*sig_alg;
	OQS_SIG		*sig;
	uint8_t		*public_key;
	uint8_t		*secret_key;
	char		 f[1024];
	int		 ok;

	/* ML-DSA-65 is the standard PQC signature algorithm (NIST FIPS 204) */
	sig_alg = OQS_SIG_alg_ml_dsa_65;
	/* Initialize the OQS SIG struct */
	sig = OQS_SIG_new(sig_alg);
	if (sig == NULL) {
		warn("Failed to initialize ML-DSA-65.");
		exit(EXIT_FAILURE);
	}
	/* Allocate buffers for keys */
	public_key = calloc(1, sig->length_public_key);
	secret_key = calloc(1, sig->length_secret_key);
	/* 1. Generate Signer Keypair (Done once by the software author) */
	if (OQS_SIG_keypair(sig, public_key, secret_key) != OQS_SUCCESS) {
		warn("Keypair generation failed.");
		goto cleanup;
	}
	printf("Generated ML-DSA keypair in '%s'\n", key);
	if (!writefile(key, ".pub", public_key, sig->length_public_key)) {
		warn("can't write public key to disk");
		goto cleanup;
	}
	/* agcXXX - need to encrypt secret key */
	fprintf(stderr, "agcXXX - ADD ENCRYPTED SECRET KEY HERE\n");
	if (!writefile(key, ".sec", secret_key, sig->length_secret_key)) {
		warn("can't write secret key to disk");
		goto cleanup;
	}
	ok = 1;
cleanup:
	free(public_key);
	free(secret_key);
	OQS_SIG_free(sig);
	return ok;
}

static int
signdata(int argc, char **argv, char *key, int keyc)
{
	const char	*sig_alg;
	OQS_SIG		*sig;
	uint8_t		*signature;
	size_t		signature_len;
	uint8_t		*seckey;
	uint8_t		 digest[SHA512_DIGEST_LENGTH];
	size_t		 seckeysize;
	char		 f[1024];
	int		 ok;
	int		 i;

	ok = 0;
	snprintf(f, sizeof(f), "%spub", key);
	/* agcXXX - unencrypt this */
	if (!readfile(f, &seckey, &seckeysize)) {
		warn("can't read secret key in '%s'", f);
		return 0;
	}
	/* ML-DSA-65 is the standard PQC signature algorithm (NIST FIPS 204) */
	sig_alg = OQS_SIG_alg_ml_dsa_65;
	/* Initialize the OQS SIG struct */
	sig = OQS_SIG_new(sig_alg);
	if (sig == NULL) {
		fprintf(stderr, "Failed to initialize ML-DSA-65.\n");
		return 1;
	}
	for (ok = 1, i = optind; i < argc ; i++) {
		if (!sumfile(argv[i], digest, sizeof(digest))) {
			warnx("can't get digest of file '%s'", argv[i]);
			ok = 0;
			continue;
		}
		/* Allocate space for the signature */
		signature = calloc(1, sig->length_signature);
		signature_len = 0;
		/* 3. Sign the software (Done by the developer before distributing) */
		if (OQS_SIG_sign(sig, signature, &signature_len, digest, sizeof(digest), seckey) != OQS_SUCCESS) {
			warnx("Signing '%s' failed.", argv[i]);
			ok = 0;
		}
		if (!writefile(argv[i], ".sig", signature, signature_len)) {
			warnx("Can't write signature file");
			ok = 0;
		}
		printf("'%s' signed successfully. Signature length: %zu bytes.\n", argv[i], signature_len);
	}
cleanup:
	free(seckey);
	free(signature);
	OQS_SIG_free(sig);
	return ok;
}

static int
verifydata(int argc, char **argv, char *key, int keyc)
{
	// ML-DSA-65 is the standard PQC signature algorithm (NIST FIPS 204)
	const char *sig_alg = OQS_SIG_alg_ml_dsa_65;

	// Initialize the OQS SIG struct
	OQS_SIG *sig = OQS_SIG_new(sig_alg);
	if (sig == NULL) {
		fprintf(stderr, "Failed to initialize ML-DSA-65.\n");
		return 1;
	}

	// Allocate buffers for keys
	uint8_t *public_key = malloc(sig->length_public_key);
	uint8_t *secret_key = malloc(sig->length_secret_key);

	// 1. Generate Signer Keypair (Done once by the software author)
	if (OQS_SIG_keypair(sig, public_key, secret_key) != OQS_SUCCESS) {
		fprintf(stderr, "Keypair generation failed.\n");
		goto cleanup;
	}
	printf("1. Generated ML-DSA keypair.\n");

	// 2. The software payload to sign 
	// In real software distribution, this would be a hash of your installer/binary
	const uint8_t *software_binary = (const uint8_t *)"my-installer-v1.0.exe payload data";
	size_t binary_len = strlen((const char *)software_binary);

	// Allocate space for the signature
	uint8_t *signature = malloc(sig->length_signature);
	size_t signature_len = 0;

	// 3. Sign the software (Done by the developer before distributing)
	if (OQS_SIG_sign(sig, signature, &signature_len, software_binary, binary_len, secret_key) != OQS_SUCCESS) {
		fprintf(stderr, "Signing failed.\n");
		goto cleanup;
	}
	printf("2. Software signed successfully. Signature length: %zu bytes.\n", signature_len);

	// --- At this point, you distribute public_key, software_binary, and signature ---

	// 4. Verify the software (Done by the end-user's device/installer)
	if (OQS_SIG_verify(sig, software_binary, binary_len, signature, signature_len, public_key) == OQS_SUCCESS) {
		printf("3. SUCCESS: The software signature is valid! Safe to install.\n");
	} else {
		printf("3. ERROR: Invalid signature! The software may have been tampered with.\n");
	}

cleanup:
	free(public_key);
	free(secret_key);
	free(signature);
	OQS_SIG_free(sig);
	return 0;
}

int
main(int argc, char **argv)
{
	const char	*sig_alg;
	OQS_SIG		*sig;
	uint8_t		*public_key;
	uint8_t		*secret_key;
	char		 key[1024];
	char		 f[1024];
	int		 ok;
	int		 keyc;
	int		 i;

	ok = 0;
	keyc = snprintf(key, sizeof(key), "%s", "ml-dsa-key");
	while ((i = getopt(argc, argv, "k:")) != -1) {
		switch(i) {
		case 'k':
			keyc = snprintf(key, sizeof(key), "%s", optarg);
			/* agcXXX - sanitise for path, / and .. */
			break;
		}
	}
	if (optind == argc) {
		fprintf(stderr, "Usage: %s [-k file] [keygen|sign|verify] file...\n", *argv);
		exit(EXIT_FAILURE);
	}
	if (strcmp(argv[optind], "keygen") == 0) {
		ok = keygen(argc, argv, key, keyc);
	} else if (strcmp(argv[optind], "sign") == 0) {
		ok = signdata(argc, argv, key, keyc);
	} else if (strcmp(argv[optind], "verify") == 0) {
		ok = signdata(argc, argv, key, keyc);
	}
	exit((ok) ? EXIT_SUCCESS : EXIT_FAILURE);
}

