/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.juneau.rest.server.mcp;

import java.util.concurrent.atomic.*;

import org.apache.juneau.bean.jsonrpc.*;
import org.apache.juneau.commons.inject.*;
import org.apache.juneau.http.Content;
import org.apache.juneau.marshall.serializer.*;
import org.apache.juneau.rest.server.*;
import org.apache.juneau.rest.server.health.*;
import org.apache.juneau.rest.server.servlet.*;

/**
 * Drop-in {@link BasicRestServlet} subclass that exposes a single MCP JSON-RPC endpoint at {@code POST /}.
 *
 * <p>
 * Subclasses provide their {@link McpServerConfig} by overriding {@link #createMcpConfig()} (called once at
 * {@link #init() init} time), or by registering an {@link McpServerConfig} bean in their {@code RestContext}
 * bean store.
 *
 * <h5 class='section'>Example:</h5>
 * <pre>
 * @Rest(path="/mcp")
 * public class MyMcpServlet extends org.apache.juneau.rest.server.mcp.v20250618.McpRestServlet {
 *     @Override
 *     protected McpServerConfig createMcpConfig() {
 *         return new McpServerConfig()
 *             .setName("my-server").setVersion("1.0.0")
 *             .addTool(new MyEchoTool());
 *     }
 * }
 * </pre>
 *
 * <p>
 * The servlet enables {@code addBeanTypes} on its serializer so that polymorphic wire types a
 * revision emits (content blocks, resource payloads) are tagged with their {@code type}
 * discriminator, and disables URI resolution ({@code uriResolution="NONE"}) so that JSON-Schema
 * JSON-Pointer {@code $ref} values (e.g. {@code "#/$defs/text"}) and absolute URIs with fragments
 * are emitted verbatim rather than rewritten against the servlet path.
 *
 * @serial exclude
 */
@Rest
@SerializerConfig(addBeanTypes = "true", uriResolution = "NONE")
public abstract class AbstractMcpRestServlet extends BasicRestServlet implements ReadinessStateAware {
   private static final long serialVersionUID = 1L;

   private final transient AtomicReference<McpServerConfig> config = new AtomicReference<>();

   /**
    * Returns the {@link McpServerConfig} backing this servlet.
    *
    * <p>
    * The default implementation lazily caches the value of {@link #createMcpConfig()}; subclasses needing
    * per-request reconfiguration may override.
    *
    * <p>
    * This is invoked on every request, so the cache is published via a lock-free {@link AtomicReference#compareAndSet}
    * rather than a {@code synchronized} block, avoiding monitor contention on the hot path. Under a first-access
    * race, {@link #createMcpConfig()} may be invoked by more than one thread, but only one resulting instance is
    * ever published to {@link #config} and returned to any caller; the others are discarded. This is safe as long
    * as {@link #createMcpConfig()} is side-effect-free and idempotent (i.e. simply constructs and returns a new
    * config), which is the documented contract of that method.
    *
    * @return The config. Never {@code null}.
    */
   public McpServerConfig getMcpConfig() {
      var c = config.get();
      if (c == null) {
         var nc = createMcpConfig();
         if (nc == null)
            throw new IllegalStateException("createMcpConfig() returned null");
         c = config.compareAndSet(null, nc) ? nc : config.get(); // HTT: the CAS-loses branch requires a genuine concurrent first-access race; untestable deterministically.
      }
      return c;
   }

   /**
    * Factory method for the {@link McpServerConfig} backing this servlet.
    *
    * <p>
    * Subclasses must override this method to declare their tools, prompts, and resources.
    *
    * @return A non-{@code null} server config.
    */
   protected abstract McpServerConfig createMcpConfig();

   /**
    * The MCP protocol revision this servlet speaks.
    *
    * <p>
    * Supplied by the revision-specific subclass a consumer extends (for example
    * {@code org.apache.juneau.rest.server.mcp.v20250618.McpRestServlet}); consumers do not normally implement this themselves.
    *
    * @return The bound revision. Never {@code null}.
    */
   protected abstract McpRevision revision();

   /**
    * Optional subscription broker backing this servlet's {@code subscriptions/listen} support.
    *
    * <p>
    * Returning <jk>null</jk> (the default) means this servlet does not support subscriptions: no
    * {@link McpSubscriptionBroker} or {@link McpSubscriptions} bean is added to the request-scoped
    * {@link BeanStore} that {@link #handleMcp} builds, so any handler or dispatch branch that looks one up
    * via {@code ctx.getBean(McpSubscriptionBroker.class)} sees {@link java.util.Optional#empty()}.
    *
    * <p>
    * A revision binding that implements {@code subscriptions/listen} (for example
    * {@code org.apache.juneau.rest.server.mcp.v20260728.McpRestServlet}) overrides this to return a
    * memoized broker instance so every request sees the same live subscription registry.
    *
    * @return The subscription broker, or <jk>null</jk> if this servlet does not support subscriptions.
    */
   protected McpSubscriptionBroker getSubscriptionBroker() {
      return null;
   }

   /**
    * Closes this servlet's open subscription streams as soon as the embedded server begins shutting down.
    *
    * <p>
    * A {@code subscriptions/listen} stream never finishes on its own, so without this the server's graceful
    * drain would wait out its entire stop timeout (30s by default) for every connected subscriber before it
    * could stop.  Closing the subscriptions makes each stream send its terminal frame and complete, letting
    * the drain finish immediately.
    *
    * @param state The per-service readiness state published by the embedded-server lifecycle component.
    */
   @SuppressWarnings({
      "java:S2654" // Not a constructor escape: the lifecycle calls this on a fully constructed servlet, and the callback only closes the broker.
   })
   @Override /* ReadinessStateAware */
   public void acceptReadinessState(ReadinessState state) {
      state.onOutOfService(this, () -> {
         var broker = getSubscriptionBroker();
         if (broker != null)
            broker.closeAll();
      });
   }

   /**
    * MCP JSON-RPC endpoint.
    *
    * @param req The parsed JSON-RPC request envelope.
    * @param restReq The HTTP request (used to surface the bean store to handlers).
    * @return The JSON-RPC response, or {@code null} for notifications (which the framework renders as an
    *         empty body).
    */
   @SuppressWarnings({
      "resource" // Request-scoped scratch BasicBeanStore; lifetime is bounded by this handler invocation, no foreign resources are captured.
   })
   @RestPost(path = "/")
   public Object handleMcp(@Content JsonRpcRequest req, RestRequest restReq) {
      var bs = new BasicBeanStore(restReq.getContext().getBeanStore())
         .addBean(RestRequest.class, restReq);
      var broker = getSubscriptionBroker();
      if (broker != null)
         bs.addBean(McpSubscriptionBroker.class, broker).addBean(McpSubscriptions.class, broker);
      var exchange = new McpExchange(req, n -> restReq.getHeaderParam(n).asString().orElse(null));
      var result = revision().dispatch(exchange, getMcpConfig(), bs);
      if (result instanceof McpResponseResult r)
         return r.response();
      if (result instanceof McpStreamResult r)
         return r.stream();
      return null;
   }
}