/*
 * 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.tracing.otel;

import static org.apache.juneau.commons.utils.Shorts.*;
import static org.apache.juneau.commons.utils.StringUtils.*;

import java.util.*;

import org.apache.juneau.http.tracing.TraceContextCarrier;
import org.apache.juneau.rest.server.*;
import org.apache.juneau.rest.server.tracing.*;
import org.apache.juneau.rest.server.tracing.Scope;

import io.opentelemetry.api.*;
import io.opentelemetry.api.common.*;
import io.opentelemetry.api.trace.*;
import io.opentelemetry.context.*;
import io.opentelemetry.context.propagation.*;

/**
 * {@link TracerHook} bridge that opens an OpenTelemetry {@link Span} for every {@code @RestOp} call.
 *
 * <h5 class='topic'>Span shape</h5>
 *
 * <p>
 * Each request becomes a single span with kind {@link SpanKind#SERVER}, named after the HTTP method
 * (e.g. {@code "GET"}, {@code "POST"}) per the OpenTelemetry HTTP semantic conventions. Three OTel
 * HTTP semantic-convention attributes are populated:
 * <ul>
 *    <li><b>{@code http.request.method}</b> &mdash; uppercased HTTP method.
 *    <li><b>{@code http.response.status_code}</b> &mdash; HTTP response status (set right before the
 *       span is closed, so error paths get the framework's resolved status).
 *    <li><b>{@code http.route}</b> &mdash; the {@code @RestOp} path template (e.g.
 *       {@code "/users/{id}"}), not the raw concrete URI. Keeps trace cardinality bounded for the same
 *       reason the Micrometer bridge uses the template as its {@code uri} tag.
 * </ul>
 *
 * <p>
 * On the exception path, the throwable is recorded via {@link Span#recordException(Throwable)}, the
 * span status is set to {@link StatusCode#ERROR}, and an {@code exception.type} attribute carries the
 * exception's simple-name.
 *
 * <h5 class='topic'>W3C trace context propagation</h5>
 *
 * <p>
 * Incoming {@code traceparent} / {@code tracestate} request headers are extracted via the configured
 * {@link TextMapPropagator} (default: {@link io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator}) so the server span continues
 * a caller-supplied distributed trace. Downstream HTTP calls made from inside the handler will pick
 * up the active span's context automatically when the consumer's outbound HTTP client honors the OTel
 * {@link Context#current() current context} (the standard OTel client instrumentations do).
 *
 * <h5 class='topic'>Wiring</h5>
 *
 * <p class='bjava'>
 *    <ja>@Configuration</ja>
 *    <jk>public class</jk> ObservabilityConfig {
 *
 *       <jc>// Option A: rely on the JVM-wide GlobalOpenTelemetry (recommended).</jc>
 *       <ja>@Bean</ja>
 *       <jk>public</jk> TracerHook tracerHook() {
 *          <jk>return new</jk> OtelTracerHook();
 *       }
 *
 *       <jc>// Option B: pass a specific OpenTelemetry instance (e.g. for tests or multi-tenant setups).</jc>
 *       <ja>@Bean</ja>
 *       <jk>public</jk> TracerHook tracerHook(OpenTelemetry <jv>otel</jv>) {
 *          <jk>return new</jk> OtelTracerHook(<jv>otel</jv>);
 *       }
 *    }
 * </p>
 *
 * <h5 class='section'>See Also:</h5><ul>
 *    <li class='jc'>{@link TracerHook}
 *    <li class='link'><a class="doclink" href="https://juneau.apache.org/docs/topics/RestServerObservability">Observability — Micrometer + OpenTelemetry</a>
 *    <li class='link'><a class="doclink" href="https://opentelemetry.io/docs/specs/semconv/http/http-spans/">OpenTelemetry HTTP semantic conventions</a>
 *    <li class='link'><a class="doclink" href="https://www.w3.org/TR/trace-context/">W3C Trace Context</a>
 * </ul>
 *
 * @since 10.0.0
 */
