#include <netinet/in.h>
#include <stddef.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <threads.h>
#include <unistd.h>
#include <netdb.h>
#include <assert.h>

// Encryption
#ifdef USE_SSL
#include <openssl/bio.h>
#include <openssl/ssl.h>
#include <openssl/err.h>
#endif

const int WRAP_WIDTH = 70;

char* USER_AGENT = "User-Agent: TinyBrowser/1.0\r\n";

int linklist_append(char* address);

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// Generic string handeling                                                  //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

char* box_substring(char* orig, int length) {
	char* new = malloc(length + 1);
	strncpy(new, orig, length + 1);
	new[length] = 0;
	return new;
}

char* box_string(char* orig) {
	int size = strlen(orig);
	return box_substring(orig, size);
}

int is_whitespace(char c) {
	switch (c) {
		case ' ': return 1;
		case '\n': return 1;
		case '\r': return 1;
		case '\t': return 1;
		default: return 0;
	}
}

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// URL Parsing                                                               //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

enum proto {
	PROTO_HTTP,
	PROTO_HTTPS,
};

struct URL {
	// Metadata
	int valid;
	char* err;
	// Parsed information
	int proto;
	char* host;
	int port;
	char* path;
	// Network data
	int content_length;
	int size;
	int status_code;
	int is_text;
};

void url_print(struct URL* url) {
	if (!url->valid) {
		printf("(invalid: %s)", url->err);
		return;
	}

	int default_port;
	switch (url->proto) {
		case PROTO_HTTP: printf("http://"); default_port = 80; break;
		case PROTO_HTTPS: printf("https://"); default_port = 443; break;
	}

	printf("%s", url->host);
	if (url->port != default_port) printf(":%d", url->port);
	printf("%s", url->path);
}

void url_free(struct URL url) {
	if (url.host) free(url.host);
	if (url.path) free(url.path);
}


char* path_append(char* base, char* new) {
	// Trival case: absolute path and empty path
	if (new[0] == '/') return box_string(new);
	if (!new[0]) box_string(base);

	// Find the index of the last slash in the base path
	int i = 0;
	int dirlen = 0;
	for (; base[i]; i++) {
		if (base[i] == '/') dirlen = i + 1;
	}

	// Copy up to and including that slash...
	char* result = box_substring(base, dirlen);
	// ... and append the new file
	result = realloc(result, dirlen + strlen(new) + 2);
	strcat(result, new);
	return result;
}

struct URL url_parse(struct URL* base, char* string) {
	// Create copy of existing URL
	struct URL url = *base;
	if (url.host) url.host = box_string(url.host);
	if (url.path) url.path = box_string(url.path);

	// Everything is ok for now
	url.valid = 1;
	url.err = NULL;

	// Check if the url contains a scheme
	char* scheme = strchr(string, ':');
	if (scheme) {
		if (strncmp(string, "http:", 5) == 0) {url.proto = PROTO_HTTP; url.port = 80;}
		#ifdef USE_SSL
		else if (strncmp(string, "https:", 6) == 0) {url.proto = PROTO_HTTPS; url.port = 443;}
		#endif
		else {
			url.err = "Unsupported protocol";
			url.valid = 0;
			return url;
		}
		string = scheme + 1;
	}

	// Check if it contains an address
	if (strncmp(string, "//", 2) == 0) {
		string += 2;
		// Location extends to the end of the string, or to a /
		int len = strlen(string);
		char* slash = strchr(string, '/');
		if (slash) len = slash - string;
		char* netloc = box_substring(string, len);
		string += len ;
		char* delim = strchr(netloc, ':');
		if (!delim) {
			if (url.host) free(url.host);
			url.host = netloc;
		} else {
			if (url.host) free(url.host);
			url.host = box_substring(netloc, delim - netloc);
			url.port = atoi(delim + 1);
			free(netloc);
		}
		// Reset most of url if address was specified
		if (url.path) free(url.path);
		url.path = box_string("/");
	}

