Goldfish/goldfish.hpp

#pragma once

#include <cstdint>
#include <string_view>
#include <string>
#include <vector>
#include <functional>
#include <memory>
#include <cstring>

namespace goldfish
{

using namespace std;

enum class HttpMethod : uint8_t
{
	GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS, UNKNOWN
};

class MimeTypes
{
public:
	static string_view get_mime_type(string_view extension)
	{
		if (!extension.empty() && extension[0] == '.') {
			extension = extension.substr(1);
		}

		switch (extension.size()) {
			case 2:
				if (extension == "js") return "application/javascript";
				break;
			case 3:
				if (extension == "css") return "text/css";
				if (extension == "htm") return "text/html";
				if (extension == "png") return "image/png";
				if (extension == "jpg") return "image/jpeg";
				if (extension == "gif") return "image/gif";
				if (extension == "svg") return "image/svg+xml";
				if (extension == "ico") return "image/x-icon";
				if (extension == "ttf") return "font/ttf";
				if (extension == "txt") return "text/plain";
				if (extension == "pdf") return "application/pdf";
				break;
			case 4:
				if (extension == "html") return "text/html";
				if (extension == "json") return "application/json";
				if (extension == "jpeg") return "image/jpeg";
				if (extension == "woff") return "font/woff";
				break;
			case 5:
				if (extension == "woff2") return "font/woff2";
				break;
		}
		return "application/octet-stream";
	}

	static bool should_compress(string_view mime_type)
	{
		return mime_type.starts_with("text/") ||
		       mime_type.starts_with("application/json") ||
		       mime_type.starts_with("application/javascript") ||
		       mime_type.starts_with("image/svg+xml");
	}
};

class IRequest
{
public:
	virtual ~IRequest() = default;
	virtual HttpMethod method() const = 0;
	virtual string_view path() const = 0;
	virtual string_view query() const = 0;
	virtual string_view body() const = 0;
	virtual string_view get_header(string_view name) const = 0;
};

class IResponse
{
public:
	virtual ~IResponse() = default;
	virtual void set_status(int status) = 0;
	virtual void set_header(string_view name, string_view value) = 0;
	virtual void set_body(string_view body) = 0;
	virtual void set_content_type(string_view content_type) = 0;
};

using Params = const vector<string_view>&;
using Handler = function<void(IRequest&, IResponse&, Params)>;

struct Node
{
	string segment;
	Handler handler;
	vector<unique_ptr<Node>> children;
	bool is_param = false;
	uint8_t max_params = 0;
};

class Router
{
private:
	unique_ptr<Node> get_root = make_unique<Node>();
	unique_ptr<Node> post_root = make_unique<Node>();
	unique_ptr<Node> put_root = make_unique<Node>();
	unique_ptr<Node> delete_root = make_unique<Node>();
	unique_ptr<Node> patch_root = make_unique<Node>();

	mutable vector<string_view> params_buffer;

	Node* method_node(HttpMethod method) const
	{
		switch (method) {
			case HttpMethod::GET: return get_root.get();
			case HttpMethod::POST: return post_root.get();
			case HttpMethod::PUT: return put_root.get();
			case HttpMethod::DELETE: return delete_root.get();
			case HttpMethod::PATCH: return patch_root.get();
			default: return nullptr;
		}
	}

	struct Segment
	{
		string_view text;
		bool has_more;
	};

	Segment read_segment(string_view path, size_t& pos) const
	{
		if (pos >= path.size()) return {"", false};

		// Skip leading slashes
		while (pos < path.size() && path[pos] == '/') ++pos;
		if (pos >= path.size()) return {"", false};

		size_t start = pos;
		while (pos < path.size() && path[pos] != '/') ++pos;

		return {path.substr(start, pos - start), pos < path.size()};
	}

	void add_route(Node* root, string_view path, Handler handler)
	{
		if (path == "/") {
			root->handler = move(handler);
			return;
		}

		Node* current = root;
		size_t pos = 0;
		uint8_t param_count = 0;

		while (true) {
			auto [segment, has_more] = read_segment(path, pos);
			if (segment.empty()) break;

			bool is_param = !segment.empty() && segment[0] == ':';
			if (is_param) {
				segment = segment.substr(1);
				param_count++;
			}

			Node* child = nullptr;
			for (auto& c : current->children) {
				if (c->segment == segment) {
					child = c.get();
					break;
				}
			}

			if (!child) {
				auto new_child = make_unique<Node>();
				new_child->segment = string(segment);
				new_child->is_param = is_param;
				child = new_child.get();
				current->children.push_back(move(new_child));
			}

			if (child->max_params < param_count) {
				child->max_params = param_count;
			}

			current = child;
			if (!has_more) break;
		}

		current->handler = move(handler);
	}

	pair<Handler, vector<string_view>> lookup(Node* root, string_view path) const
	{
		if (path == "/") {
			return {root->handler, {}};
		}

		// Resize buffer if needed
		if (params_buffer.size() < root->max_params) {
			params_buffer.resize(root->max_params);
		}
		params_buffer.clear();

		auto result = match(root, path, 0);
		if (!result.first) return {nullptr, {}};

		return {result.first, vector<string_view>(params_buffer.begin(), params_buffer.begin() + result.second)};
	}

	pair<Handler, int> match(Node* current, string_view path, size_t start) const
	{
		size_t pos = start;
		auto [segment, has_more] = read_segment(path, pos);

		if (segment.empty()) {
			return {current->handler, 0};
		}

		for (auto& child : current->children) {
			if (child->segment == segment || child->is_param) {
				int param_count = 0;
				if (child->is_param) {
					params_buffer.push_back(segment);
					param_count = 1;
				}

				if (!has_more) {
					return {child->handler, param_count};
				}

				auto [handler, nested_count] = match(child.get(), path, pos);
				if (handler) {
					return {handler, param_count + nested_count};
				}
			}
		}

		return {nullptr, 0};
	}

public:
	Router()
	{
		params_buffer.reserve(16);
	}

	void get(string_view path, Handler handler)
	{
		add_route(get_root.get(), path, move(handler));
	}

	void post(string_view path, Handler handler)
	{
		add_route(post_root.get(), path, move(handler));
	}

	void put(string_view path, Handler handler)
	{
		add_route(put_root.get(), path, move(handler));
	}

	void del(string_view path, Handler handler)
	{
		add_route(delete_root.get(), path, move(handler));
	}

	void patch(string_view path, Handler handler)
	{
		add_route(patch_root.get(), path, move(handler));
	}

	bool handle(IRequest& request, IResponse& response) const
	{
		Node* root = method_node(request.method());
		if (!root) return false;

		auto [handler, params] = lookup(root, request.path());
		if (!handler) return false;

		handler(request, response, params);
		return true;
	}
};

} // namespace goldfish