/*
 * Copyright (c) 2009 Sviatoslav Chagaev <slava@zb.lv>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#include <unistd.h>
#include <err.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <dirent.h>

#define DEFAULT_MAXLEVEL -1
#define DEFAULT_TAB "  "

typedef struct stack stack_t;
struct stack {
	int             nitems;
	char          **items;
};
/* Duplicate and s=NULL entries are not allowed */
typedef struct tree tree_t;
struct tree {
	char           *s;
	tree_t         *left;	/* lesser */
	tree_t         *right;
};
struct tree_arg {
	const char     *path;
	int             level;
};

int             dflag;
int             Fflag;
int             gflag;
int             Gflag;
int             zflag;
int             maxlevel = DEFAULT_MAXLEVEL;
char           *tab = DEFAULT_TAB;
char            ftype_chars[] = {
	'/',
	'*',
	'@',
	'=',
	'|'
};
enum ftype_chars {
	FT_DIR,
	FT_EXEC,
	FT_SYMLNK,
	FT_SOCK,
	FT_FIFO
};

void            proc_args(int argc, char **argv);
void            usage(void);
void            lstree(const char *path, int level, int maxlevel, const char *tab);
void            print_file(const char *path, const char *name, int level);
void            print_tab(int n);
char           *glue_path(const char *s1, const char *s2);
void            stack_init(stack_t * ps);
void            stack_push(stack_t * ps, const char *s);
char           *stack_pop(stack_t * ps);
tree_t         *tree_new(const char *s);
int             tree_cmp(tree_t * t1, tree_t * t2);
tree_t         *tree_add(tree_t * tree, const char *s);
void            tree_walk_lmr(tree_t * t, void *arg, int (*f) (tree_t *, void *));
void            tree_walk_lrm(tree_t * t, void *arg, int (*f) (tree_t *, void *));
int             tree_free_str_and_node(tree_t * t, void *unused);
void            tree_free(tree_t * t);
int             tree_print_node(tree_t * t, void *varg);
void            tree_print(tree_t * t, struct tree_arg * arg);

int
main(int argc, char **argv)
{
	int             i;
	int             found = 0;

	proc_args(argc, argv);

/*
	printf("dflag = %i; Fflag = %i; gflag = %i; Gflag = %i; maxlevel = %i; tab = \"%s\";\n", dflag, Fflag, gflag, Gflag, maxlevel, tab);
	printf("path_list:\n");
	for (i = 1; i < argc; ++i) {
		if (argv[i]) {
			printf("\t\"%s\"\n", argv[i]);
		}
	}

	printf("---------------\n");
*/

	for (i = 1; i < argc; ++i) {
		if (argv[i]) {
			if (found)
				putchar('\n');
			else
				found = 1;
			print_file("", argv[i], 0);
			lstree(argv[i], 1, (maxlevel < 1) ? (maxlevel) : (maxlevel + 1), tab);
		}
	}
	if (!found)
		lstree(".", 0, maxlevel, tab);
}

void
proc_args(int argc, char **argv)
{
	int             i;

	for (i = 1; i < argc; ++i) {
		if (!strcmp(argv[i], "-h") || !strcmp(argv[i], "--help"))
			usage();
		else if (!strcmp(argv[i], "-m")) {
			if (i < (argc - 1)) {
				maxlevel = atoi(argv[i + 1]);
				argv[i] = NULL;
				argv[++i] = NULL;
			} else
				usage();
		} else if (!strcmp(argv[i], "-t")) {
			if (i < (argc - 1)) {
				tab = argv[i + 1];
				argv[i] = NULL;
				argv[++i] = NULL;
			} else
				usage();
		} else if (!strcmp(argv[i], "-d")) {
			dflag = 1;
			argv[i] = NULL;
		} else if (!strcmp(argv[i], "-F")) {
			Fflag = 1;
			argv[i] = NULL;
		} else if (!strcmp(argv[i], "-g")) {
			dflag = 1;
			gflag = 1;
			argv[i] = NULL;
		} else if (!strcmp(argv[i], "-G")) {
			dflag = 1;
			Gflag = 1;
			gflag = 1;
			argv[i] = NULL;
		} else if (!strcmp(argv[i], "-z")) {
			zflag = 1;
			argv[i] = NULL;
		}
	}
}

void
usage(void)
{
	extern char    *__progname;

	printf("List files and directories in a fancy way.\n"
	       "usage:\t%s [-d] [-m <int>] [-t <string>] [path1 path2...]\n"
	       "\t-z -- reverse sort order\n"
	       "\t-d -- show directories first\n"
	       "\t-m -- max depth (inf. by default)\n"
	       "\t-t -- tabulation string (four spaces by default)\n"
	       "\t-F -- display a char denoting file's type\n"
	       "\t-g -- draw a compact pseudo-graphics tree (-d implied)\n"
	       "\t-G -- draw a pseudo-graphics tree (-d implied)\n",
	       __progname);
	exit(0);
}

void
lstree(const char *path, int level, int maxlevel, const char *tab)
{
	int             i;
	struct stat     sb;
	DIR            *dp;
	struct dirent  *de;
	char           *p = NULL;
	tree_t         *ftree = NULL;
	tree_t         *dtree = NULL;
	struct tree_arg targ;
	
	if (level > maxlevel && maxlevel > -1)
		return;

	if ((dp = opendir(path)) == NULL) {
		warn("\"%s\"", path);
		return;
	}
	while ((de = readdir(dp)) != NULL) {
		if (de->d_type == DT_DIR) {
			if (!strcmp(de->d_name, ".") || !strcmp(de->d_name, ".."))
				continue;
			dtree = tree_add(dtree, strdup(de->d_name));
		} else {
			if (dflag) {
				ftree = tree_add(ftree, strdup(de->d_name));
			} else {
				dtree = tree_add(dtree, strdup(de->d_name));
			}
		}
	}
	closedir(dp);
	targ.path = path;
	targ.level = level;
	tree_print(dtree, &targ);
	tree_free(dtree);
	if (dflag) {
		tree_print(ftree, &targ);
		tree_free(ftree);
	}
}

