/*********************************************************************
 * bfs                                                               *
 * Copyright (C) 2015-2017 Tavian Barnes <tavianator@tavianator.com> *
 *                                                                   *
 * This program is free software. It comes without any warranty, to  *
 * the extent permitted by applicable law. You can redistribute it   *
 * and/or modify it under the terms of the Do What The Fuck You Want *
 * To Public License, Version 2, as published by Sam Hocevar. See    *
 * the COPYING file or http://www.wtfpl.net/ for more details.       *
 *********************************************************************/

#include "color.h"
#include "bftw.h"
#include "util.h"
#include <errno.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

struct ext_color {
	const char *ext;
	size_t len;

	const char *color;

	struct ext_color *next;
};

struct colors {
	const char *reset;
	const char *bold;
	const char *gray;
	const char *red;
	const char *green;
	const char *yellow;
	const char *blue;
	const char *magenta;
	const char *cyan;

	const char *normal;
	const char *file;
	const char *dir;
	const char *link;
	const char *multi_hard;
	const char *pipe;
	const char *door;
	const char *block;
	const char *chardev;
	const char *orphan;
	const char *missing;
	const char *socket;
	const char *setuid;
	const char *setgid;
	const char *capable;
	const char *sticky_ow;
	const char *ow;
	const char *sticky;
	const char *exec;

	const char *warning;
	const char *error;

	struct ext_color *ext_list;

	char *data;
};

struct color_name {
	const char *name;
	size_t offset;
};

#define COLOR_NAME(name, field) {name, offsetof(struct colors, field)}

static const struct color_name color_names[] = {
	COLOR_NAME("bd", block),
	COLOR_NAME("bld", bold),
	COLOR_NAME("blu", blue),
	COLOR_NAME("ca", capable),
	COLOR_NAME("cd", chardev),
	COLOR_NAME("cyn", cyan),
	COLOR_NAME("di", dir),
	COLOR_NAME("do", door),
	COLOR_NAME("er", error),
	COLOR_NAME("ex", exec),
	COLOR_NAME("fi", file),
	COLOR_NAME("grn", green),
	COLOR_NAME("gry", gray),
	COLOR_NAME("ln", link),
	COLOR_NAME("mag", magenta),
	COLOR_NAME("mh", multi_hard),
	COLOR_NAME("mi", missing),
	COLOR_NAME("no", normal),
	COLOR_NAME("or", orphan),
	COLOR_NAME("ow", ow),
	COLOR_NAME("pi", pipe),
	COLOR_NAME("red", red),
	COLOR_NAME("rs", reset),
	COLOR_NAME("sg", setgid),
	COLOR_NAME("so", socket),
	COLOR_NAME("st", sticky),
	COLOR_NAME("su", setuid),
	COLOR_NAME("tw", sticky_ow),
	COLOR_NAME("wr", warning),
	COLOR_NAME("ylw", yellow),
	{0},
};

static const char **look_up_color(const struct colors *colors, const char *name) {
	for (const struct color_name *entry = color_names; entry->name; ++entry) {
		if (strcmp(name, entry->name) == 0) {
			return (const char **)((char *)colors + entry->offset);
		}
	}

	return NULL;
}

static const char *get_color(const struct colors *colors, const char *name) {
	const char **color = look_up_color(colors, name);
	if (color) {
		return *color;
	} else {
		return NULL;
	}
}

static void set_color(struct colors *colors, const char *name, const char *value) {
	const char **color = look_up_color(colors, name);
	if (color) {
		*color = value;
	}
}

struct colors *parse_colors(const char *ls_colors) {
	struct colors *colors = malloc(sizeof(struct colors));
	if (!colors) {
		goto done;
	}

	// From man console_codes
	colors->reset   = "0";
	colors->bold    = "01";
	colors->gray    = "01;30";
	colors->red     = "01;31";
	colors->green   = "01;32";
	colors->yellow  = "01;33";
	colors->blue    = "01;34";
	colors->magenta = "01;35";
	colors->cyan    = "01;36";

