#include <dirent.h>
#include <regex.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

//#define DENT_DUMP

struct dent {
	int type;
	size_t size;
	struct dent * parent;
	struct dent * next;
	struct dent * child;
	char name[];
};


// Create a new directory entry. Insert it into the tree if parent is given.
struct dent *
dent_create(int type, size_t size, struct dent * parent, char const * name)
{
	int namelen;
	struct dent * d;

	if (!name)
		return NULL;

	namelen = strlen(name);
	d = malloc(sizeof(struct dent) + namelen + 1);
	if (!d)
		return NULL;

	d->type = type;
	d->size = size;
	d->parent = parent;
	if (parent) {
		d->next = d->parent->child;
		d->parent->child = d;
	}
	else {
		d->next = NULL;
	}
	d->child = NULL;
	memcpy(d->name, name, namelen + 1);

	return d;
}


// Free a directory entry and its children
void
dent_free(struct dent * d)
{
	if (!d)
		return;

	dent_free(d->child);
	dent_free(d->next);
	free(d);
}


// Find a directory entry by name, in parent only (not recursive)
struct dent *
dent_find(struct dent * parent, char const * name)
{
	struct dent * d;

	for (d = parent->child; d != NULL; d = d->next) {
		if (strcmp(d->name, name) == 0) {
			return d;
		}
	}

	return NULL;
}


// Find the next entry for walking an entire directory tree.
struct dent *
dent_next(struct dent * d)
{
	if (!d)
		return NULL;

	if (d->child)
		return d->child;
	if (d->next)
		return d->next;

	while ((d = d->parent) != NULL)
		if (d->next)
			return d->next;

	return NULL;
}


// Get the recursive usage of a directory entry.
//
// This caches directory sizes in their dents.
size_t
dent_du(struct dent * d)
{
	struct dent * child;

	switch (d->type) {
	case DT_REG:
		return d->size;

	case DT_DIR:
		d->size = 0;
		for (child = d->child; child != NULL; child = child->next) {
			d->size += dent_du(child);
		}
		return d->size;
	}

	return 0;
}


#ifdef DENT_DUMP
void
dent_dump(FILE * out, struct dent * d, int depth)
{
	struct dent * child;

	if (!d)
		return;

	fprintf(out, "%*s %s", depth * 2, "-", d->name);
	switch (d->type) {
	case DT_DIR:
		fprintf(out, " (dir)");
		break;
	case DT_REG:
		fprintf(out, " (file, size=%zu)", d->size);
		break;
	default:
		fprintf(out, " (unknown!)");
		break;
	}
	fprintf(out, "\n");

	for (child = d->child; child != NULL; child = child->next) {
		dent_dump(out, child, depth + 1);
	}
}
#endif


// Parse an `ls` command.
int
cmd_ls(struct dent ** cwd, char * buf, size_t bufsiz, FILE * in, char const * arg)
{
	static regex_t dir, file;
	static int init = 0;

	if (init == 0) {
		init = 1;
		if (regcomp(&dir, "^(dir) ([[:alnum:].]+)\n?$", REG_EXTENDED) != 0
				|| regcomp(&file, "^([[:digit:]]+) ([[:alnum:].]+)\n?$", REG_EXTENDED) != 0)
		{
			fprintf(stderr, "ls: Unable to compile regexes");
			init = -1;
		}
	}
	if (init == -1) {
		return -1;
	}
	if (!buf && !in) {
		regfree(&dir);
		regfree(&file);
		init = 0;
		return 0;
	}

	if (arg) {
		fprintf(stderr, "%s: Unexpected arg to `ls`\n", __FUNCTION__);
		return -1;
	}

	while (fgets(buf, bufsiz, in)) {
		regmatch_t matches[3];
		struct dent * d;

		if (buf[0] == '$') {
			return 0;
		}
		else if (regexec(&dir, buf, 3, matches, 0) == 0) {
			buf[matches[2].rm_eo] = '\0';
			if ((d = dent_find(*cwd, buf + matches[2].rm_so)) == NULL) {
				d = dent_create(DT_DIR, 0, *cwd, buf + matches[2].rm_so);
			}
			if (d->type != DT_DIR) {
				fprintf(stderr, "%s: Mismatched type for %s\n", __FUNCTION__, d->name);
			}
		}
		else if (regexec(&file, buf, 3, matches, 0) == 0) {
			buf[matches[2].rm_eo] = '\0';
			if ((d = dent_find(*cwd, buf + matches[2].rm_so)) == NULL) {
				d = dent_create(DT_REG, atoi(buf + matches[1].rm_so),
						*cwd, buf + matches[2].rm_so);
			}

			if (d->type != DT_REG) {
				fprintf(stderr, "%s: Mismatched type for %s\n", __FUNCTION__, d->name);
			}
		}
		else {
			fprintf(stderr, "%s: Unexpected output: %s\n", __FUNCTION__, buf);
			return -1;
		}
	}

	buf[0] = '\0';
	return -1;
}