@SuppressWarnings({
   "java:S1192", // Duplicated literals (argument/property names) read more clearly inline than as constants
   "resource" // Span / Scope returned to caller for try-with-resources management; Eclipse JDT @Owning warning is by design.
})
public class OtelTracerHook implements TracerHook {

   /** Default tracer instrumentation scope name. */
   public static final String DEFAULT_INSTRUMENTATION_NAME = "org.apache.juneau.rest.server";

   /** OTel HTTP semantic-convention attribute key {@code http.request.method}. */
   public static final AttributeKey<String> ATTR_HTTP_REQUEST_METHOD = AttributeKey.stringKey("http.request.method");

   /** OTel HTTP semantic-convention attribute key {@code http.response.status_code}. */
   public static final AttributeKey<Long> ATTR_HTTP_RESPONSE_STATUS_CODE = AttributeKey.longKey("http.response.status_code");

   /** OTel HTTP semantic-convention attribute key {@code http.route}. */
   public static final AttributeKey<String> ATTR_HTTP_ROUTE = AttributeKey.stringKey("http.route");

   /** OTel attribute key {@code exception.type} used when recording a thrown exception. */
   public static final AttributeKey<String> ATTR_EXCEPTION_TYPE = AttributeKey.stringKey("exception.type");

   private final Tracer tracer;
   private final TextMapPropagator propagator;

   /**
    * Constructor using {@link GlobalOpenTelemetry#get()} as the {@link OpenTelemetry} source.
    *
    * <p>
    * Convenient when the deployment installs a single global {@link OpenTelemetry} (the dominant
    * pattern with the OTel auto-instrumentation agent). Equivalent to
    * {@code new OtelTracerHook(GlobalOpenTelemetry.get())}.
    */
   public OtelTracerHook() {
      this(GlobalOpenTelemetry.get());
   }

   /**
    * Constructor with an explicit {@link OpenTelemetry} instance.
    *
    * <p>
    * The instance is used to obtain both the {@link Tracer} (named
    * {@value #DEFAULT_INSTRUMENTATION_NAME}) and the {@link TextMapPropagator} used for incoming /
    * outgoing W3C trace-context header transfer.
    *
    * @param openTelemetry The {@link OpenTelemetry} instance. Must not be <jk>null</jk>.
    */
   public OtelTracerHook(OpenTelemetry openTelemetry) {
      reqnn("openTelemetry", openTelemetry);
      this.tracer = openTelemetry.getTracer(DEFAULT_INSTRUMENTATION_NAME);
      this.propagator = openTelemetry.getPropagators().getTextMapPropagator();
   }

   /**
    * Constructor with an explicit {@link Tracer} and {@link TextMapPropagator}.
    *
    * <p>
    * Useful for tests where in-memory tracer / propagator instances are wired directly without going
    * through a full {@link OpenTelemetry} SDK assembly.
    *
    * @param tracer The {@link Tracer} used to build per-request spans. Must not be <jk>null</jk>.
    * @param propagator The {@link TextMapPropagator} used for incoming / outgoing W3C trace-context
    *    header transfer. Must not be <jk>null</jk>.
    */
   public OtelTracerHook(Tracer tracer, TextMapPropagator propagator) {
      this.tracer = reqnn("tracer", tracer);
      this.propagator = reqnn("propagator", propagator);
   }

   @Override /* TracerHook */
   public Scope startSpan(RestRequest request) {
      return startSpan(request, null, TraceOperation.DEFAULT);
   }