	if (strlen(string)) {
		// Truncate at fragment ID
		char* hash = strchr(string, '#');
		if (hash) *hash = 0;

		char* old_path = url.path;
		url.path = path_append(url.path, string);
		free(old_path);
		// Put the hash back to avoid breaking things
		if (hash) *hash = '#';
	}

	return url;
}

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// Network                                                                   //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

void send_str(int fd, char* string) {
	send(fd, string, strlen(string), 0);
}

// Send a plaintext GET request. Does not parse response
char* raw_get_plain(struct URL* address) {
	int fd, status;

	// Resolve address
	struct addrinfo hints, *results, *rp;
	memset(&hints, 0, sizeof(hints));
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;

	int s = getaddrinfo(address->host, NULL, &hints, &results);
	if (s != 0) {
		printf("Resolve error: %s\n", gai_strerror(s));
		address->err = "Could not resolve hostname";
		address->valid = 0;
		return NULL;
	}

	// Try addresses until one works
	fd = -1;
	for (rp = results; rp != NULL; rp = rp->ai_next) {
		char ipstr[INET6_ADDRSTRLEN];
		struct sockaddr* serv_addr;
		serv_addr = rp->ai_addr;
		int type = rp->ai_addr->sa_family;

		// Set port
		if (type == AF_INET) {
			struct sockaddr_in* addr = (struct sockaddr_in*) serv_addr;
			addr->sin_port = htons(address->port);
			inet_ntop(type, &addr->sin_addr, ipstr, sizeof(ipstr));
		} else if (type == AF_INET6) {
			struct sockaddr_in6* addr = (struct sockaddr_in6*) serv_addr;
			addr->sin6_port = htons(address->port);
			inet_ntop(type, &addr->sin6_addr, ipstr, sizeof(ipstr));
		} else continue;

		printf("Trying IP %s...\n", ipstr);

		if ((fd = socket(type, SOCK_STREAM, 0)) < 0) {
			printf("Cound not create socket!\n");
			exit(1);
		}

		if ((status = connect(fd, serv_addr, rp->ai_addrlen) != 0)) {
			perror("Failed to connect.");
			close(fd);
			fd = -1;
		} else break;
 	}

	if (fd == -1) {
		address->valid = 0;
		address->err = "No addresses worked. :(";
		return NULL;
	}

	printf("Connected!\n");
	freeaddrinfo(results);

	// Send request
	send_str(fd, "GET ");
	send_str(fd, address->path);
	send_str(fd, " HTTP/1.0\r\n");
	send_str(fd, USER_AGENT);
	send_str(fd, "Host: ");
	send_str(fd, address->host);
	send_str(fd, "\r\n\r\n");

	// Read data
	char* data = box_string("");
	int size = 0;
	while (1) {
		data = realloc(data, size + 1025);
		int read = recv(fd, &data[size], 1024, 0);
		size += read;
		if (read <= 0) break;
	}

	close(fd);

	// Null terminate (ok because we allocate one byte too many while reading)
	data[size] = 0;
	address->size = size;
	return data;
}

#ifdef USE_SSL
void ssl_send(SSL* ssl, char* str) {
	size_t t;
	SSL_write_ex(ssl, str, strlen(str), &t);
}

