package io.modelcontextprotocol.conformance.server;

import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import io.modelcontextprotocol.common.McpTransportContext;
import io.modelcontextprotocol.server.CompletionsRepository;
import io.modelcontextprotocol.server.McpAsyncServer;
import io.modelcontextprotocol.server.McpAsyncServerExchange;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpStatelessAsyncServer;
import io.modelcontextprotocol.server.McpStatelessServerFeatures;
import io.modelcontextprotocol.server.PromptsRepository;
import io.modelcontextprotocol.server.ResourcesRepository;
import io.modelcontextprotocol.server.StatelessCompletionsRepository;
import io.modelcontextprotocol.server.StatelessPromptsRepository;
import io.modelcontextprotocol.server.StatelessResourcesRepository;
import io.modelcontextprotocol.server.StatelessToolsRepository;
import io.modelcontextprotocol.server.ToolsRepository;
import io.modelcontextprotocol.server.transport.HttpServletStatelessServerTransport;
import io.modelcontextprotocol.server.transport.HttpServletStreamableServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
import io.modelcontextprotocol.spec.McpSchema.CompleteRequest;
import io.modelcontextprotocol.spec.McpSchema.CompleteResult;
import io.modelcontextprotocol.spec.McpSchema.GetPromptResult;
import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult;
import io.modelcontextprotocol.spec.McpSchema.ListResourceTemplatesResult;
import io.modelcontextprotocol.spec.McpSchema.ListResourcesResult;
import io.modelcontextprotocol.spec.McpSchema.ListToolsResult;
import io.modelcontextprotocol.spec.McpSchema.Prompt;
import io.modelcontextprotocol.spec.McpSchema.PromptArgument;
import io.modelcontextprotocol.spec.McpSchema.PromptMessage;
import io.modelcontextprotocol.spec.McpSchema.ReadResourceResult;
import io.modelcontextprotocol.spec.McpSchema.Resource;
import io.modelcontextprotocol.spec.McpSchema.ResourceTemplate;
import io.modelcontextprotocol.spec.McpSchema.Role;
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
import io.modelcontextprotocol.spec.McpSchema.TextContent;
import io.modelcontextprotocol.spec.McpSchema.TextResourceContents;
import io.modelcontextprotocol.spec.McpSchema.Tool;
import jakarta.servlet.http.HttpServlet;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.apache.catalina.Context;
import org.apache.catalina.LifecycleException;
import org.apache.catalina.Wrapper;
import org.apache.catalina.startup.Tomcat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;

/**
 * Self-contained servlet demo for validating context-aware dynamic MCP primitive
 * repositories.
 */
public final class DynamicRepositoryDemoServlet {

	private static final Logger logger = LoggerFactory.getLogger(DynamicRepositoryDemoServlet.class);

	private static final String MCP_ENDPOINT = "/mcp";

	private static final String DEFAULT_MCP_ENDPOINT = "/mcp-default";

	private static final String STATEFUL_MCP_ENDPOINT = "/mcp-stateful";

	private static final String ADMIN_ENDPOINT = "/admin/*";

	private static final String USER_HEADER = "X-User";

	private static final String USER_KEY = "user";

	private static final int PAGE_SIZE = 1;

	private static final String DEFAULT_OWNER = "default";

	private static final Map<String, Object> EMPTY_JSON_SCHEMA = Map.of("type", "object", "properties", Map.of());

	private DynamicRepositoryDemoServlet() {
	}

	public static void main(String[] args) throws Exception {
		int port = Integer.getInteger("demo.port", 8081);

		HttpServletStatelessServerTransport repositoryTransport = HttpServletStatelessServerTransport.builder()
			.messageEndpoint(MCP_ENDPOINT)
			.contextExtractor(DynamicRepositoryDemoServlet::contextFrom)
			.build();

		HttpServletStatelessServerTransport defaultTransport = HttpServletStatelessServerTransport.builder()
			.messageEndpoint(DEFAULT_MCP_ENDPOINT)
			.contextExtractor(DynamicRepositoryDemoServlet::contextFrom)
			.build();

		HttpServletStreamableServerTransportProvider statefulTransport = HttpServletStreamableServerTransportProvider
			.builder()
			.mcpEndpoint(STATEFUL_MCP_ENDPOINT)
			.contextExtractor(DynamicRepositoryDemoServlet::contextFrom)
			.keepAliveInterval(Duration.ofSeconds(30))
			.build();

		McpStatelessAsyncServer repositoryServer = createRepositoryServer(repositoryTransport);
		McpStatelessAsyncServer defaultServer = createDefaultServer(defaultTransport);
		McpAsyncServer statefulServer = createStatefulRepositoryServer(statefulTransport);
		Tomcat tomcat = createEmbeddedTomcat(port, repositoryTransport, defaultTransport, statefulTransport,
				new AdminServlet(repositoryServer, defaultServer, statefulServer));

		try {
			tomcat.start();
			logger.info("Repository demo server started at http://localhost:{}{}", port, MCP_ENDPOINT);
			logger.info("Default-registration demo server started at http://localhost:{}{}", port,
					DEFAULT_MCP_ENDPOINT);
			logger.info("Stateful repository demo server started at http://localhost:{}{}", port,
					STATEFUL_MCP_ENDPOINT);
			logger.info("Admin helpers are available under http://localhost:{}/admin", port);
			logger.info("Set {} to alice or bob to change visible repository-backed MCP primitives.", USER_HEADER);
			tomcat.getServer().await();
		}
		catch (LifecycleException ex) {
			logger.error("Failed to start demo server", ex);
			throw ex;
		}
		finally {
			repositoryServer.closeGracefully().block();
			defaultServer.closeGracefully().block();
			statefulServer.closeGracefully().block();
			tomcat.stop();
			tomcat.destroy();
		}
	}

