#include "utils.h"

#include <stdarg.h>
#include <string.h>
#include <errno.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <fcntl.h>
#include <ctype.h>
#include <sys/wait.h>
#include <openssl/evp.h>
#include <dirent.h>

char ERROR_BUFFER[ERROR_BUFFER_SIZE];

result_t success() {
	return EXIT_SUCCESS;
}

result_t failure(const char* format, ...) {
	va_list args;
	va_start(args, format);

	// Write the error message to the error buffer.
	vsnprintf(ERROR_BUFFER, ERROR_BUFFER_SIZE, format, args);

	va_end(args);

	return EXIT_FAILURE;
}

const char* last_error() {
	return ERROR_BUFFER;
}

result_t execute_file(int* _exit_code, char** _stdout_buffer, size_t* _stdout_size, char** _stderr_buffer, size_t* _stderr_size, const char* working_directory, const char* file, ...) {
	// Count the number of arguments.
	va_list args;
	va_start(args, file);

	int argc = 1; // The first argument is the file itself.
	while (va_arg(args, const char*) != NULL)
		argc++;

	va_end(args);

	// Allocate the arguments array.
	const char** argv = malloc((argc + 1) * sizeof(const char*)); // The last element is the NULL terminator.
	if (argv == NULL)
		return failure("Failed to allocate memory for the arguments (%s).", strerror(errno));

	// Fill the arguments array.
	va_start(args, file);

	argv[0] = file;
	for (int i = 1; i < argc; i++)
		argv[i] = va_arg(args, const char*);
	argv[argc] = NULL;

	va_end(args);

	// Create the pipes for the standard output and error streams.
	int stdout_pipe[2];
	if (pipe(stdout_pipe) < 0) {
		free(argv);
		return failure("Failed to create the standard output pipe (%s).", strerror(errno));
	}

	int stderr_pipe[2];
	if (pipe(stderr_pipe) < 0) {
		close(stdout_pipe[0]);
		close(stdout_pipe[1]);
		free(argv);
		return failure("Failed to create the standard error pipe (%s).", strerror(errno));
	}

	// Fork the process.
	pid_t pid = fork();

	if (pid < 0) {
		close(stdout_pipe[0]);
		close(stdout_pipe[1]);
		close(stderr_pipe[0]);
		close(stderr_pipe[1]);
		free(argv);
		return failure("Failed to fork the process (%s).", strerror(errno));
	}

	if (pid == 0) {
		// Redirect the standard output and error streams to the pipes.
		dup2(stdout_pipe[1], STDOUT_FILENO);
		dup2(stderr_pipe[1], STDERR_FILENO);

		// Close the pipe file descriptors.
		close(stdout_pipe[0]);
		close(stdout_pipe[1]);
		close(stderr_pipe[0]);
		close(stderr_pipe[1]);

		// Change the working directory if necessary.
		if (working_directory != NULL) {
			if (chdir(working_directory) < 0) {
				fprintf(stderr, "Failed to change the working directory to '%s' (%s).\n", working_directory, strerror(errno));
				exit(EXIT_FAILURE);
			}
		}

		// Execute the file.
		execv(file, (char* const*)argv);

		// If execv returns, it failed.
		fprintf(stderr, "Failed to execute the file '%s' (%s).\n", file, strerror(errno));
		exit(EXIT_FAILURE);
	}

	// Close the write end of the pipes.
	close(stdout_pipe[1]);
	close(stderr_pipe[1]);

	// Free the arguments array as it is no longer needed.
	free(argv);

	// Wait for the child process to exit.
	int status;
	if (waitpid(pid, &status, 0) < 0) {
		close(stdout_pipe[0]);
		close(stderr_pipe[0]);
		return failure("Failed to wait for the child process to exit (%s).", strerror(errno));
	}

	// Read the standard output into a buffer if necessary.
	if (_stdout_buffer != NULL) {
		result_t result = read_fd(_stdout_buffer, _stdout_size, stdout_pipe[0]);
		if (result != success()) {
			close(stdout_pipe[0]);
			close(stderr_pipe[0]);
			return result;
		}
	}
	close(stdout_pipe[0]);

	// Read the standard error into a buffer if necessary.
	if (_stderr_buffer != NULL) {
		result_t result = read_fd(_stderr_buffer, _stderr_size, stderr_pipe[0]);
		if (result != success()) {
			close(stderr_pipe[0]);
			free(*_stdout_buffer);
			return result;
		}
	}
	close(stderr_pipe[0]);

	// Output the exit code if necessary.
	if (_exit_code != NULL)
		*_exit_code = WEXITSTATUS(status);

	return success();
}

