#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <string>
#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>
#include <map>
#include <sys/vfs.h>
#include <libgen.h>
#include <ctype.h>

class Space {
	public:
		Space(const char *name, float available, float total) {
			if (strlen(name) > 15) {
				strcpy(this->name, "!!!");
			} else {
				strcpy(this->name, name);
			}
			
			this->available = available;
			
			if (total > 0) {
				this->percent = available*100 / total;
				this->active = true;
			} else {
				this->percent = 100;
				this->active = false;
			}
		}
		
		char name[16];
		float available;
		float percent;
		bool active;
};

typedef struct {
	float load;
	char proc[16];
	unsigned int color;
} Sysload;

typedef struct {
	unsigned short int m1, m2;
	bool long_names;
} Config;

/// Data gatherers
void get_available_space(std::vector<Space> &space);
void get_memswap(std::vector<Space> &space);
void get_sysload(Sysload &sysload);

/// Formatting and parsing
void append(char *output, const char *string, int color=-1);
const char *float_to_str(float value, char *output);
int space_to_color(Space &space);
float grepf(const std::string &data, const char *marker);
std::vector<std::string> read_lines_from_file(std::string filename);
bool column_contains_substr(std::string line, unsigned int column, std::string substr, int position=0);
std::string get_column(std::string line, unsigned int column);
bool str_in_list(std::string str, std::vector<std::string> list);

/// Global
std::map<int, std::string> color_map;
Config config;
std::vector<std::string> ignored;

enum {
	COL_GREEN,
	COL_DARKGREEN,
	COL_YELLOW,
	COL_BLUE,
	COL_RED,
	COL_RFLASH,
	COL_NONE
};

int main(int argc, char **argv) {
	char output[512];
	char tmp[16];
	
	std::vector<Space> space;
	Sysload sysload;
	
	output[0] = 0;
	
	color_map[COL_GREEN] = "|GREEN|";
	color_map[COL_DARKGREEN] = "|DARKGREEN|";
	color_map[COL_YELLOW] = "|YELLOW|";
	color_map[COL_BLUE] = "|BLUE|";
	color_map[COL_RED] = "|RED|";
	color_map[COL_RFLASH] = "|RFLASH|";
	color_map[COL_NONE] = "|NONE|";
	
	ignored.push_back("rootfs");
	ignored.push_back("proc");
	ignored.push_back("sysfs");
	ignored.push_back("udev");
	ignored.push_back("devpts");
	ignored.push_back("shm");
	ignored.push_back("usbfs");
	ignored.push_back("binfmt_misc");
	ignored.push_back("fusectl");
	
	if (argc == 4) {
		config.m1 = atoi(argv[1]);
		config.m2 = atoi(argv[2]);
		if (strcmp(argv[3], "long") == 0) {
			config.long_names = true;
		} else {
			config.long_names = false;
		}
	} else {
		config.m1 = 20;
		config.m2 = 10;
		config.long_names = false;
	}
	
	get_available_space(space);
	get_memswap(space);
	get_sysload(sysload);
	
	append(output, "|HOST|", sysload.color);
	append(output, ":");
	append(output, "|TTY|", COL_DARKGREEN);
	append(output, " ");
	append(output, "|CWD|", COL_DARKGREEN);
	append(output, ":::");
	
	for (unsigned int i = 0; i < space.size(); i++) {
		append(output, space[i].name);
		// TODO Active only
		append(output, float_to_str(space[i].available, tmp), space_to_color(space[i]));
		append(output, " ");
	}
	
	append(output, "| ");
	append(output, float_to_str(sysload.load, tmp), sysload.color);
	append(output, " ");
	append(output, sysload.proc, sysload.color);
	
	std::cout << output;
	
	return 0;
}

void append(char *output, const char *input, int color) {
	unsigned int len_o, len_i, len_c = 0;
	
	len_o = strlen(output);
	len_i = strlen(input);
	if (color != -1) len_c = strlen(color_map[color].c_str());
	
	if (color != -1) strcpy(output + len_o, color_map[color].c_str());
	strcpy(output + len_o + len_c, input);
	if (color != -1) strcpy(output + len_o + len_i + len_c, color_map[COL_NONE].c_str());
}

const char *float_to_str(float value, char *output) {
	char suffix;
	
	output[0] = 0;
	
	if (value > 2*1099511627776.0) { // 2T
		value /= 1099511627776.0;
		suffix = 'T';
	} else if (value > 2*1073741824.0) { // 2G
		value /= 1073741824.0;
		suffix = 'G';
	} else if (value > 2*1048576.0) { // 2M
		value /= 1048576.0;
		suffix = 'M';
	} else if (value > 2*1024.0) { // 2k
		value /= 1024.0;
		suffix = 'k';
	} else {
		suffix = 0;
	}
	
	sprintf(output, "%.2f%c", value, suffix);
	
	return output;
}

int space_to_color(Space &space) {
	if (space.percent >= config.m1) {
		return COL_GREEN;
	} else if (space.percent >= config.m2) {
		return COL_YELLOW;
	} else if (space.percent > (config.m2/32.0)) {
		return COL_RED;
	} else {
		return COL_RFLASH;
	}
}

void get_sysload(Sysload &sysload) {
	float dump;
	std::ifstream in;
	
	in.open("/proc/loadavg", std::ifstream::in);
	
	in >> sysload.load;
	in >> dump;
	in >> dump;
	in >> sysload.proc;
	
	in.close();
	
	if (sysload.load > 12) {
		sysload.color = COL_RFLASH;
	} else if (sysload.load > 8) {
		sysload.color = COL_RED;
	} else if (sysload.load > 4) {
		sysload.color = COL_YELLOW;
	} else {
		sysload.color = COL_GREEN;
	}
}

