#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <unistr.h>
#include <time.h>

#define ATTR_NONE 0
#define ATTR_BOLD 1
#define ATTR_UNDER 4
#define ATTR_BLINK 5

#define SYM_GIT "\ue0a0"
#define SYM_LN "\ue0a1"
#define SYM_LOCK "\ue0a2"
#define SYM_RIGHTBLK "\ue0b0"
#define SYM_RIGHT "\ue0b1"
#define SYM_LEFTBLK "\ue0b2"
#define SYM_LEFT "\ue0b3"

#define COL_DEFAULT_FG 7
#define COL_DEFAULT_BG 0
#define COL_LOGIN_USER_FG 202
#define COL_LOGIN_USER_BG 232
#define COL_LOGIN_ROOT_FG 232
#define COL_LOGIN_ROOT_BG 124
#define COL_SCREEN_FG 232
#define COL_SCREEN_BG 27
#define COL_PATH_FG 202
#define COL_PATH2_FG 232
#define COL_PATH_BG 232
#define COL_GIT_CLEAN_FG 232
#define COL_GIT_CLEAN_BG 88
#define COL_GIT_DIRTY_FG 232
#define COL_GIT_DIRTY_BG 46
#define COL_LOCK_FG 196
#define COL_TIME_FG 202
#define COL_TIME_BG 232
#define COL_LOAD_OK_FG 202
#define COL_LOAD_OK_BG 232
#define COL_LOAD_WARN_FG 232
#define COL_LOAD_WARN_BG 3
#define COL_LOAD_HIGH_FG 232
#define COL_LOAD_HIGH_BG 202
#define COL_LOAD_OVER_FG 15
#define COL_LOAD_OVER_BG 196
#define COL_PROMPT_OK_FG 232
#define COL_PROMPT_OK_BG 34
#define COL_ERR_FG 15
#define COL_ERR_BG 196


#define PATH_MAX 4096

#ifndef uint
#define uint unsigned int
#endif

/* -------------------------------------------------- */

class Config {
public:
	Config(int argc, char **argv) {
		if (argc != 5) {
			fprintf(stderr, "Usage: %s red_threshold yellow_threshold terminal_width last_exit_status\n", argv[0]);
			this->failed = true;
			return;
		}
		
		this->red_thresh = strtof(argv[1], NULL);
		this->yellow_thresh = strtof(argv[2], NULL);
		this->term_width = strtol(argv[3], NULL, 10);
		this->last_exit_status = strtol(argv[4], NULL, 10);
		this->is_root = (geteuid() == 0);
		
		this->failed = false;
	}
	
	uint term_width;
	float red_thresh;
	float yellow_thresh;
	uint last_exit_status;
	bool is_root;
	
	bool failed;
};

class Segment {
public:
	Segment(const char *str, bool printable=true) {
		this->raw_length = strlen(str);
		this->raw_str = new char[(this->raw_length + 1) * sizeof(char)]; // +1 for null
		strcpy(this->raw_str, str);
		
		if (printable) {
			this->printable_length = u8_mbsnlen((const uint8_t*)str, this->raw_length);
		} else {
			this->printable_length = 0;
		}
	}
	
	~Segment(void) {
		delete[] this->raw_str;
	}
	
	char *raw_str;
	size_t raw_length;
	size_t printable_length;
};

class Output {
public:
	Output(void) {
		this->printable_length = 0;
		this->raw_length = 0;
		this->ptr = 0;
		this->max = 256; // a reasonable amount?
		
		this->segments = new Segment*[this->max];
		
		this->set_colors(COL_DEFAULT_FG, COL_DEFAULT_BG);
	}
	
	~Output(void) {
		for (uint i = 0; i < this->ptr; i++) {
			delete this->segments[i];
		}
		
		delete[] this->segments;
	}
	
	void output(FILE *stream) {
		for (uint i = 0; i < this->ptr; i++) {
			fputs(this->segments[i]->raw_str, stream);
		}
	}
	
	void set_colors(uint fg, uint bg, const char *transition=NULL, bool left_to_right=true) {
		if (transition) {
			this->push(" ");
		}
		
		if (left_to_right) this->_push_color(bg, true);
		
		if (transition) {
			this->_push_color(left_to_right? this->col_bg : bg);
			this->push(transition);
			if (left_to_right) this->push(" ");
		}
		
		if (not left_to_right) this->_push_color(bg, true);
		this->_push_color(fg);
		
		if (transition and not left_to_right) this->push(" ");
		
		this->col_fg = fg;
		this->col_bg = bg;
	}
	
