#pragma once

#include "common.hpp"
#include "http_response.hpp"
#include "http_request.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 param_name;
	std::unordered_map<HttpMethod, Handler> handlers;
	std::unordered_map<std::string, std::unique_ptr<TrieNode>> children;
	std::unique_ptr<TrieNode> param_child;

	TrieNode() = default;
};

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++;

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

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

			if (is_param) {
				if (!current->param_child) {
					current->param_child = std::make_unique<TrieNode>();
				}
				current->param_child->param_name = segment;
				current = current->param_child.get();
			} else {
				auto it = current->children.find(segment);
				if (it == current->children.end()) {
					current->children[segment] = std::make_unique<TrieNode>();
				}
				current = current->children[segment].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);
			std::string segment_str(segment);

			// Try exact match first
			auto it = current->children.find(segment_str);
			if (it != current->children.end()) {
				current = it->second.get();
				continue;
			}

			// Try parameter match
			if (current->param_child) {
				params[current->param_child->param_name] = segment_str;
				current = current->param_child.get();
				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;
	}
};