	private static McpStatelessAsyncServer createRepositoryServer(HttpServletStatelessServerTransport transport) {
		DemoToolsRepository toolsRepository = new DemoToolsRepository();
		DemoResourcesRepository resourcesRepository = new DemoResourcesRepository();
		DemoPromptsRepository promptsRepository = new DemoPromptsRepository();
		DemoCompletionsRepository completionsRepository = new DemoCompletionsRepository(promptsRepository,
				resourcesRepository);
		return McpServer.async(transport)
			.serverInfo("dynamic-repository-demo", "1.0.0")
			.capabilities(ServerCapabilities.builder()
				.tools(false)
				.resources(false, false)
				.prompts(false)
				.completions()
				.build())
			.toolsRepository(toolsRepository)
			.resourcesRepository(resourcesRepository)
			.promptsRepository(promptsRepository)
			.completionsRepository(completionsRepository)
			.build();
	}

	private static McpStatelessAsyncServer createDefaultServer(HttpServletStatelessServerTransport transport) {
		return McpServer.async(transport)
			.serverInfo("default-registration-demo", "1.0.0")
			.capabilities(ServerCapabilities.builder()
				.tools(false)
				.resources(false, false)
				.prompts(false)
				.completions()
				.build())
			.tools(defaultTools())
			.resources(defaultResources())
			.resourceTemplates(defaultResourceTemplates())
			.prompts(defaultPrompts())
			.completions(defaultCompletions())
			.build();
	}

	private static McpAsyncServer createStatefulRepositoryServer(
			HttpServletStreamableServerTransportProvider transportProvider) {
		DemoStatefulToolsRepository toolsRepository = new DemoStatefulToolsRepository();
		DemoStatefulResourcesRepository resourcesRepository = new DemoStatefulResourcesRepository();
		DemoStatefulPromptsRepository promptsRepository = new DemoStatefulPromptsRepository();
		DemoStatefulCompletionsRepository completionsRepository = new DemoStatefulCompletionsRepository(
				promptsRepository, resourcesRepository);
		return McpServer.async(transportProvider)
			.serverInfo("stateful-dynamic-repository-demo", "1.0.0")
			.capabilities(
					ServerCapabilities.builder().tools(true).resources(false, true).prompts(true).completions().build())
			.toolsRepository(toolsRepository)
			.resourcesRepository(resourcesRepository)
			.promptsRepository(promptsRepository)
			.completionsRepository(completionsRepository)
			.build();
	}

	private static McpTransportContext contextFrom(HttpServletRequest request) {
		String user = Optional.ofNullable(request.getHeader(USER_HEADER))
			.filter(value -> !value.isBlank())
			.orElse("anonymous");
		return McpTransportContext.create(Map.of(USER_KEY, user));
	}

	private static Tomcat createEmbeddedTomcat(int port, HttpServletStatelessServerTransport repositoryTransport,
			HttpServletStatelessServerTransport defaultTransport,
			HttpServletStreamableServerTransportProvider statefulTransport, HttpServlet adminServlet) {
		Tomcat tomcat = new Tomcat();
		tomcat.setPort(port);

		String baseDir = System.getProperty("java.io.tmpdir");
		tomcat.setBaseDir(baseDir);

		Context context = tomcat.addContext("", baseDir);

		addServlet(context, "repositoryMcpServlet", repositoryTransport, MCP_ENDPOINT);
		addServlet(context, "defaultMcpServlet", defaultTransport, DEFAULT_MCP_ENDPOINT);
		addServlet(context, "statefulMcpServlet", statefulTransport, STATEFUL_MCP_ENDPOINT);
		addServlet(context, "dynamicRepositoryAdminServlet", adminServlet, ADMIN_ENDPOINT);

		tomcat.getConnector().setAsyncTimeout(30000);
		return tomcat;
	}

	private static void addServlet(Context context, String name, HttpServlet servlet, String mapping) {
		Wrapper wrapper = context.createWrapper();
		wrapper.setName(name);
		wrapper.setServlet(servlet);
		wrapper.setLoadOnStartup(1);
		wrapper.setAsyncSupported(true);
		context.addChild(wrapper);
		context.addServletMappingDecoded(mapping, name);
	}

	private static List<McpStatelessServerFeatures.AsyncToolSpecification> defaultTools() {
		return List.of(DemoToolsRepository.toolSpecification(DEFAULT_OWNER, "primary"),
				DemoToolsRepository.toolSpecification(DEFAULT_OWNER, "secondary"));
	}

	private static List<McpStatelessServerFeatures.AsyncResourceSpecification> defaultResources() {
		return List.of(DemoResourcesRepository.resourceSpecification(DEFAULT_OWNER, "profile"),
				DemoResourcesRepository.resourceSpecification(DEFAULT_OWNER, "settings"));
	}

	private static List<McpStatelessServerFeatures.AsyncResourceTemplateSpecification> defaultResourceTemplates() {
		return List.of(DemoResourcesRepository.resourceTemplateSpecification(DEFAULT_OWNER, "docs", "doc"),
				DemoResourcesRepository.resourceTemplateSpecification(DEFAULT_OWNER, "reports", "report"));
	}

	private static List<McpStatelessServerFeatures.AsyncPromptSpecification> defaultPrompts() {
		return List.of(DemoPromptsRepository.promptSpecification(DEFAULT_OWNER, "primary", "topic"),
				DemoPromptsRepository.promptSpecification(DEFAULT_OWNER, "secondary", "tone"));
	}

	private static List<McpStatelessServerFeatures.AsyncCompletionSpecification> defaultCompletions() {
		return List.of(
				DemoCompletionsRepository.fixedOwnerCompletionSpecification(
						new McpSchema.PromptReference(promptName(DEFAULT_OWNER, "primary")), DEFAULT_OWNER),
				DemoCompletionsRepository.fixedOwnerCompletionSpecification(
						new McpSchema.ResourceReference(resourceTemplateUri(DEFAULT_OWNER, "docs", "doc")),
						DEFAULT_OWNER));
	}