	void _push_color(uint color=ATTR_NONE, bool bg=false) {
		if (color == ATTR_NONE) {
			this->set_attr();
		} else {
			char tmp[16];
			snprintf(tmp, 16, "\e[%d;5;%dm", (bg? 48:38), color);
			
			this->push(new Segment(tmp, false));
		}
	}
	
	void set_attr(uint attr=ATTR_NONE) {
		char tmp[10];
		snprintf(tmp, 10, "\e[%dm", attr);
		
		this->push(new Segment(tmp, false));
	}
	
	void push(Segment *segment) {
		if (this->ptr + 1 == this->max) {
			fprintf(stderr, "Output.segments full, this will probably fail.\n");
		}
		
		this->segments[this->ptr++] = segment;
		this->printable_length += segment->printable_length;
		this->raw_length += segment->raw_length;
	}
	
	void push(const char *str) {
		this->push(new Segment(str));
	}
	
	void push(uint n) {
		char tmp[1024];
		sprintf(tmp, "%d", n);
		this->push(tmp);
	}
	
	Segment **segments;
	size_t ptr, max;
	size_t raw_length;
	size_t printable_length;
	
	uint col_fg, col_bg;
};

void render_screen(Output *output) {
	char *tmp = getenv("WINDOW");
	
	if (tmp) {
		output->set_colors(COL_SCREEN_FG, COL_SCREEN_BG, SYM_RIGHTBLK);
		output->push(tmp);
	}
}

void render_path(Output *output) {
	char buf[PATH_MAX];
	
	getcwd(buf, PATH_MAX);
	
	// tokenize path TODO replace homes with ~
	char *ptr = buf;
	char *ptr_end;
	char *tokens[100]; // a sensible default?
	char token[4096];
	uint token_count = 0;
	uint token_length;
	
	while (*ptr) { // until NULL terminator
		ptr = strchrnul(ptr, '/') + 1;
		ptr_end = strchrnul(ptr, '/');
		token_length = ptr_end - ptr;
		memcpy(token, ptr, token_length);
		token[token_length] = 0;
		
		if (token_length) {
			tokens[token_count] = new char[(token_length + 1) * sizeof(char)];
			memcpy(tokens[token_count], token, token_length);
			tokens[token_count][token_length] = 0;
			token_count++;
		}
	}
	
	// render path with increasing background brightness
	uint fg = COL_PATH_FG;
	uint bg = COL_PATH_BG;
	
	for (uint i = 0; i < token_count; i++) {
		if (bg == 241) fg = COL_PATH2_FG;
		
		if (bg < 255) {
			output->set_colors(fg, bg++, SYM_RIGHTBLK);
		} else {
			output->push(" ");
			output->push(SYM_RIGHT);
			output->push(" ");
		}
		
		output->push(tokens[i]);
		delete[] tokens[i];
	}
}

void render_git(Output *output) {
	FILE *pipe;
	char branch_name[1024];
	
	pipe = popen("git rev-parse --abbrev-ref HEAD 2> /dev/null", "r");
	
	if (pipe) {
		if (fgets(branch_name, 1024, pipe)) {
			branch_name[strlen(branch_name) - 1] = 0; // kill the newline
			
			// rebuild index
			system("git update-index -q --ignore-submodules --refresh");
			
			// so that I can check if there are unstaged changes
			if (system("git diff-files --quiet --ignore-submodules")) {
				output->set_colors(COL_GIT_CLEAN_FG, COL_GIT_CLEAN_BG, SYM_RIGHTBLK);
			} else {
				output->set_colors(COL_GIT_DIRTY_FG, COL_GIT_DIRTY_BG, SYM_RIGHTBLK);
			}
			
			output->push(SYM_GIT);
			output->push(" ");
			output->push(branch_name);
		}
		
		pclose(pipe);
	} else {
		fprintf(stderr, "--- failed to run git\n");
	}
}