result_t read_fd(char** _buffer, size_t* _size, int fd) {
	size_t size = 0;
	size_t capacity = 4096;

	// Allocate the buffer.
	char* buffer = malloc(capacity);
	if (buffer == NULL)
		return failure("Failed to allocate memory for the buffer (%s).", strerror(errno));

	// Read the file descriptor into the buffer.
	ssize_t bytes_read;
	while ((bytes_read = read(fd, buffer + size, capacity - size)) > 0) {
		size += bytes_read;

		// Resize the buffer if necessary.
		if (size == capacity) {
			// Basic doubling strategy.
			capacity *= 2;
			char* new_buffer = realloc(buffer, capacity);
			if (new_buffer == NULL) {
				free(buffer);
				return failure("Failed to reallocate memory for the buffer (%s).", strerror(errno));
			}
			buffer = new_buffer;
		}
	}

	// Check for read errors.
	if (bytes_read < 0) {
		free(buffer);
		return failure("Failed to read from the file descriptor (%s).", strerror(errno));
	}

	// Resize the buffer to the actual size.
	char* new_buffer = realloc(buffer, size + 1);
	if (new_buffer == NULL) {
		free(buffer);
		return failure("Failed to reallocate memory for the buffer (%s).", strerror(errno));
	}
	buffer = new_buffer;

	// Null-terminate the buffer.
	buffer[size] = '\0';

	// Output the buffer and size.
	*_buffer = buffer;

	if (_size != NULL)
		*_size = size;

	return success();
}

result_t write_fd(int fd, const char* buffer, size_t size) {
	size_t bytes_written = 0;

	// Write the buffer to the file descriptor.
	while (bytes_written < size) {
		ssize_t written = write(fd, buffer + bytes_written, size - bytes_written);
		if (written < 0)
			return failure("Failed to write to the file descriptor (%s).", strerror(errno));
		bytes_written += written;
	}

	return success();
}

result_t read_file(char** _buffer, size_t* _size, const char* file) {
	// Open the file.
	int fd = open(file, O_RDONLY);
	if (fd < 0)
		return failure("Failed to open the file '%s' (%s).", file, strerror(errno));

	// Read the file descriptor into a buffer.
	result_t result = read_fd(_buffer, _size, fd);

	// Close the file.
	close(fd);

	return result;
}

result_t write_file(const char* file, const char* buffer, size_t size, mode_t mode) {
	// Open the file.
	int fd = open(file, O_WRONLY | O_CREAT | O_TRUNC, mode);
	if (fd < 0)
		return failure("Failed to open the file '%s' (%s).", file, strerror(errno));

	// Write the buffer to the file descriptor.
	result_t result = write_fd(fd, buffer, size);

	// Close the file.
	close(fd);

	return result;
}

result_t copy_fd(int source_fd, int destination_fd) {
	char buffer[4096];
	ssize_t bytes_read;

	// Read the source file descriptor into the buffer.
	while ((bytes_read = read(source_fd, buffer, sizeof(buffer))) > 0) {
		// Try to write the buffer to the destination file descriptor.
		ssize_t bytes_written = 0;
		while (bytes_written < bytes_read) {
			ssize_t written = write(destination_fd, buffer + bytes_written, bytes_read - bytes_written);
			if (written < 0)
				return failure("Failed to write to the file descriptor (%s).", strerror(errno));
			bytes_written += written;
		}
	}

	// Check for read errors.
	if (bytes_read < 0)
		return failure("Failed to read from the file descriptor (%s).", strerror(errno));

	return success();
}

result_t copy_file(const char* source, const char* destination, mode_t mode) {
	// Open the source file.
	int source_fd = open(source, O_RDONLY);
	if (source_fd < 0)
		return failure("Failed to open the source file '%s' (%s).", source, strerror(errno));

	// Open the destination file.
	int destination_fd = open(destination, O_WRONLY | O_CREAT | O_TRUNC, mode);
	if (destination_fd < 0) {
		close(source_fd);
		return failure("Failed to open the destination file '%s' (%s).", destination, strerror(errno));
	}

	// Copy the source file to the destination file.
	result_t result = copy_fd(source_fd, destination_fd);

	// Close the files.
	close(source_fd);
	close(destination_fd);

	return result;
}

char* trim(const char* string) {
	// Check for NULL.
	if (string == NULL)
		return NULL;

	// Skip leading whitespace.
	while (*string != '\0' && isspace(*string))
		string++;

	// Find the end of the string.
	const char* end = string;
	while (*end != '\0')
		end++;

	// Skip trailing whitespace.
	while (end > string && isspace(*(end - 1)))
		end--;

	// Allocate a new string and copy the trimmed string into it.
	size_t length = end - string;
	char* trimmed = malloc(length + 1);
	if (trimmed == NULL)
		return NULL;

	memcpy(trimmed, string, length);
	trimmed[length] = '\0';

	return trimmed;
}

char* file_name(const char* path) {
	// Check for NULL.
	if (path == NULL)
		return NULL;

	// Find the last occurrence of the path separator.
	const char* file_name = strrchr(path, '/');

	// If the path separator was found, return the file name after it.
	if (file_name != NULL)
		return strdup(file_name + 1);

	// Otherwise, return the entire path.
	return strdup(path);
}