void
print_file(const char *path, const char *name, int level)
{
	struct stat     sb;
	char           *p;

	print_tab(level);
	printf("%s", name);
	if (Fflag) {
		if (path[0] == '\0')
			p = strdup(name);
		else
			p = glue_path(path, name);
		if (stat(p, &sb) == -1) {
			warn("\"%s\"", p);
			free(p);
			return;
		}
		if (S_ISDIR(sb.st_mode)) {
			putchar(ftype_chars[FT_DIR]);
		} else if (S_ISFIFO(sb.st_mode)) {
			putchar(ftype_chars[FT_FIFO]);
		} else if (S_ISLNK(sb.st_mode)) {
			putchar(ftype_chars[FT_SYMLNK]);
		} else if (S_ISSOCK(sb.st_mode)) {
			putchar(ftype_chars[FT_SOCK]);
		} else if (S_ISREG(sb.st_mode)) {
			if (sb.st_mode & (S_IXUSR | S_IXGRP | S_IXOTH))
				putchar(ftype_chars[FT_EXEC]);
		}
		free(p);
	}
	putchar('\n');
}

void
print_tab(int n)
{
	int             i;

	if (gflag) {
		if (n == 0)
			return;
		if (Gflag) {
			for (i = 0; i < n; ++i)
				printf("| ");
			putchar('\n');
			for (i = 0; i < (n - 1); ++i)
				printf("| ");
			printf("+-");
		} else {
			for (i = 0; i < (n - 1); ++i)
				printf("| ");
			printf("`-");
		}
	} else {
		for (i = 0; i < n; ++i)
			printf("%s", tab);
	}
}

char           *
glue_path(const char *s1, const char *s2)
{
	size_t          s1len, s2len, len;
	char           *buf;

	s1len = strlen(s1);
	s2len = strlen(s2);
	len = s1len + s2len;

	if ((buf = malloc(len + 1 + 1)) == NULL)
		return NULL;

	strcpy(buf, s1);
	buf[s1len] = '/';
	strcpy(buf + s1len + 1, s2);

	return buf;
}

tree_t         *
tree_new(const char *s)
{
	tree_t         *t;

	t = malloc(sizeof(tree_t));
	t->s = s;
	t->left = NULL;
	t->right = NULL;

	return t;
}

int
tree_cmp(tree_t * t1, tree_t * t2)
{
	return strcmp(t1->s, t2->s);
}

tree_t         *
tree_add(tree_t * tree, const char *s)
{
	int             cmp;
	tree_t         *t, *pt;

	t = tree_new(s);
	if (tree != NULL) {
		pt = tree;
		while (cmp = tree_cmp(t, pt)) {
			if (cmp < 0) {
				if (pt->left == NULL) {
					pt->left = t;
					break;
				} else {
					pt = pt->left;
				}
			} else if (cmp > 0) {
				if (pt->right == NULL) {
					pt->right = t;
					break;
				} else {
					pt = pt->right;
				}
			}
		}
		return tree;
	} else {
		return t;
	}
}

/*
 * left,middle,right
 * Stops if f returns not 0
 */
void
tree_walk_lmr(tree_t * t, void *arg, int (*f) (tree_t *, void *))
{
	if (t == NULL)
		return;

	tree_walk_lmr(t->left, arg, f);
	if ((*f) (t, arg))
		return;
	tree_walk_lmr(t->right, arg, f);
}

void
tree_walk_lrm(tree_t * t, void *arg, int (*f) (tree_t *, void *))
{
	if (t == NULL)
		return;

	tree_walk_lrm(t->left, arg, f);
	tree_walk_lrm(t->right, arg, f);
	if ((*f) (t, arg))
		return;
}

void
tree_walk_rlm(tree_t * t, void *arg, int (*f) (tree_t *, void *))
{
	if (t == NULL)
		return;

	tree_walk_rlm(t->right, arg, f);
	tree_walk_rlm(t->left, arg, f);
	if ((*f) (t, arg))
		return;
}

int
tree_free_str_and_node(tree_t * t, void *unused)
{
	free(t->s);
	free(t);
	return 0;
}

void
tree_free(tree_t * t)
{
	tree_walk_lrm(t, NULL, tree_free_str_and_node);
}

int
tree_print_node(tree_t * t, void *varg)
{
	struct stat sb;
	struct tree_arg *arg = varg;
	char *p;

	p = glue_path(arg->path, t->s);
	if (stat(p, &sb) == -1) {
		warn("\"%s\"", p);
		free(p);
		return 0;
	}
	if (S_ISDIR(sb.st_mode)) {
		print_file(arg->path, t->s, arg->level);
		lstree(p, arg->level + 1, maxlevel, tab);
	} else {
		print_file(arg->path, t->s, arg->level);
	}
	free(p);

	return 0;
}

void
tree_print(tree_t * t, struct tree_arg * arg)
{
	if (zflag)
		tree_walk_lrm(t, (void *) arg, tree_print_node);
	else
		tree_walk_lmr(t, (void *) arg, tree_print_node);
}