// Send an encrypted GET request. Does not parse the response
char* raw_get_ssl(struct URL* address) {
	SSL_CTX* ctx = SSL_CTX_new(TLS_client_method());
	assert(ctx);
	// Fail on bad ssl
	SSL_CTX_set_verify(ctx, SSL_VERIFY_PEER, NULL);
	// Use default trusted certs
	assert(SSL_CTX_set_default_verify_paths(ctx));
	// Don't use broken crypto
	assert(SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION));

	int sock = -1;
	BIO_ADDRINFO *res;
	const BIO_ADDRINFO *ai = NULL;

	// Resolve hostname
	char port[10];
	snprintf(port, 10, "%d", address->port); // Why is this a string?!
	int status;
	status = BIO_lookup_ex(address->host, port, BIO_LOOKUP_CLIENT, AF_UNSPEC, SOCK_STREAM, 0, &res);
	if (!status) {
		address->valid = 0;
		address->err = "Failed to resolve hostname (via openssl)";
		return 0;
	}

	// Attempt to connect to all adresses until one works
	for (ai = res; ai != NULL; ai = BIO_ADDRINFO_next(ai)) {
		sock = BIO_socket(BIO_ADDRINFO_family(ai), SOCK_STREAM, 0, 0);
		if (sock == -1) continue;

		if (!BIO_connect(sock, BIO_ADDRINFO_address(ai), BIO_SOCK_NODELAY)) {
			BIO_closesocket(sock);
			sock = -1;
			continue;
		}
		break;
	}

	// Error out if none worked
	if (sock == -1) {
		address->valid = 0;
		address->err = "No addresses worked (via openssl)";
		return 0;
	}

	// We don't need the addresses anymore
	BIO_ADDRINFO_free(res);

	BIO *bio;

	// Wrap the socket in openssl.
	bio = BIO_new(BIO_s_socket());
	if (bio == NULL) {
		BIO_closesocket(sock);
		return NULL;
	}
	BIO_set_fd(bio, sock, BIO_CLOSE);

	// Assign ssl object
	SSL* ssl = SSL_new(ctx);
	SSL_set_bio(ssl, bio, bio);

	// Set server name (non-critcal)
	SSL_set_tlsext_host_name(ssl, address->host);

	// Connect!
	if (!SSL_set1_host(ssl, address->host)) {
		address->valid = 0;
		address->err = "Wrong SSL hostname";
		return 0;
	}

	if (SSL_connect(ssl) < 1) {
		address->valid = 0;
		address->err = "failed to start SSL";
		return 0;
	}

	// Send request
	ssl_send(ssl, "GET ");
	ssl_send(ssl, address->path);
	ssl_send(ssl, " HTTP/1.0\r\n");
	ssl_send(ssl, USER_AGENT);
	ssl_send(ssl, "Host: ");
	ssl_send(ssl, address->host);
	ssl_send(ssl, "\r\n\r\n");

	// Read data
	char* data = box_string("");
	int size = 0;
	while (1) {
		data = realloc(data, size + 1025);
		size_t read = 0;
		status = SSL_read_ex(ssl, &data[size], 1024, &read);
		size += read;
		if (read <= 0) break;
	}

	SSL_shutdown(ssl);
	SSL_free(ssl);
	SSL_CTX_free(ctx);

	// Null terminate (ok because we allocate one byte too many while reading)
	data[size] = 0;
	address->size = size;
	return data;
}
#endif

// Attempts to download a given url
// Returns bytes on heap if sucessful,
uint8_t* url_get(struct URL* address) {
	char* data;
	if (address->proto == PROTO_HTTP) {
		data = raw_get_plain(address);
	} else if (address->proto == PROTO_HTTPS) {
		#ifdef USE_SSL
		data = raw_get_ssl(address);
		#else
		data = 0;
		#endif
	}
	if (data == 0) return 0;

	// Header parser
	char* line = data;
	if (strncmp(line, "HTTP/", 5) != 0) {
		bad_status:
		free(data);
		address->valid = 0;
		address->err = "Response is not HTTP";
		return 0;
	}

	char* code = strchr(line, ' ');
	if (code) {
		address->status_code = atoi(code + 1);
		printf("Status code is: %d\n", address->status_code);
	} else {
		goto bad_status;
	}

	line = strchr(line, '\n') + 1;
	if (!line) {
		bad_headers:
		address->valid = 0;
		address->err = "Bad headers";
		free(data);
		return 0;
	}

	char* location = NULL;
	address->content_length = -1;
	address->is_text = 1;
	while (1) {
		int is_header = strncmp("\r\n", line, 2) && strncmp("\n", line, 1);
		if (!is_header) {
			line = strchr(line, '\n') + 1;
			break;
		}

		if (strncmp(line, "Content-Length: ", 16) == 0) {
			address->content_length = atoi(line + 16);
		} else if (strncmp(line, "Content-Type: ", 12) == 0) {
			address->is_text = 0;
			if (strncmp(line+14, "text/plain", 10) == 0) address->is_text = 1;
			if (strncmp(line+14, "text/html", 9) == 0) address->is_text = 1;
		} else if (strncmp(line, "Location: ", 10) == 0) {
			line += 10;
			char* end = strchr(line, '\n');
			char* end2 = strchr(line, '\r');
			if (end2 < end) end = end2;
			if (!end) {
				goto bad_headers;
			}
			location = box_substring(line, end - line);
		}
		line = strchr(line, '\n') + 1;
	}

	if (location) {
		uint8_t* data;
		struct URL new = url_parse(address, location);
		printf("Using ");
		url_print(&new);
		printf("...\n");
		if (new.valid) {
			struct URL old = *address;
			*address = new;
			data = url_get(address);
			url_free(old);
		} else {
			address->err = new.err;
			address->valid = 0;
			free(location);
			return 0;
		}
		free(location);
		return data;
	}

	if (address->content_length == -1) address->content_length = strlen(line);
	uint8_t* file = malloc(address->content_length + 1);
	memcpy(file, line, address->content_length);
	free(data);
	file[address->content_length] = 0;
	return file;
}

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// HTML                                                                      //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