	private static String user(McpTransportContext context) {
		Object user = context.get(USER_KEY);
		return (user != null) ? user.toString() : "anonymous";
	}

	private static String user(McpAsyncServerExchange exchange) {
		return user((exchange != null) ? exchange.transportContext() : McpTransportContext.EMPTY);
	}

	private static String userParam(HttpServletRequest request) {
		return Optional.ofNullable(request.getParameter("user")).filter(value -> !value.isBlank()).orElse("alice");
	}

	private static String toolName(String user, String name) {
		return user + "-tool-" + name;
	}

	private static String resourceUri(String user, String name) {
		return "demo://users/" + user + "/" + name;
	}

	private static String resourceTemplateUri(String user, String group, String variable) {
		return "demo://users/" + user + "/" + group + "/{" + variable + "}";
	}

	private static String promptName(String user, String name) {
		return user + "-prompt-" + name;
	}

	private static int pageStart(McpSchema.PaginatedRequest request) {
		if (request == null || request.cursor() == null || request.cursor().isBlank()) {
			return 0;
		}
		try {
			return Math.max(0, Integer.parseInt(request.cursor()));
		}
		catch (NumberFormatException ex) {
			return 0;
		}
	}

	private static String nextCursor(List<?> items, int start) {
		int next = start + PAGE_SIZE;
		return (next < items.size()) ? Integer.toString(next) : null;
	}

	private static <T> List<T> page(List<T> items, int start) {
		if (start >= items.size()) {
			return List.of();
		}
		return List.copyOf(items.subList(start, Math.min(items.size(), start + PAGE_SIZE)));
	}

	private static boolean belongsToUser(String primitiveName, String user) {
		return primitiveName.startsWith(user + "-");
	}

	private static boolean uriBelongsToUser(String uri, String user) {
		return uri.startsWith("demo://users/" + user + "/");
	}

	private static boolean matchesUriTemplate(String uriTemplate, String uri) {
		String[] templateSegments = uriTemplate.split("/", -1);
		String[] uriSegments = uri.split("/", -1);
		if (templateSegments.length != uriSegments.length) {
			return false;
		}
		for (int i = 0; i < templateSegments.length; i++) {
			String templateSegment = templateSegments[i];
			if (templateSegment.startsWith("{") && templateSegment.endsWith("}")) {
				continue;
			}
			if (!templateSegment.equals(uriSegments[i])) {
				return false;
			}
		}
		return true;
	}

	private static final class DemoToolsRepository implements StatelessToolsRepository {

		private final ConcurrentHashMap<String, McpStatelessServerFeatures.AsyncToolSpecification> runtimeTools = new ConcurrentHashMap<>();

		@Override
		public Mono<ListToolsResult> listTools(McpTransportContext transportContext,
				McpSchema.PaginatedRequest request) {
			List<Tool> tools = toolSpecifications(user(transportContext)).stream()
				.map(McpStatelessServerFeatures.AsyncToolSpecification::tool)
				.toList();
			int start = pageStart(request);
			return Mono.just(ListToolsResult.builder(page(tools, start)).nextCursor(nextCursor(tools, start)).build());
		}

		@Override
		public Mono<McpStatelessServerFeatures.AsyncToolSpecification> resolveTool(String name,
				McpTransportContext transportContext) {
			return Mono.justOrEmpty(toolSpecifications(user(transportContext)).stream()
				.filter(spec -> spec.tool().name().equals(name))
				.findFirst());
		}

		@Override
		public void addTool(McpStatelessServerFeatures.AsyncToolSpecification toolSpecification) {
			this.runtimeTools.put(toolSpecification.tool().name(), toolSpecification);
		}

		@Override
		public boolean removeTool(String name) {
			return this.runtimeTools.remove(name) != null;
		}

		private List<McpStatelessServerFeatures.AsyncToolSpecification> toolSpecifications(String user) {
			List<McpStatelessServerFeatures.AsyncToolSpecification> tools = new ArrayList<>();
			tools.add(toolSpecification(user, "primary"));
			tools.add(toolSpecification(user, "secondary"));
			this.runtimeTools.values()
				.stream()
				.filter(spec -> belongsToUser(spec.tool().name(), user))
				.sorted(Comparator.comparing(spec -> spec.tool().name()))
				.forEach(tools::add);
			return tools;
		}

		private static McpStatelessServerFeatures.AsyncToolSpecification toolSpecification(String owner, String name) {
			return new McpStatelessServerFeatures.AsyncToolSpecification(tool(owner, name),
					(transportContext,
							request) -> Mono.just(CallToolResult.builder()
								.addTextContent("tool=" + request.name() + ", owner=" + owner + ", request user="
										+ user(transportContext))
								.build()));
		}

		private static Tool tool(String user, String name) {
			return Tool.builder(toolName(user, name), EMPTY_JSON_SCHEMA).description("Visible only to " + user).build();
		}

	}

	private static final class DemoResourcesRepository implements StatelessResourcesRepository {

		private final ConcurrentHashMap<String, McpStatelessServerFeatures.AsyncResourceSpecification> runtimeResources = new ConcurrentHashMap<>();

		private final ConcurrentHashMap<String, McpStatelessServerFeatures.AsyncResourceTemplateSpecification> runtimeTemplates = new ConcurrentHashMap<>();

		@Override
		public Mono<ListResourcesResult> listResources(McpTransportContext transportContext,
				McpSchema.PaginatedRequest request) {
			List<Resource> resources = resourceSpecifications(user(transportContext)).stream()
				.map(McpStatelessServerFeatures.AsyncResourceSpecification::resource)
				.toList();
			int start = pageStart(request);
			return Mono.just(ListResourcesResult.builder(page(resources, start))
				.nextCursor(nextCursor(resources, start))
				.build());
		}

