cpp_server/router.hpp

#pragma once

#include "http_common.hpp"
#include "http_parser.hpp"
#include "http_response.hpp"
#include <string>
#include <unordered_map>
#include <vector>
#include <functional>
#include <string_view>
#include <memory>

using std::string_view;

using Handler = std::function<void(const HttpRequest&, HttpResponse&)>;

struct TrieNode {
	std::string prefix;
	std::string param_name; // For parameter nodes
	std::unordered_map<HttpMethod, Handler> handlers;
	std::unordered_map<char, std::unique_ptr<TrieNode>> children;
	bool is_param = false;

	TrieNode(string_view p = "") : prefix(p) {}
};

class Router {
private:
	std::unique_ptr<TrieNode> root = std::make_unique<TrieNode>();

	TrieNode* insert_path(string_view path) {
		TrieNode* current = root.get();
		size_t i = 0;

		while (i < path.length()) {
			if (path[i] == '/') {
				i++;
				continue;
			}

			bool is_param = path[i] == ':';
			if (is_param) i++; // skip ':'

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

			std::string segment(path.substr(start, i - start));

			if (is_param) {
				// All parameters share the same node
				auto it = current->children.find(':');
				if (it == current->children.end()) {
					auto new_node = std::make_unique<TrieNode>();
					new_node->is_param = true;
					new_node->param_name = segment;
					current->children[':'] = std::move(new_node);
				}
				current = current->children[':'].get();
			} else {
				char first_char = segment[0];
				auto it = current->children.find(first_char);
				if (it == current->children.end()) {
					auto new_node = std::make_unique<TrieNode>(segment);
					current->children[first_char] = std::move(new_node);
					current = current->children[first_char].get();
				} else {
					current = it->second.get();
					if (current->prefix != segment) {
						// Handle prefix mismatch - split node if needed
						auto new_node = std::make_unique<TrieNode>(segment);
						current->children[first_char] = std::move(new_node);
						current = current->children[first_char].get();
					}
				}
			}
		}
		return current;
	}

	TrieNode* find_path(string_view path, std::unordered_map<std::string, std::string>& params) const {
		TrieNode* current = root.get();
		size_t i = 0;

		while (i < path.length() && current) {
			if (path[i] == '/') {
				i++;
				continue;
			}

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

			string_view segment = path.substr(start, i - start);

			// Try exact match first
			auto it = current->children.find(segment[0]);
			if (it != current->children.end() && it->second->prefix == segment) {
				current = it->second.get();
				continue;
			}

			// Try parameter match
			auto param_it = current->children.find(':');
			if (param_it != current->children.end()) {
				current = param_it->second.get();
				params[current->param_name] = std::string(segment);
				continue;
			}

			return nullptr;
		}
		return current;
	}

public:
	void get(string_view path, Handler handler) {
		insert_path(path)->handlers[HttpMethod::GET] = std::move(handler);
	}

	void post(string_view path, Handler handler) {
		insert_path(path)->handlers[HttpMethod::POST] = std::move(handler);
	}

	void put(string_view path, Handler handler) {
		insert_path(path)->handlers[HttpMethod::PUT] = std::move(handler);
	}

	void del(string_view path, Handler handler) {
		insert_path(path)->handlers[HttpMethod::DELETE] = std::move(handler);
	}

	void patch(string_view path, Handler handler) {
		insert_path(path)->handlers[HttpMethod::PATCH] = std::move(handler);
	}

	bool route(HttpMethod method, string_view path, Handler& out_handler, std::unordered_map<std::string, std::string>& params) const {
		TrieNode* node = find_path(path, params);
		if (!node) return false;

		auto it = node->handlers.find(method);
		if (it == node->handlers.end()) return false;

		out_handler = it->second;
		return true;
	}

	bool handle(HttpRequest& request, HttpResponse& response) const {
		Handler handler;
		if (route(request.method, request.path, handler, request.params)) {
			handler(request, response);
			return true;
		}
		return false;
	}
};

// Usage example:
/*
Router router;

router.get("/users", [](const HttpRequest& req, HttpResponse& res) {
	res.set_json("{\"users\": []}");
});

router.get("/users/:id", [](const HttpRequest& req, HttpResponse& res) {
	std::string id = req.params.at("id");
	res.set_json("{\"id\": \"" + id + "\"}");
});

router.post("/users", [](const HttpRequest& req, HttpResponse& res) {
	// Use req.body to get POST data
	res.status = 201;
	res.set_json("{\"created\": true}");
});

router.get("/users/:id/posts/:postId", [](const HttpRequest& req, HttpResponse& res) {
	std::string userId = req.params.at("id");
	std::string postId = req.params.at("postId");
	res.set_json("{\"userId\": \"" + userId + "\", \"postId\": \"" + postId + "\"}");
});
*/