#define NUM_ENT 12
char* entities[NUM_ENT][2] = {
	{"amp;", "&"},
	{"gt;", ">"},
	{"lt;", "<"},
	{"shy;", ""},
	{"nbsp;", " "},
	{"mdash;", "---"},
	{"ndash;", "--"},
	{"rsquo;", "'"},
	{"lsquo;", "'"},
	{"rdquo;", "\""},
	{"ldquo;", "\""},
	{"mdlr;", "..."}
};

void bad_html(char* err) {
	printf("\n######### [BAD HTML] ##########\n");
	printf("%s\n", err);
	printf("######### [BAD HTML] ##########\n");
}

// Parses a single HTML element.
// Takes a pointer to right *after* the opening tag, and returns
// a pointer to right after the closing tag.
//
// State for collapsing paragraphs and whitespace
int was_whitespace;
int is_new_para;

// Word wrapping
int current_line_length;

// Formatting: incremented at an opening tag, decremented at closing.
int fmt_link = 0;
int fmt_italics = 0;
int fmt_bold = 0;
int fmt_pre = 0;

void html_reset() {
	fmt_link = 0;
	fmt_italics = 0;
	fmt_bold = 0;
	fmt_pre = 0;
	current_line_length = 0;

	is_new_para = 1;
	was_whitespace = 1;
}

void format() {
	printf("\x1b[0m");
	if (fmt_pre) printf("\x1b[2m");
	if (fmt_italics) printf("\x1b[33m");
	if (fmt_bold) printf("\x1b[1m");
	if (fmt_link) printf("\x1b[4;36m");
}

// These tags don't contain anything
char* NULLTAG[] = {
	"link", "meta", "br", "link", "img",
	"circle", "path", "use", "hr"
};

int is_self_closing(char* name) {
	for (int i = 0; i < sizeof(NULLTAG)/sizeof(char*); i++) {
		if (strcmp(NULLTAG[i], name) == 0) return 1;
	}
	return 0;
}

// These should start a new paragraph
char* BLOCK[] = {
	"img", "p", "legend", "fieldset",
	"div", "ul", "ol", "h1", "h2", "h3",
	"h4", "h5", "pre", "hr"
};

int is_block(char* name) {
	for (int i = 0; i < sizeof(BLOCK)/sizeof(char*); i++) {
		if (strcmp(BLOCK[i], name) == 0) return 1;
	}
	return 0;
}

char* read_attr_value(char* html, char** result) {
	int len = 0;
	char* value = NULL;
	if (*html == '"' || *html == '\'') {
		// Quoted value
		char quote = *html;
		html++;
		while (*html != quote && *html) {
			len++;
			value = realloc(value, len + 1);
			if (*html == '\\') {
				html ++;
			}
			value[len - 1] = *html;
			html ++;
		}
		if (*html) html ++;
	} else {
		// Unquoted value
		while (!is_whitespace(*html) && *html != '>' && *html != '=' && *html) {
			len++;
			value = realloc(value, len + 1);
			value[len - 1] = *html;
			html ++;
		}
	}

	if (value) value[len] = 0;
	*result = value;
	return html;
}