   @Override /* TracerHook */
   public Scope startSpan(RestRequest request, TraceContextCarrier carrier, TraceOperation operation) {
      Context extracted = extractContext(request, carrier);

      String httpMethod = defaultIfBlank(request.getMethod(), "UNKNOWN");
      String route = resolveRoute(request);
      String spanName = defaultIfBlank(operation.getSpanName(), httpMethod);

      var spanBuilder = tracer.spanBuilder(spanName)
         .setSpanKind(SpanKind.SERVER)
         .setParent(extracted)
         .setAttribute(ATTR_HTTP_REQUEST_METHOD, httpMethod);
      if (! route.isEmpty())
         spanBuilder.setAttribute(ATTR_HTTP_ROUTE, route);
      operation.getAttributes().forEach((k, v) -> spanBuilder.setAttribute(AttributeKey.stringKey(k), v));

      Span span = spanBuilder.startSpan();
      Context spanContext = extracted.with(span);

      // Render the W3C trace-context headers (and baggage, when a baggage propagator is configured)
      // from the server-started span's context (the only point where it's reliably active) and stash
      // them as request attributes for TraceContextResponseProcessor / a dated MCP result to read
      // later. The configured propagator only injects traceparent when the span context is valid, and
      // tracestate/baggage only when non-empty, so both guards are handled here for free.
      stashTraceContext(request, spanContext);

      io.opentelemetry.context.Scope otelScope = spanContext.makeCurrent();

      return new OtelScope(span, otelScope);
   }

   /**
    * Builds the effective parent {@link Context} for a new server span, giving a non-HTTP
    * {@link TraceContextCarrier} &mdash; recognized by an active {@code TraceContextExtractor} among
    * the resolved operation arguments (for example an MCP request's {@code params._meta}) &mdash;
    * precedence over HTTP request headers.
    *
    * <p>
    * Extraction runs in up to two passes against the same configured {@link #propagator}: first
    * against ordinary HTTP headers (via {@link RestRequestTextMapGetter}), then &mdash; only when
    * <c>carrier</c> is non-<jk>null</jk> &mdash; a second pass against <c>carrier</c> (via
    * {@link TraceContextCarrierTextMapGetter}), seeded with the first pass's result. Per the
    * {@link TextMapPropagator} contract, extraction that finds nothing usable returns the given
    * {@link Context} unmodified, so an explicit carrier value wins while an absent or invalid one
    * falls back to whatever the HTTP pass already produced. This applies independently to the
    * {@code traceparent}/{@code tracestate} pair (extracted together as one {@code SpanContext}) and
    * to {@code baggage} (extracted separately whenever the configured propagator includes a baggage
    * propagator).
    *
    * @param request The in-flight request; source of the HTTP-header fallback. Never <jk>null</jk>.
    * @param carrier The extractor-recognized carrier, or <jk>null</jk> if none was recognized (HTTP
    *    headers only).
    * @return The effective parent {@link Context}. Never <jk>null</jk>.
    */
   private Context extractContext(RestRequest request, TraceContextCarrier carrier) {
      Context httpContext = propagator.extract(Context.current(), request, RestRequestTextMapGetter.INSTANCE);
      if (carrier == null)
         return httpContext;
      return propagator.extract(httpContext, carrier, TraceContextCarrierTextMapGetter.INSTANCE);
   }

   @Override /* TracerHook */
   public Scope startSpan(String spanName) {
      // Custom (non-request) observation: open an INTERNAL span parented to the current context
      // (so it nests under an active request span when one exists), with no HTTP attributes.
      Span span = tracer.spanBuilder(spanName)
         .setSpanKind(SpanKind.INTERNAL)
         .setParent(Context.current())
         .startSpan();
      io.opentelemetry.context.Scope otelScope = Context.current().with(span).makeCurrent();
      return new OtelScope(span, otelScope);
   }

   private void stashTraceContext(RestRequest request, Context spanContext) {
      var carrier = new HashMap<String,String>(4);
      propagator.inject(spanContext, carrier, MAP_SETTER);
      var traceparent = carrier.get("traceparent");
      if (nn(traceparent))
         request.setAttribute(TraceContextResponseProcessor.ATTR_TRACEPARENT, traceparent);
      var tracestate = carrier.get("tracestate");
      if (nn(tracestate) && ! tracestate.isEmpty())
         request.setAttribute(TraceContextResponseProcessor.ATTR_TRACESTATE, tracestate);
      var baggage = carrier.get("baggage");
      if (nn(baggage) && ! baggage.isEmpty())
         request.setAttribute(TraceContextResponseProcessor.ATTR_BAGGAGE, baggage);
   }