	// Defaults generated by dircolors --print-database
	colors->normal     = NULL;
	colors->file       = NULL;
	colors->dir        = "01;34";
	colors->link       = "01;36";
	colors->multi_hard = NULL;
	colors->pipe       = "40;33";
	colors->socket     = "01;35";
	colors->door       = "01;35";
	colors->block      = "40;33;01";
	colors->chardev    = "40;33;01";
	colors->orphan     = "40;31;01";
	colors->setuid     = "37;41";
	colors->setgid     = "30;43";
	colors->capable    = "30;41";
	colors->sticky_ow  = "30;42";
	colors->ow         = "34;42";
	colors->sticky     = "37;44";
	colors->exec       = "01;32";
	colors->warning    = "40;33;01";
	colors->error      = "40;31;01";
	colors->ext_list   = NULL;
	colors->data       = NULL;

	if (ls_colors) {
		colors->data = strdup(ls_colors);
	}

	if (!colors->data) {
		goto done;
	}

	char *start = colors->data;
	char *end;
	struct ext_color *ext;
	for (end = strchr(start, ':'); *start && end; start = end + 1, end = strchr(start, ':')) {
		char *equals = strchr(start, '=');
		if (!equals) {
			continue;
		}

		*equals = '\0';
		*end = '\0';

		const char *key = start;
		const char *value = equals + 1;

		// Ignore all-zero values
		if (strspn(value, "0") == strlen(value)) {
			continue;
		}

		if (key[0] == '*') {
			ext = malloc(sizeof(struct ext_color));
			if (ext) {
				ext->ext = key + 1;
				ext->len = strlen(ext->ext);
				ext->color = value;
				ext->next = colors->ext_list;
				colors->ext_list = ext;
			}
		} else {
			set_color(colors, key, value);
		}
	}

done:
	return colors;
}

void free_colors(struct colors *colors) {
	if (colors) {
		struct ext_color *ext = colors->ext_list;
		while (ext) {
			struct ext_color *saved = ext;
			ext = ext->next;
			free(saved);
		}

		free(colors->data);
		free(colors);
	}
}

CFILE *cfopen(const char *path, const struct colors *colors) {
	CFILE *cfile = malloc(sizeof(*cfile));
	if (!cfile) {
		return NULL;
	}

	cfile->close = false;
	cfile->file = fopen(path, "wb");
	if (!cfile->file) {
		cfclose(cfile);
		return NULL;
	}
	cfile->close = true;

	if (isatty(fileno(cfile->file))) {
		cfile->colors = colors;
	} else {
		cfile->colors = NULL;
	}

	return cfile;
}

CFILE *cfdup(FILE *file, const struct colors *colors) {
	CFILE *cfile = malloc(sizeof(*cfile));
	if (!cfile) {
		return NULL;
	}

	cfile->close = false;
	cfile->file = file;

	if (isatty(fileno(file))) {
		cfile->colors = colors;
	} else {
		cfile->colors = NULL;
	}

	return cfile;
}

int cfclose(CFILE *cfile) {
	int ret = 0;
	if (cfile) {
		if (cfile->close) {
			ret = fclose(cfile->file);
		}
		free(cfile);
	}
	return ret;
}

static const char *file_color(const struct colors *colors, const char *filename, const struct BFTW *ftwbuf) {
	const struct stat *sb = ftwbuf->statbuf;
	if (!sb) {
		return colors->orphan;
	}

	const char *color = NULL;

	switch (sb->st_mode & S_IFMT) {
	case S_IFREG:
		if (sb->st_mode & S_ISUID) {
			color = colors->setuid;
		} else if (sb->st_mode & S_ISGID) {
			color = colors->setgid;
		} else if (sb->st_mode & 0111) {
			color = colors->exec;
		}

		if (!color && sb->st_nlink > 1) {
			color = colors->multi_hard;
		}

		if (!color) {
			size_t namelen = strlen(filename);

			for (struct ext_color *ext = colors->ext_list; ext; ext = ext->next) {
				if (namelen >= ext->len && memcmp(filename + namelen - ext->len, ext->ext, ext->len) == 0) {
					color = ext->color;
					break;
				}
			}
		}

		if (!color) {
			color = colors->file;
		}

		break;

	case S_IFDIR:
		if (sb->st_mode & S_ISVTX) {
			if (sb->st_mode & S_IWOTH) {
				color = colors->sticky_ow;
			} else {
				color = colors->sticky;
			}
		} else if (sb->st_mode & S_IWOTH) {
			color = colors->ow;
		}

		if (!color) {
			color = colors->dir;
		}

		break;

	case S_IFLNK:
		if (faccessat(ftwbuf->at_fd, ftwbuf->at_path, F_OK, 0) == 0) {
			color = colors->link;
		} else {
			color = colors->orphan;
		}
		break;

	case S_IFBLK:
		color = colors->block;
		break;
	case S_IFCHR:
		color = colors->chardev;
		break;
	case S_IFIFO:
		color = colors->pipe;
		break;
	case S_IFSOCK:
		color = colors->socket;
		break;

#ifdef S_IFDOOR
	case S_IFDOOR:
		color = colors->door;
		break;
#endif
	}

	if (!color) {
		color = colors->normal;
	}

	return color;
}