float grepf(const std::string &data, const char *marker) {
	unsigned int pos, offset;
	char value[16];
	float out;
	
	offset = 0;
	pos = data.find(marker) + strlen(marker);
	
	// skip through spaces
	while (data[pos] == ' ') {
		pos++;
	}
	
	// read the number
	for (; data[pos] != ' '; pos++, offset++) {
		value[offset] = data[pos];
	}
	
	// terminator
	value[offset] = 0;
	
	out = atof(value);
	
	return out;
}

void get_memswap(std::vector<Space> &space) {
	std::string s;
	std::ifstream in;
	char c;
	
	in.open("/proc/meminfo", std::ifstream::in);
	
	while (true) {
		c = in.get();
		
		if (in.good()) {
			s += c;
		} else {
			break;
		}
	}
	
	in.close();
	
	// memory
	space.push_back(Space("mem", grepf(s, "MemFree:")*1024 + grepf(s, "Buffers:")*1024 + grepf(s, "Cached:")*1024, grepf(s, "MemTotal:")*1024));
	
	// swap
	space.push_back(Space("swp", grepf(s, "SwapFree:")*1024.0, grepf(s, "SwapTotal:")*1024.0));
}

void get_available_space(std::vector<Space> &space) {
	// Go through physical RW mounts and query them re: free space
	// Ignore /dev and everything below it
	// Generate a unique name: mountpoint basename? or maybe shortened to smallest unique part.
	
	std::vector<std::string> mounts;
	struct statfs stats;
	std::string fullname, name, mountpoint;
	fsblkcnt_t free_space, total_space;
	int ptr, len;
	size_t str_pos;
	bool OK;
	
	mounts = read_lines_from_file("/etc/mtab");
	
	for (unsigned int i = 0; i < mounts.size(); i++) {
		if (column_contains_substr(mounts[i], 3, "rw") and column_contains_substr(mounts[i], 0, "/dev", -1)) {
			
			// no CDs - loops have rw set
			if (column_contains_substr(mounts[i], 2, "iso9660")) {
				continue;
			}
			
			// no binds
			if (column_contains_substr(mounts[i], 3, "bind")) {
				continue;
			}
			
			
			mountpoint = get_column(mounts[i], 1);
			
			// handle spaces in mount point names - /proc/mounts uses a 4-byte '\\040' sequence to signify a space
			while (true) {
				str_pos = mountpoint.find("\\040");
				
				if (str_pos == std::string::npos) {
					break;
				} else {
					mountpoint.replace(str_pos, 4, " ");
				}
			}
			
			// generate name FIXME basenames are unique on systems with obsessive-compulsive roots only
			fullname = basename((char*)mountpoint.c_str());
			if (config.long_names or fullname == "/") {
				name = fullname;
			} else {
				ptr = 1;
				len = fullname.size();
				name = fullname.substr(0, ptr);
				
				OK = false;
				while (not OK) {
					OK = true;
					for (unsigned int j = 0; j < space.size(); j++) {
						if (name == space[j].name) {
							OK = false;
							if (ptr == len) {
								// too bad, we give up
								name += "!";
								OK = true;
							} else {
								name = fullname.substr(0, ++ptr);
							}
							
							break;
						}
					}
				}
			}
			
			// figure out free and total space
			statfs(mountpoint.c_str(), &stats);
			free_space = stats.f_bfree * stats.f_bsize;
			total_space = stats.f_blocks * stats.f_bsize;
			
			// pass it on
			space.push_back(Space(name.c_str(), free_space, total_space));
		}
	}
}

std::string get_column(std::string line, unsigned int column) {
	std::string output;
	unsigned int col_counter;
	bool eating_chars; // else eating whitespace
	char c;
	
	col_counter = -1; // Dirty? Maybe. It's called 'creative' where I come from.
	eating_chars = false; // ignore leading whitespace
	
	// start eating bytes
	for (unsigned int i = 0; i < line.size(); i++) {
		c = line[i];
		
		if (eating_chars and isspace(c)) { // end of a column
			eating_chars = false;
			
			if (col_counter == column) {
				break;
			}
		
		} else if (not (eating_chars or isspace(c))) { // beginning of a column
			eating_chars = true;
			col_counter++;
		}
		
		if (col_counter == column) {
			output += c;
		}
	}
	
	return output;
}

bool column_contains_substr(std::string line, unsigned int column, std::string substr, int position) {
	std::string str;
	
	str = get_column(line, column);
	
	switch (position) {
		case -1:
			// beginning
			if (str.substr(0, substr.size()) == substr) {
				return true;
			} else {
				return false;
			}
			
			break;
		case 0:
			// anywhere
			if (str.find(substr) == std::string::npos) {
				return false;
			} else {
				return true;
			}
			
			break;
		case 1:
			// end
			if (str.substr(str.size() - substr.size() - 1, substr.size()) == substr) {
				return true;
			} else {
				return false;
			}
			
			break;
	};
	
	return false;
}

std::vector<std::string> read_lines_from_file(std::string filename) {
	char buf[512];
	std::ifstream in;
	std::vector<std::string> lines;
	
	in.open(filename.c_str(), std::ifstream::in);
	
	while (in.good()) {
		in.getline(buf, 512);
		
		if (in.gcount()) {
			lines.push_back(buf);
		}
	}
	
	return lines;
}

bool str_in_list(std::string str, std::vector<std::string> list) {
	for (unsigned int i = 0; i < list.size(); i++) {
		if (list[i] == str) {
			return true;
		}
	}
	
	return false;
}