char* directory_name(const char* path) {
	// Check for NULL.
	if (path == NULL)
		return NULL;

	// Find the last occurrence of the path separator.
	const char* file_name = strrchr(path, '/');

	// If the path separator was found, return the directory name before it.
	if (file_name != NULL) {
		size_t length = file_name - path;
		char* directory_name = malloc(length + 1);
		if (directory_name == NULL)
			return NULL;

		memcpy(directory_name, path, length);
		directory_name[length] = '\0';

		return directory_name;
	}

	// Otherwise, return the current directory.
	return strdup(".");
}

result_t list_files(char*** _files, const char* directory, result_t (*filter)(const char*)) {
	// Open the directory.
	DIR* dir = opendir(directory);
	if (dir == NULL)
		return failure("Failed to open the directory '%s' (%s).", directory, strerror(errno));

	// Count the files in the directory.
	size_t count = 0;
	size_t capacity = 16;
	
	char** files = malloc(capacity * sizeof(char*));
	if (files == NULL) {
		closedir(dir);
		return failure("Failed to allocate memory for the files.");
	}

	// Read the files in the directory.
	struct dirent* entry;
	while ((entry = readdir(dir)) != NULL) {
		// Skip the current and parent directories.
		if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0)
			continue;

		// Filter the file if necessary.
		if (filter != NULL) {
			result_t result = filter(entry->d_name);
			if (result != success())
				continue;
		}

		// Add the file to the list.
		if (count == capacity) {
			// Basic doubling strategy.
			capacity *= 2;
			char** new_files = realloc(files, capacity * sizeof(char*));
			if (new_files == NULL) {
				closedir(dir);
				for (size_t i = 0; i < count; i++)
					free(files[i]);
				free(files);
				return failure("Failed to reallocate memory for the files.");
			}
			files = new_files;
		}

		files[count] = strdup(entry->d_name);
		if (files[count] == NULL) {
			closedir(dir);
			for (size_t i = 0; i < count; i++)
				free(files[i]);
			free(files);
			return failure("Failed to allocate memory for the file.");
		}

		count++;
	}

	// Close the directory.
	closedir(dir);

	// Null-terminate the list of files.
	char** new_files = realloc(files, (count + 1) * sizeof(char*));
	if (new_files == NULL) {
		for (size_t i = 0; i < count; i++)
			free(files[i]);
		free(files);
		return failure("Failed to reallocate memory for the files.");
	}

	new_files[count] = NULL;

	// Output the list of files.
	*_files = new_files;

	return success();
}

result_t md5_fd(char** _hash, int fd) {
	// Initialize the digest context.
	EVP_MD_CTX* context = EVP_MD_CTX_new();
	if (context == NULL)
		return failure("Failed to create the MD5 digest context.");

	// Initialize the MD5 algorithm.
	const EVP_MD* algorithm = EVP_md5();
	if (algorithm == NULL) {
		EVP_MD_CTX_free(context);
		return failure("Failed to initialize the MD5 algorithm.");
	}

	// Initialize the digest.
	if (EVP_DigestInit_ex(context, algorithm, NULL) != 1) {
		EVP_MD_CTX_free(context);
		return failure("Failed to initialize the MD5 digest.");
	}

	char buffer[4096];
	ssize_t bytes_read;

	// Read the file descriptor into the buffer.
	while ((bytes_read = read(fd, buffer, sizeof(buffer))) > 0) {
		// Update the digest with the buffer.
		if (EVP_DigestUpdate(context, buffer, bytes_read) != 1) {
			EVP_MD_CTX_free(context);
			return failure("Failed to update the MD5 digest.");
		}
	}

	// Check for read errors.
	if (bytes_read < 0) {
		EVP_MD_CTX_free(context);
		return failure("Failed to read from the file descriptor (%s).", strerror(errno));
	}

	// Finalize the digest.
	unsigned char hash[EVP_MAX_MD_SIZE];
	unsigned int length;
	if (EVP_DigestFinal_ex(context, hash, &length) != 1) {
		EVP_MD_CTX_free(context);
		return failure("Failed to finalize the MD5 digest.");
	}

	// Free the digest context.
	EVP_MD_CTX_free(context);

	// Convert the hash to a hexadecimal string.
	char* hex_hash = malloc(length * 2 + 1);
	if (hex_hash == NULL)
		return failure("Failed to allocate memory for the hash.");

	for (unsigned int i = 0; i < length; i++)
		sprintf(hex_hash + i * 2, "%02x", hash[i]);

	// Output the hash.
	*_hash = hex_hash;

	return success();
}

result_t md5_file(char** _hash, const char* file) {
	// Open the file.
	int fd = open(file, O_RDONLY);
	if (fd < 0)
		return failure("Failed to open the file '%s' (%s).", file, strerror(errno));

	// Calculate the MD5 hash of the file.
	result_t result = md5_fd(_hash, fd);

	// Close the file.
	close(fd);

	return result;
}