static int print_esc(const char *esc, FILE *file) {
	if (fputs("\033[", file) == EOF) {
		return -1;
	}
	if (fputs(esc, file) == EOF) {
		return -1;
	}
	if (fputs("m", file) == EOF) {
		return -1;
	}

	return 0;
}

static int print_path(CFILE *cfile, const struct BFTW *ftwbuf) {
	const struct colors *colors = cfile->colors;
	FILE *file = cfile->file;
	const char *path = ftwbuf->path;

	if (!colors) {
		return fputs(path, file) == EOF ? -1 : 0;
	}

	const char *filename = path + ftwbuf->nameoff;

	if (colors->dir) {
		if (print_esc(colors->dir, file) != 0) {
			return -1;
		}
	}
	if (fwrite(path, 1, ftwbuf->nameoff, file) != ftwbuf->nameoff) {
		return -1;
	}
	if (colors->dir) {
		if (print_esc(colors->reset, file) != 0) {
			return -1;
		}
	}

	const char *color = file_color(colors, filename, ftwbuf);
	if (color) {
		if (print_esc(color, file) != 0) {
			return -1;
		}
	}
	if (fputs(filename, file) == EOF) {
		return -1;
	}
	if (color) {
		if (print_esc(colors->reset, file) != 0) {
			return -1;
		}
	}

	return 0;
}

static int print_link(CFILE *cfile, const struct BFTW *ftwbuf) {
	int ret = -1;

	char *target = xreadlinkat(ftwbuf->at_fd, ftwbuf->at_path, 0);
	if (!target) {
		goto done;
	}

	struct BFTW altbuf = *ftwbuf;
	altbuf.path = target;
	altbuf.nameoff = xbasename(target) - target;

	struct stat statbuf;
	if (fstatat(ftwbuf->at_fd, ftwbuf->at_path, &statbuf, 0) == 0) {
		altbuf.statbuf = &statbuf;
	} else {
		altbuf.statbuf = NULL;
	}

	ret = print_path(cfile, &altbuf);

done:
	free(target);
	return ret;
}

int cfprintf(CFILE *cfile, const char *format, ...) {
	const struct colors *colors = cfile->colors;
	FILE *file = cfile->file;

	int ret = -1;

	va_list args;
	va_start(args, format);

	for (const char *i = format; *i; ++i) {
		if (*i == '%') {
			switch (*++i) {
			case '%':
				goto one_char;

			case 'c':
				if (fputc(va_arg(args, int), file) == EOF) {
					goto done;
				}
				break;

			case 's':
				if (fputs(va_arg(args, const char *), file) == EOF) {
					goto done;
				}
				break;

			case 'P':
				if (print_path(cfile, va_arg(args, const struct BFTW *)) != 0) {
					goto done;
				}
				break;

			case 'L':
				if (print_link(cfile, va_arg(args, const struct BFTW *)) != 0) {
					goto done;
				}
				break;

			case '{': {
				++i;
				const char *end = strchr(i, '}');
				if (!end) {
					goto invalid;
				}
				if (!colors) {
					i = end;
					break;
				}

				size_t len = end - i;
				char name[len + 1];
				memcpy(name, i, len);
				name[len] = '\0';

				const char *esc = get_color(colors, name);
				if (!esc) {
					goto invalid;
				}
				if (print_esc(esc, file) != 0) {
					goto done;
				}

				i = end;
				break;
			}

			default:
			invalid:
				errno = EINVAL;
				goto done;
			}

			continue;
		}

	one_char:
		if (fputc(*i, file) == EOF) {
			goto done;
		}
	}

	ret = 0;

done:
	va_end(args);
	return ret;
}