		@Override
		public Mono<ListResourceTemplatesResult> listResourceTemplates(McpTransportContext transportContext,
				McpSchema.PaginatedRequest request) {
			List<ResourceTemplate> templates = resourceTemplateSpecifications(user(transportContext)).stream()
				.map(McpStatelessServerFeatures.AsyncResourceTemplateSpecification::resourceTemplate)
				.toList();
			int start = pageStart(request);
			return Mono.just(ListResourceTemplatesResult.builder(page(templates, start))
				.nextCursor(nextCursor(templates, start))
				.build());
		}

		@Override
		public Mono<McpStatelessServerFeatures.AsyncResourceSpecification> resolveResource(String uri,
				McpTransportContext transportContext) {
			return Mono.justOrEmpty(resourceSpecifications(user(transportContext)).stream()
				.filter(spec -> spec.resource().uri().equals(uri))
				.findFirst());
		}

		@Override
		public Mono<McpStatelessServerFeatures.AsyncResourceTemplateSpecification> resolveResourceTemplate(String uri,
				McpTransportContext transportContext) {
			return Mono.justOrEmpty(resourceTemplateSpecifications(user(transportContext)).stream()
				.filter(spec -> matchesUriTemplate(spec.resourceTemplate().uriTemplate(), uri))
				.findFirst());
		}

		@Override
		public void addResource(McpStatelessServerFeatures.AsyncResourceSpecification resourceSpecification) {
			this.runtimeResources.put(resourceSpecification.resource().uri(), resourceSpecification);
		}

		@Override
		public boolean removeResource(String uri) {
			return this.runtimeResources.remove(uri) != null;
		}

		@Override
		public void addResourceTemplate(
				McpStatelessServerFeatures.AsyncResourceTemplateSpecification resourceTemplateSpecification) {
			this.runtimeTemplates.put(resourceTemplateSpecification.resourceTemplate().uriTemplate(),
					resourceTemplateSpecification);
		}

		@Override
		public boolean removeResourceTemplate(String uriTemplate) {
			return this.runtimeTemplates.remove(uriTemplate) != null;
		}

		private List<McpStatelessServerFeatures.AsyncResourceSpecification> resourceSpecifications(String user) {
			List<McpStatelessServerFeatures.AsyncResourceSpecification> resources = new ArrayList<>();
			resources.add(resourceSpecification(user, "profile"));
			resources.add(resourceSpecification(user, "settings"));
			this.runtimeResources.values()
				.stream()
				.filter(spec -> uriBelongsToUser(spec.resource().uri(), user))
				.sorted(Comparator.comparing(spec -> spec.resource().uri()))
				.forEach(resources::add);
			return resources;
		}

		private List<McpStatelessServerFeatures.AsyncResourceTemplateSpecification> resourceTemplateSpecifications(
				String user) {
			List<McpStatelessServerFeatures.AsyncResourceTemplateSpecification> templates = new ArrayList<>();
			templates.add(resourceTemplateSpecification(user, "docs", "doc"));
			templates.add(resourceTemplateSpecification(user, "reports", "report"));
			this.runtimeTemplates.values()
				.stream()
				.filter(spec -> uriBelongsToUser(spec.resourceTemplate().uriTemplate(), user))
				.sorted(Comparator.comparing(spec -> spec.resourceTemplate().uriTemplate()))
				.forEach(templates::add);
			return templates;
		}

		private boolean hasResourceReference(String uri, String user) {
			return resourceSpecifications(user).stream().anyMatch(spec -> spec.resource().uri().equals(uri))
					|| resourceTemplateSpecifications(user).stream()
						.anyMatch(spec -> spec.resourceTemplate().uriTemplate().equals(uri));
		}

		private static McpStatelessServerFeatures.AsyncResourceSpecification resourceSpecification(String owner,
				String name) {
			return new McpStatelessServerFeatures.AsyncResourceSpecification(resource(owner, name),
					(transportContext, request) -> Mono.just(
							ReadResourceResult
								.builder(
										List.of(TextResourceContents
											.builder(request.uri(),
													"resource=" + name + ", owner=" + owner + ", request user="
															+ user(transportContext))
											.mimeType("text/plain")
											.build()))
								.build()));
		}

		private static McpStatelessServerFeatures.AsyncResourceTemplateSpecification resourceTemplateSpecification(
				String owner, String group, String variable) {
			return new McpStatelessServerFeatures.AsyncResourceTemplateSpecification(
					resourceTemplate(owner, group, variable),
					(transportContext, request) -> Mono.just(
							ReadResourceResult
								.builder(
										List.of(TextResourceContents
											.builder(request.uri(),
													"resource template=" + group + ", owner=" + owner
															+ ", request user=" + user(transportContext))
											.mimeType("text/plain")
											.build()))
								.build()));
		}

		private static Resource resource(String user, String name) {
			return Resource.builder(resourceUri(user, name), user + " " + name).mimeType("text/plain").build();
		}

		private static ResourceTemplate resourceTemplate(String user, String group, String variable) {
			return ResourceTemplate.builder(resourceTemplateUri(user, group, variable), user + " " + group)
				.mimeType("text/plain")
				.build();
		}

	}

	private static final class DemoPromptsRepository implements StatelessPromptsRepository {

		private final ConcurrentHashMap<String, McpStatelessServerFeatures.AsyncPromptSpecification> runtimePrompts = new ConcurrentHashMap<>();

		@Override
		public Mono<ListPromptsResult> listPrompts(McpTransportContext transportContext,
				McpSchema.PaginatedRequest request) {
			List<Prompt> prompts = promptSpecifications(user(transportContext)).stream()
				.map(McpStatelessServerFeatures.AsyncPromptSpecification::prompt)
				.toList();
			int start = pageStart(request);
			return Mono
				.just(ListPromptsResult.builder(page(prompts, start)).nextCursor(nextCursor(prompts, start)).build());
		}