// Parse a `cd` command.
int
cmd_cd(struct dent ** cwd, char * buf, size_t bufsiz, FILE * in, char const * arg)
{
	if (!arg) {
		fprintf(stderr, "%s: Missing parameter\n", __FUNCTION__);
		return -1;
	}

	if (strcmp(arg, "/") == 0) {
		if (!*cwd)
			*cwd = dent_create(DT_DIR, 0, NULL, "/");

		while (*cwd && (*cwd)->parent)
			*cwd = (*cwd)->parent;
	}
	else if (strcmp(arg, "..") == 0) {
		if (!*cwd) {
			fprintf(stderr, "%s: No cwd!\n", __FUNCTION__);
			return -1;
		}
		if (!(*cwd)->parent) {
			fprintf(stderr, "%s: No parent dir!\n", __FUNCTION__);
			return -1;
		}
		*cwd = (*cwd)->parent;
	}
	else {
		struct dent * d = dent_find(*cwd, arg);
		if (!d) {
			fprintf(stderr, "%s: No subdirectory: %s\n", __FUNCTION__, arg);
			return -1;
		}
		if (d->type != DT_DIR) {
			fprintf(stderr, "%s: Is not a directory: %s\n", __FUNCTION__, arg);
			return -1;
		}
		*cwd = d;
	}

	return fgets(buf, bufsiz, in) ? 0 : -1;
}


int
main(int argc, char ** argv)
{
	char buf[BUFSIZ];
	struct dent * cwd = NULL, * root, * d;
	int i, part = 1;
	size_t s;

	while ((i = getopt(argc, argv, "p:")) != -1) {
		switch (i) {
		case 'p':
			part = atoi(optarg);
			break;

		default:
			return -1;
		}
	}

	i = fgets(buf, sizeof(buf), stdin) ? 0 : -1;
	while (i == 0) {
		char * tok, * toksave = NULL;

		tok = strtok_r(buf, " \n", &toksave);
		if (strcmp(tok, "$") != 0) {
			fprintf(stderr, "Not a command: %s\n", buf);
			return -1;
		}

		// Get the next command.
		//
		// Command parsers are responsible for parsing the command argument,
		// and any output of the command, and reading the line containing the
		// next command into `buf`.
		//
		// They should return 0 for success, -1 on error.
		tok = strtok_r(NULL, " \n", &toksave);
		if (strcmp(tok, "cd") == 0) {
			tok = strtok_r(NULL, "\n", &toksave);
			i = cmd_cd(&cwd, buf, sizeof(buf), stdin, tok);
		}
		else if (strcmp(tok, "ls") == 0) {
			tok = strtok_r(NULL, "\n", &toksave);
			i = cmd_ls(&cwd, buf, sizeof(buf), stdin, tok);
		}
		else {
			fprintf(stderr, "Unexpected command: %s\n", tok);
			return -1;
		}
	}

	root = cwd;
	while (root && root->parent)
		root = root->parent;

#ifdef DENT_DUMP
	dent_dump(stderr, root, 0);
#endif

	// Calculate disk usage sizes;
	dent_du(root);

	switch (part) {
	case 1:
		// Find directories whose size is less than a limit, and sum them.
		s = 0;
		for (d = root; d != NULL; d = dent_next(d)) {
			if (d->type == DT_DIR && d->size <= 100000)
				s += d->size;
		}
		printf("Sum of dir sizes where size <= 100000: %zu\n", s);
		break;

	case 2:
		// Find smallest directory to delete that drops disk space to 70M - 30M
		s = 70000000;
		for (d = root; d != NULL; d = dent_next(d)) {
			if (d->type == DT_DIR
					&& d->size < s
					&& root->size - d->size < 40000000)
				s = d->size;
		}
		printf("Smallest dir that can drop usage below 40000000: %zu\n", s);
		break;

	default:
		fprintf(stderr, "Unexpected part %d\n", part);
		return -1;
	}

	// Tidy up
	dent_free(root);
	cmd_ls(NULL, NULL, 0, NULL, NULL);

	return 0;
}