#include <stdio.h>
#include <string.h>
#include <time.h>
#include <stdlib.h>
#include <stdbool.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <limits.h>

#include "djb2.h"
#include "entropy.h"
#include "md5.h"
#include "md5/global.h"
#include "md5/md5.h"
#include "sha256.h"
#include "hash_ledger.h"


static void hexify(const unsigned char *in, size_t len, char *out) {
	for (size_t i = 0; i < len; i++) {
		sprintf(out + (i * 2), "%02x", in[i]);
	}

	out[len * 2] = '\0';
}

size_t hash_ledger_bucket(struct hash_ledger *ledger, const char *path) {
	return djb2(path) % ledger->num_buckets;
}

struct hash_ledger *hash_ledger_init(size_t num_buckets) {
	struct hash_ledger *ledger = calloc(1, sizeof(struct hash_ledger));
	if (!ledger) {
		return NULL;
	}

	ledger->num_buckets = num_buckets;
	ledger->buckets = calloc(num_buckets, sizeof(struct hash_entry *));
	if (!ledger->buckets) {
		free(ledger);
		return NULL;
	}

	return ledger;
}

void hash_ledger_destroy(struct hash_ledger *ledger) {
	if (!ledger) {
		return;
	}

	for (size_t i = 0; i < ledger->num_buckets; i++) {
		struct hash_entry *entry = ledger->buckets[i];
		while (entry) {
			struct hash_entry *next = entry->next;
			free(entry);
			entry = next;
		}
	}

	free(ledger->buckets);
	free(ledger);
}

struct hash_entry *hash_ledger_find(struct hash_ledger *ledger, const char *path) {
	if (!ledger || !path) {
		return NULL;
	}

	size_t idx = hash_ledger_bucket(ledger, path);
	struct hash_entry *entry = ledger->buckets[idx];

	while (entry) {
		if (strcmp(entry->path, path) == 0) {
			return entry;
		}

		entry = entry->next;
	}

	return NULL;
}

struct hash_entry *hash_ledger_add_or_update(struct hash_ledger *ledger,
											 const char *path,
											 struct stat *sb) {
	if (!ledger || !path || !sb) {
		return NULL;
	}

	struct hash_entry *entry = hash_ledger_find(ledger, path);

	if (entry) {
		if (entry->sb.st_ino != sb->st_ino ||
			entry->sb.st_mtime != sb->st_mtime ||
			entry->sb.st_size != sb->st_size) {
			// file changed. must rehash
			struct multihash hashes;
			if (multihash_file(path, &hashes)) {
				entry->entropy = hashes.entropy;
				memcpy(entry->md5, hashes.md5, sizeof(entry->md5));
				memcpy(entry->sha256, hashes.sha256, sizeof(entry->sha256));
			}
			memcpy(&entry->sb, sb, sizeof(struct stat));
			entry->last_scanned = 0; // force rescan
			entry->scan_count = 0;
			entry->verdict = VERDICT_UNKNOWN;
			//entry->last_scanned = time(NULL);
		}

		// hash found. increment scan count
		//entry->scan_count++;
		return entry;
	}

	// not found. create new entry
	entry = calloc(1, sizeof(struct hash_entry));
	if (!entry) {
		return NULL;
	}

	strncpy(entry->path, path, sizeof(entry->path) - 1);
	memcpy(&entry->sb, sb, sizeof(struct stat));
	//entry->last_scanned = time(NULL);
	//entry->scan_count = 1;
	entry->last_scanned = 0;
	entry->scan_count = 0;
	entry->verdict = VERDICT_UNKNOWN;
	memset(entry->matched_rules, 0, sizeof(entry->matched_rules));

	struct multihash hashes;
	if (multihash_file(path, &hashes)) {
		entry->entropy = hashes.entropy;
		memcpy(entry->md5, hashes.md5, sizeof(entry->md5));
		memcpy(entry->sha256, hashes.sha256, sizeof(entry->sha256));
	}

	size_t idx = hash_ledger_bucket(ledger, path);
	entry->next = ledger->buckets[idx];
	ledger->buckets[idx] = entry;

	return entry;
}

bool multihash_file(const char *path, struct multihash *out) {
	if (!path || !out) {
		return false;
	}

	FILE *fp = fopen(path, "rb");
	if (!fp) {
		return false;
	}

	unsigned char  buf[8192];
	size_t         n;
	MD5_CTX        md5_ctx;
	SHA256_CTX     sha256_ctx;
	entropy_ctx    ent_ctx;

	MD5Init(&md5_ctx);
	sha256_init(&sha256_ctx);
	entropy_init(&ent_ctx);

	while((n = fread(buf, 1, sizeof(buf), fp)) > 0) {
		MD5Update(&md5_ctx, buf, n);
		sha256_update(&sha256_ctx, buf, n);
		entropy_update(&ent_ctx, buf, n);
	}

	fclose(fp);

	out->entropy = entropy_final(&ent_ctx);

	unsigned char md5_result[MD5_DIGEST_LENGTH];
	MD5Final(md5_result, &md5_ctx);
	hexify(md5_result, sizeof(md5_result), out->md5);

	unsigned char sha256_result[SHA256_DIGEST_LENGTH];
	sha256_final(&sha256_ctx, sha256_result);
	hexify(sha256_result, sizeof(sha256_result), out->sha256);

	return true;
}