void render_lock(Output *output) {
	if (access(".", W_OK) != 0) {
// 		output->set_colors(COL_LOCK_FG, COL_LOCK_BG, SYM_RIGHTBLK);
		output->_push_color(COL_LOCK_FG);
		output->push(" ");
		output->push(SYM_LOCK);
	}
}

void render_padding(const uint term_width, const uint length, FILE *stream) {
	int delta = term_width - length;
	
	if (delta >= 0) {
		delta--; // fish refuses to write to the last column
		
		for (uint i = 0; i < (uint) delta; i++) {
			fputc(' ', stream);
		}
	} else {
		fputc('!', stream);
	}
}

void render_time(Output *output) {
	time_t rawtime;
	struct tm *info;
	char buffer[1024];
	
	time(&rawtime);
	info = localtime(&rawtime);
	
	strftime(buffer, 1024, "%H:%M", info);
	
	output->push(buffer);
}

void render_sysload(Output *output) {
	double load;
	char tmp[16];
	uint bg, fg;
	
	if (getloadavg(&load, 1) == 1) {
		if (load > 12) {
			fg = COL_LOAD_OVER_FG;
			bg = COL_LOAD_OVER_BG;
		} else if (load > 8) {
			fg = COL_LOAD_HIGH_FG;
			bg = COL_LOAD_HIGH_BG;
		} else if (load > 4) {
			fg = COL_LOAD_WARN_FG;
			bg = COL_LOAD_WARN_BG;
		} else {
			fg = COL_LOAD_OK_FG;
			bg = COL_LOAD_OK_BG;
		}
		
		output->set_colors(fg, bg, SYM_RIGHTBLK);
		sprintf(tmp, "%.2f", load);
		output->push(tmp);
	} else {
		output->set_colors(COL_ERR_FG, COL_ERR_BG, SYM_RIGHTBLK);
		output->push("ERR");
	}
}

/* -------------------------------------------------- */

int main(int argc, char **argv) {
	Config config(argc, argv);
	if (config.failed) return 1;
	
	/* --------------------------------------------------
	 * output init
	 */
	
	Output top_left;
	Output top_right;
	Output prompt;
	
	/* --------------------------------------------------
	 * top_left: user@host[:screen] path > to > cwd [git branch] [lock]
	 */
	
	char buf[PATH_MAX];
	
	if (config.is_root) {
		top_left.set_colors(COL_LOGIN_ROOT_FG, COL_LOGIN_ROOT_BG);
	} else {
		top_left.set_colors(COL_LOGIN_USER_FG, COL_LOGIN_USER_BG);
	}
	
	top_left.push(" ");
	top_left.push(getlogin());
	top_left.push("@");
	
	gethostname(buf, 256);
	top_left.push(buf);
	
	render_screen(&top_left);
	render_path(&top_left);
	render_lock(&top_left);
	render_git(&top_left);
	
	top_left.set_colors(COL_DEFAULT_FG, COL_DEFAULT_BG, SYM_RIGHTBLK);
	
	/* --------------------------------------------------
	 * prompt: [last_exit_status:]$|# >
	 */
	
	if (config.last_exit_status) {
		prompt.set_colors(COL_ERR_FG, COL_ERR_BG);
		prompt.push(" ");
		prompt.push(config.last_exit_status);
	} else {
		prompt.set_colors(COL_PROMPT_OK_FG, COL_PROMPT_OK_BG);
		prompt.push(" ");
		
		if (config.is_root) {
			prompt.push("#");
		} else {
			prompt.push("$");
		}
	}
	
	prompt.set_colors(COL_DEFAULT_FG, COL_DEFAULT_BG, SYM_RIGHTBLK);
	
	/* --------------------------------------------------
	 * top_right: < time < spaces < sysload
	 */
	
	top_right.set_colors(COL_TIME_FG, COL_TIME_BG, SYM_LEFTBLK, false);
	render_time(&top_right);
	render_sysload(&top_right);
	top_right.set_colors(COL_DEFAULT_FG, COL_DEFAULT_BG, SYM_RIGHTBLK);
	
	/* --------------------------------------------------
	 * output the prompt
	 */
	
	top_left.output(stdout);
	render_padding(config.term_width, top_left.printable_length + top_right.printable_length, stdout);
	top_right.output(stdout);
	fputs("\n", stdout);
	prompt.output(stdout);
	
	return 0;
}