		@Override
		public Mono<McpStatelessServerFeatures.AsyncPromptSpecification> resolvePrompt(String name,
				McpTransportContext transportContext) {
			return Mono.justOrEmpty(promptSpecifications(user(transportContext)).stream()
				.filter(spec -> spec.prompt().name().equals(name))
				.findFirst());
		}

		@Override
		public void addPrompt(McpStatelessServerFeatures.AsyncPromptSpecification promptSpecification) {
			this.runtimePrompts.put(promptSpecification.prompt().name(), promptSpecification);
		}

		@Override
		public boolean removePrompt(String name) {
			return this.runtimePrompts.remove(name) != null;
		}

		private List<McpStatelessServerFeatures.AsyncPromptSpecification> promptSpecifications(String user) {
			List<McpStatelessServerFeatures.AsyncPromptSpecification> prompts = new ArrayList<>();
			prompts.add(promptSpecification(user, "primary", "topic"));
			prompts.add(promptSpecification(user, "secondary", "tone"));
			this.runtimePrompts.values()
				.stream()
				.filter(spec -> belongsToUser(spec.prompt().name(), user))
				.sorted(Comparator.comparing(spec -> spec.prompt().name()))
				.forEach(prompts::add);
			return prompts;
		}

		private boolean hasPrompt(String name, String user) {
			return promptSpecifications(user).stream().anyMatch(spec -> spec.prompt().name().equals(name));
		}

		private static McpStatelessServerFeatures.AsyncPromptSpecification promptSpecification(String owner,
				String name, String argumentName) {
			return new McpStatelessServerFeatures.AsyncPromptSpecification(prompt(owner, name, argumentName),
					(transportContext, request) -> Mono.just(
							GetPromptResult
								.builder(List.of(
										PromptMessage.builder(Role.USER, TextContent
											.builder("prompt=" + request.name() + ", owner=" + owner + ", request user="
													+ user(transportContext) + ", args=" + request.arguments())
											.build()).build()))
								.build()));
		}

		private static Prompt prompt(String user, String name, String argumentName) {
			return Prompt.builder(promptName(user, name))
				.description("Visible only to " + user)
				.arguments(List.of(PromptArgument.builder(argumentName).required(false).build()))
				.build();
		}

	}

	private static final class DemoCompletionsRepository implements StatelessCompletionsRepository {

		private final DemoPromptsRepository promptsRepository;

		private final DemoResourcesRepository resourcesRepository;

		private DemoCompletionsRepository(DemoPromptsRepository promptsRepository,
				DemoResourcesRepository resourcesRepository) {
			this.promptsRepository = promptsRepository;
			this.resourcesRepository = resourcesRepository;
		}

		@Override
		public Mono<McpStatelessServerFeatures.AsyncCompletionSpecification> resolveCompletion(
				McpSchema.CompleteReference reference, McpTransportContext transportContext) {
			String user = user(transportContext);
			if (reference instanceof McpSchema.PromptReference promptReference
					&& this.promptsRepository.hasPrompt(promptReference.name(), user)) {
				return Mono.just(completionSpecification(reference));
			}
			if (reference instanceof McpSchema.ResourceReference resourceReference
					&& this.resourcesRepository.hasResourceReference(resourceReference.uri(), user)) {
				return Mono.just(completionSpecification(reference));
			}
			return Mono.empty();
		}

		private static McpStatelessServerFeatures.AsyncCompletionSpecification completionSpecification(
				McpSchema.CompleteReference reference) {
			return new McpStatelessServerFeatures.AsyncCompletionSpecification(reference,
					(transportContext, request) -> {
						String user = user(transportContext);
						return completeForOwner(user, request);
					});
		}

		private static McpStatelessServerFeatures.AsyncCompletionSpecification fixedOwnerCompletionSpecification(
				McpSchema.CompleteReference reference, String owner) {
			return new McpStatelessServerFeatures.AsyncCompletionSpecification(reference,
					(transportContext, request) -> completeForOwner(owner, request));
		}

		private static Mono<CompleteResult> completeForOwner(String owner, CompleteRequest request) {
			String argumentName = request.argument().name();
			String value = request.argument().value();
			List<String> candidates = List.of(owner + "-" + argumentName + "-alpha",
					owner + "-" + argumentName + "-beta", owner + "-" + argumentName + "-gamma");
			List<String> values = candidates.stream().filter(candidate -> candidate.contains(value)).toList();
			if (values.isEmpty()) {
				values = candidates;
			}
			return Mono
				.just(new CompleteResult(new CompleteResult.CompleteCompletion(values, candidates.size(), false)));
		}

	}

	private static final class DemoStatefulToolsRepository implements ToolsRepository {

		private final ConcurrentHashMap<String, McpServerFeatures.AsyncToolSpecification> runtimeTools = new ConcurrentHashMap<>();

		@Override
		public Mono<ListToolsResult> listTools(McpAsyncServerExchange exchange, McpSchema.PaginatedRequest request) {
			List<Tool> tools = toolSpecifications(user(exchange)).stream()
				.map(McpServerFeatures.AsyncToolSpecification::tool)
				.toList();
			int start = pageStart(request);
			return Mono.just(ListToolsResult.builder(page(tools, start)).nextCursor(nextCursor(tools, start)).build());
		}

		@Override
		public Mono<McpServerFeatures.AsyncToolSpecification> resolveTool(String name,
				McpAsyncServerExchange exchange) {
			return Mono.justOrEmpty(toolSpecifications(user(exchange)).stream()
				.filter(spec -> spec.tool().name().equals(name))
				.findFirst());
		}

		@Override
		public void addTool(McpServerFeatures.AsyncToolSpecification toolSpecification) {
			this.runtimeTools.put(toolSpecification.tool().name(), toolSpecification);
		}

		@Override
		public boolean removeTool(String name) {
			return this.runtimeTools.remove(name) != null;
		}

