#pragma once

#include "common.hpp"
#include "response.hpp"
#include "request.hpp"
#include <string>
#include <unordered_map>
#include <vector>
#include <functional>
#include <string_view>
#include <memory>
#include <cstring>

using std::string_view;

using Handler = std::function<void(const Request&, Response&)>;

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>();

	// Zero-allocation path segment iterator
	struct PathIterator {
		const char* current;
		const char* end;

		PathIterator(string_view path) : current(path.data()), end(path.data() + path.size()) {
			skip_slash();
		}

		bool has_next() const { return current < end; }

		string_view next() {
			if (!has_next()) return {};

			const char* start = current;
			while (current < end && *current != '/') ++current;
			string_view segment(start, current - start);
			skip_slash();
			return segment;
		}

	private:
		void skip_slash() {
			while (current < end && *current == '/') ++current;
		}
	};

	TrieNode* insert_path(string_view path) {
		TrieNode* current = root.get();
		PathIterator it(path);

		while (it.has_next()) {
			string_view segment = it.next();
			if (segment.empty()) continue;

			bool is_param = segment[0] == ':';
			if (is_param) {
				segment = segment.substr(1); // Remove ':'
				if (!current->param_child) {
					current->param_child = std::make_unique<TrieNode>();
				}
				current->param_child->param_name = std::string(segment);
				current = current->param_child.get();
			} else {
				std::string segment_key(segment);
				auto it = current->children.find(segment_key);
				if (it == current->children.end()) {
					current->children[segment_key] = std::make_unique<TrieNode>();
				}
				current = current->children[segment_key].get();
			}
		}
		return current;
	}

	TrieNode* find_path(string_view path, std::unordered_map<string_view, string_view>& params) const {
		TrieNode* current = root.get();
		PathIterator it(path);

		while (it.has_next() && current) {
			string_view segment = it.next();
			if (segment.empty()) continue;

			// Try exact match first - zero allocation lookup
			auto child_it = current->children.end();
			for (auto iter = current->children.begin(); iter != current->children.end(); ++iter) {
				if (likely(iter->first.size() == segment.size()) &&
				    memcmp(iter->first.data(), segment.data(), segment.size()) == 0) {
					child_it = iter;
					break;
				}
			}

			if (child_it != current->children.end()) {
				current = child_it->second.get();
				continue;
			}

			// Try parameter match
			if (current->param_child) {
				// Zero-allocation parameter storage - points to original path
				string_view param_name(current->param_child->param_name);
				params[param_name] = segment;
				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<string_view, string_view>& params) const {
		TrieNode* node = find_path(path, params);
		if (unlikely(!node)) return false;

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

		out_handler = it->second;
		return true;
	}

	bool handle(Request& request, Response& response) const {
		// Clear old params without allocation
		request.params.clear();

		Handler handler;
		if (route(request.method, request.path, handler, request.params)) {
			handler(request, response);
			return true;
		}
		return false;
	}
};