char* html_display(char* html, char* tag, int depth) {
	int is_visible = 1;
	int is_html = 1;
	if (strcmp("script", tag) == 0 || strcmp("style", tag) == 0) {
		is_visible = 0;
		is_html = 0;
	}
	if (strcmp("title", tag) == 0) is_visible = 0;
	if (strcmp("head", tag) == 0) is_visible = 0;

	while (1) {
		switch (*html) {
			// End of string
			case 0: return html; break;
			// Entities
			case '&':
				was_whitespace = 0;
				html++;
				// Don't attempt to resolve entities in CSS and JS
				if (is_visible) {
					int found = 0;
					for (int i = 0; i < NUM_ENT; i++) {
						char* name = entities[i][0];
						if (strncmp(html, name, strlen(name)) == 0) {
							current_line_length += strlen(entities[i][1]);
							printf("%s", entities[i][1]);
							html += strlen(name);
							found = 1;
							break;
						}
					}
					if (!found) printf("&");
					is_new_para = 0;
					break;
				}
			// Tags!
			case '<':
				html++;

				char* endoftag = strchr(html, '>');
				if (!endoftag) {
					bad_html("No right bracket for closing tag!");
					return html;
				}
				if (*html == '/') {
					// Stricter rules for javascript and stylesheets (for substack)
					if (!is_html) {
						int real_closing = 0;
						if (strncmp("/script", html, 7) == 0) real_closing = 1;
						if (strncmp("/style", html, 6) == 0) real_closing = 1;
						if (!real_closing) break;
					}
					// Closing tag. Tags are always closed in reverse order regardless of name.
					html = endoftag + 1;
					// ... ensure the's a blank line after block elements
					if (is_block(tag)) {
						if (!is_new_para) {
							printf("\n\n");
							was_whitespace = 1;
							is_new_para = 1;
							current_line_length = 0;
						}
				}
					return html;
				} else if (*html == '!') {
					// For comments, read until the closing delimiter
					if (strncmp("!--", html, 3) == 0) {
						while (strncmp("-->", html, 3) != 0 && *html) html++;
						if (strncmp("-->", html, 3) == 0) html += 3;
						break;
					} else {
						// Doctype:
						html = endoftag + 1;
						break;
					}
				}
				// Still here? It's a opening tag.

				// Don't allow tags in <script> and <style>
				if (!is_html) break;

				// Search for the end of a tags name.
				int namelen = 0;
				while (1) {
						if (html[namelen] == '>') break;
						if (html[namelen] == '/') break;
						if (is_whitespace(html[namelen])) break;
						namelen++;
				}

				// Read tag name into heap
				char* name = box_substring(html, namelen);

				// Read attributes
				char* link_address = NULL;
				html += namelen;
				while (1) {
					char* value = NULL;
					char* key = NULL;
					// Skip whitespace
					while (is_whitespace(*html)) html++;
					if (*html == 0 || *html == '/' || *html == '>') break;
					// Read key
					html = read_attr_value(html, &key);
					while (is_whitespace(*html)) html++;
					// Read value
					if (*html != '=') {free(key); continue;} else html++;
					while (is_whitespace(*html)) html++;
					html = read_attr_value(html, &value);
					while (is_whitespace(*html)) html++;
					// Extract hrefs
					if (strcmp(key, "href") == 0) {
						link_address = value;
						free(key);
					} else {
						free(key); free(value);
					}
				}

				// handle <br> tags
				if (strcmp("br", name) == 0) {
					was_whitespace = 1;
					current_line_length = 0;
					printf("\n");
				}

				// handle <hr> tags
				if (strcmp("hr", name) == 0) {
					for (int i = 0; i < WRAP_WIDTH; i++) putc('-', stdout);
					printf("\n");
					was_whitespace = 1;
					current_line_length = 0;
				}

				// Ensure there's a blank line before block elements
				if (is_block(name)) {
					if (!is_new_para) {
						printf("\n\n");
						is_new_para = 1;
						was_whitespace = 1;
						current_line_length = 0;
					}
				}

				if (link_address) fmt_link++;
				if (strcmp(name, "em") == 0 || strcmp(name, "i") == 0) fmt_italics++;
				if (strcmp(name, "b") == 0) fmt_bold++;
				if (strcmp(name, "pre") == 0) fmt_pre++;
				format();

				// ... and recurse to handle what's inside.
				html = endoftag + 1;
				if (!is_self_closing(name))
					html = html_display(html, name, depth + 1);

				if (link_address) fmt_link--;
				if (strcmp(name, "em") == 0 || strcmp(name, "i") == 0) fmt_italics--;
				if (strcmp(name, "b") == 0) fmt_bold--;
				if (strcmp(name, "pre") == 0) fmt_pre--;
				format();

				if (link_address) {
					int link_no = linklist_append(link_address);
					printf("[%d]", link_no);
				}
				free(name);

				break;

			// Everything else: print it.
			default: if (is_visible) {
				char c = *html;
				// Code blocks: No collapsing
				if (fmt_pre) {
					putc(*html, stdout);
					html++;
					is_new_para = 0;
					was_whitespace = 0;
					break;
				}
				// Whitespace collasping and word wrap
				if (is_whitespace(c)) {
					if (!was_whitespace) {
						if (current_line_length > WRAP_WIDTH) {
							putc('\n', stdout);
							current_line_length = 0;
						} else {
							current_line_length++;
							putc(' ', stdout);
						}
						was_whitespace = 1;
					}
				} else {
					was_whitespace = 0;
					is_new_para = 0;
					current_line_length++;
					putc(*html, stdout);
				}
			} html++; break;
		}
	}
}

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// UI                                                                        //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