		private List<McpServerFeatures.AsyncToolSpecification> toolSpecifications(String user) {
			List<McpServerFeatures.AsyncToolSpecification> tools = new ArrayList<>();
			tools.add(toolSpecification(user, "primary"));
			tools.add(toolSpecification(user, "secondary"));
			this.runtimeTools.values()
				.stream()
				.filter(spec -> belongsToUser(spec.tool().name(), user))
				.sorted(Comparator.comparing(spec -> spec.tool().name()))
				.forEach(tools::add);
			return tools;
		}

		private static McpServerFeatures.AsyncToolSpecification toolSpecification(String owner, String name) {
			return new McpServerFeatures.AsyncToolSpecification(DemoToolsRepository.tool(owner, name),
					(exchange, request) -> Mono.just(CallToolResult.builder()
						.addTextContent(
								"tool=" + request.name() + ", owner=" + owner + ", request user=" + user(exchange))
						.build()));
		}

	}

	private static final class DemoStatefulResourcesRepository implements ResourcesRepository {

		private final ConcurrentHashMap<String, McpServerFeatures.AsyncResourceSpecification> runtimeResources = new ConcurrentHashMap<>();

		private final ConcurrentHashMap<String, McpServerFeatures.AsyncResourceTemplateSpecification> runtimeTemplates = new ConcurrentHashMap<>();

		@Override
		public Mono<ListResourcesResult> listResources(McpAsyncServerExchange exchange,
				McpSchema.PaginatedRequest request) {
			List<Resource> resources = resourceSpecifications(user(exchange)).stream()
				.map(McpServerFeatures.AsyncResourceSpecification::resource)
				.toList();
			int start = pageStart(request);
			return Mono.just(ListResourcesResult.builder(page(resources, start))
				.nextCursor(nextCursor(resources, start))
				.build());
		}

		@Override
		public Mono<ListResourceTemplatesResult> listResourceTemplates(McpAsyncServerExchange exchange,
				McpSchema.PaginatedRequest request) {
			List<ResourceTemplate> templates = resourceTemplateSpecifications(user(exchange)).stream()
				.map(McpServerFeatures.AsyncResourceTemplateSpecification::resourceTemplate)
				.toList();
			int start = pageStart(request);
			return Mono.just(ListResourceTemplatesResult.builder(page(templates, start))
				.nextCursor(nextCursor(templates, start))
				.build());
		}

		@Override
		public Mono<McpServerFeatures.AsyncResourceSpecification> resolveResource(String uri,
				McpAsyncServerExchange exchange) {
			return Mono.justOrEmpty(resourceSpecifications(user(exchange)).stream()
				.filter(spec -> spec.resource().uri().equals(uri))
				.findFirst());
		}

		@Override
		public Mono<McpServerFeatures.AsyncResourceTemplateSpecification> resolveResourceTemplate(String uri,
				McpAsyncServerExchange exchange) {
			return Mono.justOrEmpty(resourceTemplateSpecifications(user(exchange)).stream()
				.filter(spec -> matchesUriTemplate(spec.resourceTemplate().uriTemplate(), uri))
				.findFirst());
		}

		@Override
		public void addResource(McpServerFeatures.AsyncResourceSpecification resourceSpecification) {
			this.runtimeResources.put(resourceSpecification.resource().uri(), resourceSpecification);
		}

		@Override
		public boolean removeResource(String uri) {
			return this.runtimeResources.remove(uri) != null;
		}

		@Override
		public void addResourceTemplate(
				McpServerFeatures.AsyncResourceTemplateSpecification resourceTemplateSpecification) {
			this.runtimeTemplates.put(resourceTemplateSpecification.resourceTemplate().uriTemplate(),
					resourceTemplateSpecification);
		}

		@Override
		public boolean removeResourceTemplate(String uriTemplate) {
			return this.runtimeTemplates.remove(uriTemplate) != null;
		}

		private List<McpServerFeatures.AsyncResourceSpecification> resourceSpecifications(String user) {
			List<McpServerFeatures.AsyncResourceSpecification> resources = new ArrayList<>();
			resources.add(resourceSpecification(user, "profile"));
			resources.add(resourceSpecification(user, "settings"));
			this.runtimeResources.values()
				.stream()
				.filter(spec -> uriBelongsToUser(spec.resource().uri(), user))
				.sorted(Comparator.comparing(spec -> spec.resource().uri()))
				.forEach(resources::add);
			return resources;
		}

		private List<McpServerFeatures.AsyncResourceTemplateSpecification> resourceTemplateSpecifications(String user) {
			List<McpServerFeatures.AsyncResourceTemplateSpecification> templates = new ArrayList<>();
			templates.add(resourceTemplateSpecification(user, "docs", "doc"));
			templates.add(resourceTemplateSpecification(user, "reports", "report"));
			this.runtimeTemplates.values()
				.stream()
				.filter(spec -> uriBelongsToUser(spec.resourceTemplate().uriTemplate(), user))
				.sorted(Comparator.comparing(spec -> spec.resourceTemplate().uriTemplate()))
				.forEach(templates::add);
			return templates;
		}

		private boolean hasResourceReference(String uri, String user) {
			return resourceSpecifications(user).stream().anyMatch(spec -> spec.resource().uri().equals(uri))
					|| resourceTemplateSpecifications(user).stream()
						.anyMatch(spec -> spec.resourceTemplate().uriTemplate().equals(uri));
		}

		private static McpServerFeatures.AsyncResourceSpecification resourceSpecification(String owner, String name) {
			return new McpServerFeatures.AsyncResourceSpecification(DemoResourcesRepository.resource(owner, name),
					(exchange, request) -> Mono.just(
							ReadResourceResult
								.builder(
										List.of(TextResourceContents
											.builder(request.uri(),
													"resource=" + name + ", owner=" + owner + ", request user="
															+ user(exchange))
											.mimeType("text/plain")
											.build()))
								.build()));
		}

