Nextus-Web/src/main/java/dev/coph/nextusweb/server/websocket/WebSocketRouter.java

package dev.coph.nextusweb.server.websocket;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * A trie-based router that maps WebSocket upgrade paths to {@link WebSocketHandler}s.
 *
 * <p>It mirrors the HTTP {@link dev.coph.nextusweb.server.router.Router Router} but is simpler:
 * a path resolves to a single handler (there is no HTTP method dimension) and only static and
 * {@code {param}} path-parameter segments are supported (no wildcards). Registration mutates the
 * shared trie at start-up; {@link #resolve(String)} is safe to call concurrently afterwards.</p>
 */
public final class WebSocketRouter {

    /** Root of the routing trie. */
    private final Node root = new Node();

    /**
     * Creates an empty WebSocket router with no registered handlers.
     */
    public WebSocketRouter() {
    }

    /**
     * Registers a handler at the given path, creating any missing trie nodes. Segments wrapped
     * in braces (e.g. {@code /chat/{room}}) are treated as path parameters.
     *
     * @param path    the WebSocket path to mount the handler at
     * @param handler the handler to invoke for connections on that path
     * @return this router, for fluent chaining
     */
    public WebSocketRouter on(String path, WebSocketHandler handler) {
        Node node = root;
        for (String segment : split(path)) {
            if (segment.startsWith("{") && segment.endsWith("}")) {
                if (node.paramChild == null) {
                    node.paramChild = new Node();
                    node.paramName = segment.substring(1, segment.length() - 1);
                }
                node = node.paramChild;
            } else {
                node = node.children.computeIfAbsent(segment, k -> new Node());
            }
        }
        node.handler = handler;
        return this;
    }

    /**
     * Resolves a path to its handler, capturing any path parameters along the way.
     *
     * @param path the request path
     * @return a {@link Resolution} carrying the handler and captured parameters, or {@code null}
     *         if no handler is registered for the path
     */
    public Resolution resolve(String path) {
        Map<String, String> params = new HashMap<>(4);
        Node node = root;
        for (String segment : split(path)) {
            Node next = node.children.get(segment);
            if (next != null) {
                node = next;
            } else if (node.paramChild != null) {
                params.put(node.paramName, segment);
                node = node.paramChild;
            } else {
                return null;
            }
        }
        if (node.handler == null) return null;
        return new Resolution(node.handler, params);
    }

    /**
     * Splits a path into its non-empty segments, ignoring leading and collapsing internal
     * slashes.
     *
     * @param path the raw path
     * @return the ordered list of path segments
     */
    private static List<String> split(String path) {
        List<String> out = new ArrayList<>();
        int start = path.startsWith("/") ? 1 : 0;
        for (int i = start; i < path.length(); i++) {
            if (path.charAt(i) == '/') {
                if (i > start) out.add(path.substring(start, i));
                start = i + 1;
            }
        }
        if (start < path.length()) out.add(path.substring(start));
        return out;
    }

    /**
     * A successful path resolution.
     *
     * @param handler    the handler bound to the matched path
     * @param pathParams the path parameters captured while matching, keyed by name
     */
    public record Resolution(WebSocketHandler handler, Map<String, String> pathParams) {
    }

    /**
     * A single node in the WebSocket routing trie. Holds static children keyed by segment, an
     * optional path-parameter child, and the handler (if any) registered at this node.
     */
    private static final class Node {
        /** Static child nodes keyed by their literal path segment. */
        final Map<String, Node> children = new ConcurrentHashMap<>();
        /** Child matching any single segment as a path parameter, or {@code null} if none. */
        Node paramChild;
        /** Name under which {@link #paramChild} captures the matched segment. */
        String paramName;
        /** Handler registered at this node, or {@code null} if the path is only a prefix. */
        WebSocketHandler handler;
    }
}