int link_count = 0;
char** link_list = NULL;

int linklist_append(char* link) {
	link_count += 1;
	link_list = realloc(link_list, sizeof(char*) * link_count);
	link_list[link_count - 1] = link;
	return link_count;
}

int main(int argc, char** argv) {
	char* url;
	#ifdef USE_SSL
	if (argc == 1) url = "https://maurycyz.com/real_pages/";
	#else
	if (argc == 1) url = "http://maurycyz.com/real_pages/";
	#endif
	else if (argc == 2) url = argv[1];
	else {
		printf("usage: %s [url]\n", argv[0]);
		return 1;
	}

	struct URL current_page = {
		.proto = PROTO_HTTP,
		.port = 80,
		.host = NULL,
		.path = box_string("/"),
	};
	current_page = url_parse(&current_page, url);

	while (1) {
		// Fetch URL
		printf("Connecting to "); url_print(&current_page); printf(" ...\n");

		if (!current_page.valid) {
			printf("Invalid URL!: %s\n", current_page.err);
			return 1;
		}

		uint8_t* document = url_get(&current_page);

		if (!current_page.valid) {
			printf("GET failed!: %s\n", current_page.err);
			return 1;
		}

		if (!current_page.is_text) {
			printf("Content does not apear to be HTML or text. Refusing to display.\n");
			return 1;
		}

		// Display and extract links
		for (int i = 0; i < link_count; i++) free(link_list[i]);
		if (link_list) free(link_list);
		link_count = 0;
		link_list = NULL;

		printf("######\n\n");
		char* html = (char*)document;
		while (*html) {
			html_reset();
			html = html_display(html, ":root", 0);
		}
		printf("######\n\n");

		// List all the links
		if (link_count == 0) {
			printf("\n\nDead end page :(\n");
			return 0;
		}

		for (int i = 0; i < link_count; i++) {
			printf("[%d]: ", i+1);
			struct URL dest = url_parse(&current_page, link_list[i]);
			url_print(&dest);
			url_free(dest);
			printf("\n");
		}

		// Ask the user
		int selection = 0;
		while (!selection) {
			printf("What link do you want to follow? ");
			char input[128];
			fgets(input, 128, stdin);
			selection = atoi(input);
			if (selection > link_count) selection = 0;
		}

		// Navigate to the selected page
		char* path = link_list[selection - 1];
		struct URL next = url_parse(&current_page, path);
		url_free(current_page);
		current_page = next;
	}

	url_free(current_page);
	return 0;
}