		private static McpServerFeatures.AsyncResourceTemplateSpecification resourceTemplateSpecification(String owner,
				String group, String variable) {
			return new McpServerFeatures.AsyncResourceTemplateSpecification(
					DemoResourcesRepository.resourceTemplate(owner, group, variable),
					(exchange, request) -> Mono.just(
							ReadResourceResult
								.builder(
										List.of(TextResourceContents
											.builder(request.uri(),
													"resource template=" + group + ", owner=" + owner
															+ ", request user=" + user(exchange))
											.mimeType("text/plain")
											.build()))
								.build()));
		}

	}

	private static final class DemoStatefulPromptsRepository implements PromptsRepository {

		private final ConcurrentHashMap<String, McpServerFeatures.AsyncPromptSpecification> runtimePrompts = new ConcurrentHashMap<>();

		@Override
		public Mono<ListPromptsResult> listPrompts(McpAsyncServerExchange exchange,
				McpSchema.PaginatedRequest request) {
			List<Prompt> prompts = promptSpecifications(user(exchange)).stream()
				.map(McpServerFeatures.AsyncPromptSpecification::prompt)
				.toList();
			int start = pageStart(request);
			return Mono
				.just(ListPromptsResult.builder(page(prompts, start)).nextCursor(nextCursor(prompts, start)).build());
		}

		@Override
		public Mono<McpServerFeatures.AsyncPromptSpecification> resolvePrompt(String name,
				McpAsyncServerExchange exchange) {
			return Mono.justOrEmpty(promptSpecifications(user(exchange)).stream()
				.filter(spec -> spec.prompt().name().equals(name))
				.findFirst());
		}

		@Override
		public void addPrompt(McpServerFeatures.AsyncPromptSpecification promptSpecification) {
			this.runtimePrompts.put(promptSpecification.prompt().name(), promptSpecification);
		}

		@Override
		public boolean removePrompt(String name) {
			return this.runtimePrompts.remove(name) != null;
		}

		private List<McpServerFeatures.AsyncPromptSpecification> promptSpecifications(String user) {
			List<McpServerFeatures.AsyncPromptSpecification> prompts = new ArrayList<>();
			prompts.add(promptSpecification(user, "primary", "topic"));
			prompts.add(promptSpecification(user, "secondary", "tone"));
			this.runtimePrompts.values()
				.stream()
				.filter(spec -> belongsToUser(spec.prompt().name(), user))
				.sorted(Comparator.comparing(spec -> spec.prompt().name()))
				.forEach(prompts::add);
			return prompts;
		}

		private boolean hasPrompt(String name, String user) {
			return promptSpecifications(user).stream().anyMatch(spec -> spec.prompt().name().equals(name));
		}

		private static McpServerFeatures.AsyncPromptSpecification promptSpecification(String owner, String name,
				String argumentName) {
			return new McpServerFeatures.AsyncPromptSpecification(
					DemoPromptsRepository.prompt(owner, name, argumentName),
					(exchange, request) -> Mono.just(
							GetPromptResult
								.builder(List.of(
										PromptMessage.builder(Role.USER, TextContent
											.builder("prompt=" + request.name() + ", owner=" + owner + ", request user="
													+ user(exchange) + ", args=" + request.arguments())
											.build()).build()))
								.build()));
		}

	}

	private static final class DemoStatefulCompletionsRepository implements CompletionsRepository {

		private final DemoStatefulPromptsRepository promptsRepository;

		private final DemoStatefulResourcesRepository resourcesRepository;

		private DemoStatefulCompletionsRepository(DemoStatefulPromptsRepository promptsRepository,
				DemoStatefulResourcesRepository resourcesRepository) {
			this.promptsRepository = promptsRepository;
			this.resourcesRepository = resourcesRepository;
		}

		@Override
		public Mono<McpServerFeatures.AsyncCompletionSpecification> resolveCompletion(
				McpSchema.CompleteReference reference, McpAsyncServerExchange exchange) {
			String user = user(exchange);
			if (reference instanceof McpSchema.PromptReference promptReference
					&& this.promptsRepository.hasPrompt(promptReference.name(), user)) {
				return Mono.just(completionSpecification(reference));
			}
			if (reference instanceof McpSchema.ResourceReference resourceReference
					&& this.resourcesRepository.hasResourceReference(resourceReference.uri(), user)) {
				return Mono.just(completionSpecification(reference));
			}
			return Mono.empty();
		}

		private static McpServerFeatures.AsyncCompletionSpecification completionSpecification(
				McpSchema.CompleteReference reference) {
			return new McpServerFeatures.AsyncCompletionSpecification(reference,
					(exchange, request) -> DemoCompletionsRepository.completeForOwner(user(exchange), request));
		}

	}

	private static final class AdminServlet extends HttpServlet {

		private final McpStatelessAsyncServer repositoryServer;

		private final McpStatelessAsyncServer defaultServer;

		private final McpAsyncServer statefulServer;

		AdminServlet(McpStatelessAsyncServer repositoryServer, McpStatelessAsyncServer defaultServer,
				McpAsyncServer statefulServer) {
			this.repositoryServer = repositoryServer;
			this.defaultServer = defaultServer;
			this.statefulServer = statefulServer;
		}

		@Override
		protected void doGet(HttpServletRequest request, HttpServletResponse response) throws IOException {
			String path = Optional.ofNullable(request.getPathInfo()).orElse("/help");
			if ("/context-free".equals(path)) {
				writeJson(response, contextFreeJson(this.repositoryServer));
				return;
			}
			if ("/default/context-free".equals(path)) {
				writeJson(response, contextFreeJson(this.defaultServer));
				return;
			}
			if ("/stateful/context-free".equals(path)) {
				writeJson(response, contextFreeJson(this.statefulServer));
				return;
			}
			writeJson(response,
					"{\"mcp\":{\"repository\":\"POST /mcp\",\"default\":\"POST /mcp-default\",\"stateful\":\"POST /mcp-stateful\"},\"endpoints\":[\"POST /admin/runtime/add?user=alice\",\"POST /admin/runtime/remove?user=alice\",\"GET /admin/context-free\",\"POST /admin/default/runtime/add\",\"POST /admin/default/runtime/remove\",\"GET /admin/default/context-free\",\"POST /admin/stateful/runtime/add?user=alice\",\"POST /admin/stateful/runtime/remove?user=alice\",\"GET /admin/stateful/context-free\"]}");
		}