   private static final TextMapSetter<HashMap<String,String>> MAP_SETTER = (carrier, key, value) -> {
      if (nn(carrier))
         carrier.put(key, value);
   };

   @Override /* TracerHook */
   public void inject(TraceContextCarrier carrier) {
      propagator.inject(Context.current(), carrier, TraceContextCarrierTextMapSetter.INSTANCE);
   }

   /**
    * Returns the {@link Tracer} this hook publishes spans to.
    *
    * @return The {@link Tracer}.
    */
   public Tracer getTracer() { return tracer; }

   /**
    * Returns the {@link TextMapPropagator} used to extract W3C trace-context headers from
    * incoming requests.
    *
    * @return The {@link TextMapPropagator}.
    */
   public TextMapPropagator getPropagator() { return propagator; }

   private static String resolveRoute(RestRequest request) {
      try {
         var opCtx = request.getOpContext();
         return opCtx == null ? "" : defaultIfBlank(opCtx.getPathPattern(), "");
      } catch (RuntimeException e) {
         // Defensive: opCtx.getPathPattern() reaches into the path-matcher array; protect against NPE / array bounds in odd routing setups.
         return "";
      }
   }

   private static String defaultIfBlank(String value, String fallback) {
      return isEmpty(value) ? fallback : value;
   }

   private static final class OtelScope implements Scope {

      private final Span span;
      private final io.opentelemetry.context.Scope scope;

      OtelScope(Span span, io.opentelemetry.context.Scope scope) {
         this.span = span;
         this.scope = scope;
      }

      @Override /* Scope */
      public void setStatusCode(int statusCode) {
         span.setAttribute(ATTR_HTTP_RESPONSE_STATUS_CODE, (long) statusCode);
         if (statusCode >= 500)
            span.setStatus(StatusCode.ERROR);
      }

      @Override /* Scope */
      public void setError(Throwable error) {
         if (nn(error)) {
            span.recordException(error);
            span.setAttribute(ATTR_EXCEPTION_TYPE, cns(error));
            span.setStatus(StatusCode.ERROR);
         }
      }

      @Override /* Scope */
      public void close() {
         try {
            scope.close();
         } finally {
            span.end();
         }
      }

      @Override /* Scope */
      public void recordRpcError(int code, String message) {
         span.setAttribute(TraceOperation.ATTR_RPC_RESPONSE_STATUS_CODE, String.valueOf(code));
         span.setAttribute(TraceOperation.ATTR_ERROR_TYPE, errorType(code));
         span.setStatus(StatusCode.ERROR, message == null ? "" : message);
      }

      /**
       * Maps a JSON-RPC 2.0 error code to a low-cardinality {@code error.type} category.
       *
       * <p>
       * The five codes reserved by the JSON-RPC 2.0 specification map to their exact category name;
       * any code in the specification's reserved-for-implementation-defined-server-errors range
       * ({@code -32000} to {@code -32099}, inclusive) maps to the generic {@code "server_error"}
       * category. An application-defined code outside every reserved range falls back to its own
       * decimal string so the attribute stays populated (if not maximally low-cardinality) rather than
       * silently empty.
       *
       * @param code The JSON-RPC error code (e.g. {@code -32601}).
       * @return The {@code error.type} category. Never <jk>null</jk>.
       */
      private static String errorType(int code) {
         switch (code) {
            case -32700: return "parse_error";
            case -32600: return "invalid_request";
            case -32601: return "method_not_found";
            case -32602: return "invalid_params";
            case -32603: return "internal_error";
            default:
               return (code <= -32000 && code >= -32099) ? "server_error" : String.valueOf(code);
         }
      }
   }
}