		@Override
		protected void doPost(HttpServletRequest request, HttpServletResponse response) throws IOException {
			String path = Optional.ofNullable(request.getPathInfo()).orElse("");
			String user = userParam(request);
			if ("/runtime/add".equals(path)) {
				addRuntimePrimitives(this.repositoryServer, user);
				writeJson(response, "{\"status\":\"added\",\"user\":" + quote(user) + "}");
				return;
			}
			if ("/runtime/remove".equals(path)) {
				removeRuntimePrimitives(this.repositoryServer, user);
				writeJson(response, "{\"status\":\"removed\",\"user\":" + quote(user) + "}");
				return;
			}
			if ("/default/runtime/add".equals(path)) {
				addRuntimePrimitives(this.defaultServer, DEFAULT_OWNER);
				writeJson(response, "{\"status\":\"added\",\"owner\":" + quote(DEFAULT_OWNER) + "}");
				return;
			}
			if ("/default/runtime/remove".equals(path)) {
				removeRuntimePrimitives(this.defaultServer, DEFAULT_OWNER);
				writeJson(response, "{\"status\":\"removed\",\"owner\":" + quote(DEFAULT_OWNER) + "}");
				return;
			}
			if ("/stateful/runtime/add".equals(path)) {
				addRuntimePrimitives(this.statefulServer, user);
				writeJson(response, "{\"status\":\"added\",\"user\":" + quote(user) + "}");
				return;
			}
			if ("/stateful/runtime/remove".equals(path)) {
				removeRuntimePrimitives(this.statefulServer, user);
				writeJson(response, "{\"status\":\"removed\",\"user\":" + quote(user) + "}");
				return;
			}
			response.sendError(HttpServletResponse.SC_NOT_FOUND, "Unknown admin endpoint: " + path);
		}

		private static void addRuntimePrimitives(McpStatelessAsyncServer server, String user) {
			server.addTool(DemoToolsRepository.toolSpecification(user, "runtime")).block();
			server.addResource(DemoResourcesRepository.resourceSpecification(user, "runtime-resource")).block();
			server
				.addResourceTemplate(
						DemoResourcesRepository.resourceTemplateSpecification(user, "runtime-docs", "item"))
				.block();
			server.addPrompt(DemoPromptsRepository.promptSpecification(user, "runtime", "item")).block();
		}

		private static void addRuntimePrimitives(McpAsyncServer server, String user) {
			server.addTool(DemoStatefulToolsRepository.toolSpecification(user, "runtime")).block();
			server.addResource(DemoStatefulResourcesRepository.resourceSpecification(user, "runtime-resource")).block();
			server
				.addResourceTemplate(
						DemoStatefulResourcesRepository.resourceTemplateSpecification(user, "runtime-docs", "item"))
				.block();
			server.addPrompt(DemoStatefulPromptsRepository.promptSpecification(user, "runtime", "item")).block();
		}

		private static void removeRuntimePrimitives(McpStatelessAsyncServer server, String user) {
			server.removeTool(toolName(user, "runtime")).block();
			server.removeResource(resourceUri(user, "runtime-resource")).block();
			server.removeResourceTemplate(resourceTemplateUri(user, "runtime-docs", "item")).block();
			server.removePrompt(promptName(user, "runtime")).block();
		}

		private static void removeRuntimePrimitives(McpAsyncServer server, String user) {
			server.removeTool(toolName(user, "runtime")).block();
			server.removeResource(resourceUri(user, "runtime-resource")).block();
			server.removeResourceTemplate(resourceTemplateUri(user, "runtime-docs", "item")).block();
			server.removePrompt(promptName(user, "runtime")).block();
		}

		private static String contextFreeJson(McpStatelessAsyncServer server) {
			List<String> tools = server.listTools().map(Tool::name).collectList().block();
			List<String> resources = server.listResources().map(Resource::uri).collectList().block();
			List<String> templates = server.listResourceTemplates()
				.map(ResourceTemplate::uriTemplate)
				.collectList()
				.block();
			List<String> prompts = server.listPrompts().map(Prompt::name).collectList().block();
			return "{\"tools\":" + jsonArray(tools) + ",\"resources\":" + jsonArray(resources) + ",\"templates\":"
					+ jsonArray(templates) + ",\"prompts\":" + jsonArray(prompts) + "}";
		}

		private static String contextFreeJson(McpAsyncServer server) {
			List<String> tools = server.listTools().map(Tool::name).collectList().block();
			List<String> resources = server.listResources().map(Resource::uri).collectList().block();
			List<String> templates = server.listResourceTemplates()
				.map(ResourceTemplate::uriTemplate)
				.collectList()
				.block();
			List<String> prompts = server.listPrompts().map(Prompt::name).collectList().block();
			return "{\"tools\":" + jsonArray(tools) + ",\"resources\":" + jsonArray(resources) + ",\"templates\":"
					+ jsonArray(templates) + ",\"prompts\":" + jsonArray(prompts) + "}";
		}

		private static void writeJson(HttpServletResponse response, String json) throws IOException {
			response.setStatus(HttpServletResponse.SC_OK);
			response.setContentType("application/json");
			response.setCharacterEncoding("UTF-8");
			response.getWriter().write(json);
		}

		private static String jsonArray(List<String> values) {
			return values.stream().map(AdminServlet::quote).collect(Collectors.joining(",", "[", "]"));
		}

		private static String quote(String value) {
			return "\"" + value.replace("\\", "\\\\").replace("\"", "\\\"") + "\"";
		}

	}

}