// Copyright 2025 The NativeLink Authors. All rights reserved.

//

// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//    See LICENSE file for details

//

// 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.

use core::default::Default;

use std::collections::HashMap;

use std::env;

use std::sync::{Arc, OnceLock};

use base64::Engine;

use base64::prelude::BASE64_STANDARD_NO_PAD;

use ginepro::LoadBalancedChannel;

use hyper::http::Response;

use nativelink_error::{Code, ResultExt, make_err};

use nativelink_proto::build::bazel::remote::execution::v2::RequestMetadata;

use opentelemetry::propagation::TextMapCompositePropagator;

use opentelemetry::trace::{TraceContextExt, Tracer, TracerProvider};

use opentelemetry::{KeyValue, global};

use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;

use opentelemetry_http::HeaderExtractor;

use opentelemetry_otlp::{

    LogExporter, MetricExporter, Protocol, SpanExporter, WithExportConfig, WithTonicConfig,

};

use opentelemetry_sdk::Resource;

use opentelemetry_sdk::logs::SdkLoggerProvider;

use opentelemetry_sdk::metrics::SdkMeterProvider;

use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};

use opentelemetry_sdk::trace::SdkTracerProvider;

use opentelemetry_semantic_conventions::attribute::ENDUSER_ID;

use prost::Message;

use tracing::debug;

use tracing::metadata::LevelFilter;

use tracing_opentelemetry::{MetricsLayer, layer};

use tracing_subscriber::filter::Directive;

use tracing_subscriber::prelude::__tracing_subscriber_SubscriberExt;

use tracing_subscriber::util::SubscriberInitExt;

use tracing_subscriber::{EnvFilter, Layer, Registry, fmt, registry};

use uuid::Uuid;

use crate::metrics::record_rpc_served;

use crate::origin_event::{BAZEL_METADATA_KEY, request_metadata_to_baggage};

/// The OTLP "service.name" field for all nativelink services.

const NATIVELINK_SERVICE_NAME: &str = "nativelink";

// An `EnvFilter` to filter out non-nativelink information.

//

// See: https://github.com/open-telemetry/opentelemetry-rust/issues/2877

//

// Note that `EnvFilter` doesn't implement `clone`, so create a new one for

// each telemetry kind.

fn otlp_filter() -> EnvFilter {

    fn expect_parse(directive: &str) -> Directive {

        directive

            .parse()

            .unwrap_or_else(|_| panic!("Static directive '{directive}' failed to parse"))

    }

    EnvFilter::builder()

        .with_default_directive(LevelFilter::INFO.into())

        .from_env_lossy()

        .add_directive(expect_parse("hyper=off"))

        .add_directive(expect_parse("tonic=off"))

        .add_directive(expect_parse("h2=off"))

        .add_directive(expect_parse("reqwest=off"))

        .add_directive(expect_parse("tower=off"))

}

// Create a tracing layer intended for stdout printing.

//

// The output of this layer is configurable via the `NL_LOG` environment

// variable.

fn tracing_stdout_layer() -> impl Layer<Registry> {

    let nl_log_fmt = env::var("NL_LOG").unwrap_or_else(|_| "pretty".to_string());

    let stdout_filter = otlp_filter();

    match nl_log_fmt.as_str() {

        "compact" => fmt::layer()

            .compact()

            .with_timer(fmt::time::time())

            .with_filter(stdout_filter)

            .boxed(),

        "json" => fmt::layer()

            .json()

            .with_timer(fmt::time::time())

            .with_filter(stdout_filter)

            .boxed(),

        _ => fmt::layer()

            .pretty()

            .with_timer(fmt::time::time())

            .with_filter(stdout_filter)

            .boxed(),

    }

}

/// Initialize tracing with OpenTelemetry support.

///

/// # Errors

///

/// Returns `Err` if logging was already initialized or if the exporters can't

/// be initialized.

pub async fn init_tracing() -> Result<(), nativelink_error::Error> {

    static INITIALIZED: OnceLock<()> = OnceLock::new();

    if INITIALIZED.get().is_some() {

        return Err(make_err!(Code::Internal, "Logging already initialized"));

    }

    // We currently use a UUIDv4 for "service.instance.id" as per:

    // https://opentelemetry.io/docs/specs/semconv/attributes-registry/service/

    // This might change as we get a better understanding of its use cases in the

    // context of broader observability infrastructure.

    let resource = Resource::builder()

        .with_service_name(NATIVELINK_SERVICE_NAME)

        .with_attribute(KeyValue::new(

            "service.instance.id",

            Uuid::new_v4().to_string(),

        ))

        .build();

    let propagator = TextMapCompositePropagator::new(vec![

        Box::new(BaggagePropagator::new()),

        Box::new(TraceContextPropagator::new()),

    ]);

    global::set_text_map_propagator(propagator);

    let maybe_channel = maybe_load_balanced_channel().await;

    // Logs

    let mut log_exporter_builder = LogExporter::builder().with_tonic();

    if let Some(channel) = maybe_channel.clone() {

        log_exporter_builder = log_exporter_builder.with_channel(channel.into());

    }

    let otlp_log_layer = OpenTelemetryTracingBridge::new(

        &SdkLoggerProvider::builder()

            .with_resource(resource.clone())

            .with_batch_exporter(

                log_exporter_builder

                    .with_protocol(Protocol::Grpc)

                    .build()

                    .map_err(|e| make_err!(Code::Internal, "{e}"))

                    .err_tip(|| "While creating OpenTelemetry OTLP Log exporter")?,

            )

            .build(),

    )

    .with_filter(otlp_filter());

    // Traces

    let mut span_exporter_builder = SpanExporter::builder().with_tonic();

    if let Some(channel) = maybe_channel.clone() {

        span_exporter_builder = span_exporter_builder.with_channel(channel.into());

    }

    let otlp_trace_layer = layer()

        .with_tracer(

            SdkTracerProvider::builder()

                .with_resource(resource.clone())

                .with_batch_exporter(

                    span_exporter_builder

                        .with_protocol(Protocol::Grpc)

                        .build()

                        .map_err(|e| make_err!(Code::Internal, "{e}"))

                        .err_tip(|| "While creating OpenTelemetry OTLP Span exporter")?,

                )

                .build()

                .tracer(NATIVELINK_SERVICE_NAME),

        )

        .with_filter(otlp_filter());

    // Metrics

    let mut metric_exporter_builder = MetricExporter::builder().with_tonic();

    if let Some(channel) = maybe_channel {

        metric_exporter_builder = metric_exporter_builder.with_channel(channel.into());

    }

    let meter_provider = SdkMeterProvider::builder()

        .with_resource(resource)

        .with_periodic_exporter(

            metric_exporter_builder

                .with_protocol(Protocol::Grpc)

                .build()

                .map_err(|e| make_err!(Code::Internal, "{e}"))

                .err_tip(|| "While creating OpenTelemetry OTLP Metric exporter")?,

        )

        .build();

    global::set_meter_provider(meter_provider.clone());

    let otlp_metrics_layer = MetricsLayer::new(meter_provider).with_filter(otlp_filter());

    registry()

        .with(tracing_stdout_layer())

        .with(otlp_log_layer)

        .with(otlp_trace_layer)

        .with(otlp_metrics_layer)

        .init();

    INITIALIZED.set(()).unwrap_or(());

    Ok(())

}

/// Environment variable pointing OTLP exporters at a load-balanced endpoint.

pub const NL_OTEL_ENDPOINT: &str = "NL_OTEL_ENDPOINT";

/// Creates a load-balanced gRPC channel for the endpoint configured via

/// [`NL_OTEL_ENDPOINT`], or returns `None` if the variable isn't set.

///

/// Public so that integration tests can verify that exporters behave the

/// same with and without the load-balanced channel.

pub async fn maybe_load_balanced_channel() -> Option<LoadBalancedChannel> {

    match env::var(NL_OTEL_ENDPOINT) {

        Ok(endpoint) => {

            let url = Url::parse(endpoint.as_str())

                .map_err(|e| {

                    make_err!(Code::Internal, "Unable to parse endpoint {endpoint}: {e:?}")

                })

                .unwrap();

            let host = url

                .host()

                .err_tip(|| format!("Unable to get host from endpoint {endpoint}"))

                .unwrap();

            let port = url

                .port()

                .err_tip(|| format!("Unable to get port from endpoint {endpoint}"))

                .unwrap();

            Some(

                LoadBalancedChannel::builder((host.to_string(), port))

                    .channel()

                    .await

                    .map_err(|e| make_err!(Code::Internal, "Invalid hostname '{endpoint}': {e}"))

                    .unwrap(),

            )

        }

        Err(_) => None,

    }

}

/// This is the header that bazel sends when using the `--remote_header` flag.

/// TODO(palfrey): Bazel supports other headers, and we should optimize their usage.

const BAZEL_REQUESTMETADATA_HEADER: &str = "build.bazel.remote.execution.v2.requestmetadata-bin";

use opentelemetry::baggage::BaggageExt;

use opentelemetry::context::FutureExt;

use url::Url;

/// ASCII headers from an inbound client request, stored in the task context

/// so that outgoing upstream calls can forward them (e.g. JWT auth tokens).

#[derive(Clone, Debug, Default)]

pub struct ClientHeaders(pub Arc<HashMap<String, String>>);

#[derive(Debug, Clone)]

pub struct OtlpMiddleware<S> {

    inner: S,

    identity_required: bool,

}

impl<S> OtlpMiddleware<S> {

    const fn new(inner: S, identity_required: bool) -> Self {

        Self {

            inner,

            identity_required,

        }

    }

}

impl<S, ReqBody, ResBody> tower::Service<hyper::http::Request<ReqBody>> for OtlpMiddleware<S>

where

    S: tower::Service<hyper::http::Request<ReqBody>, Response = Response<ResBody>>

        + Clone

        + Send

        + 'static,

    S::Future: Send + 'static,

    ReqBody: core::fmt::Debug + Send + 'static,

    ResBody: From<String> + Send + 'static + Default,

{

    type Response = S::Response;

    type Error = S::Error;

    type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(

        &mut self,

        cx: &mut core::task::Context<'_>,

    ) -> core::task::Poll<Result<(), Self::Error>> {

        self.inner.poll_ready(cx)

    }

    fn call(&mut self, req: hyper::http::Request<ReqBody>) -> Self::Future {

        // We must take the current `inner` and not the clone.

        // See: <https://docs.rs/tower/latest/tower/trait.Service.html#be-careful-when-cloning-inner-services>

        let clone = self.inner.clone();

        let mut inner = core::mem::replace(&mut self.inner, clone);

        // Capture all ASCII-valued request headers before req is consumed, so

        // they can be forwarded to upstream services (e.g. JWT auth tokens).

        let client_headers = ClientHeaders(Arc::new(

            req.headers()

                .iter()

                .filter_map(|(name, value)| 
{2

                    value

                        .to_str()

                        .ok()

                        .map(|v| (name.as_str().to_lowercase(), v.to_string()))

                })

                .collect(),

        ));

        let parent_cx = global::get_text_map_propagator(|propagator| {

            propagator.extract(&HeaderExtractor(req.headers()))

        });

        let identity = parent_cx

            .baggage()

            .get(ENDUSER_ID)

            .map(|value| 
value.as_str()1
.
to_string1
())

            .unwrap_or_default();

        if identity.is_empty() {

            if self.identity_required {

                return Box::pin(async move {

                    Ok(tonic::Status::failed_precondition(

                        r"

NativeLink instance configured to require this OpenTelemetry Baggage header:

    `Baggage: enduser.id=YOUR_IDENTITY`

",

                    )

                    .into_http())

                });

            }

        } else {

            debug!("Baggage enduser.id: {identity}");

        }

        let tracer = global::tracer("origin_middleware");

        let span = tracer

            .span_builder("origin_request")

            .with_kind(opentelemetry::trace::SpanKind::Server)

            .start_with_context(&tracer, &parent_cx);

        let mut cx = parent_cx.with_span(span);

        if let Some(
bazel_header0
) = req.headers().get(BAZEL_REQUESTMETADATA_HEADER)

            && let Ok(decoded) = BASE64_STANDARD_NO_PAD.decode(bazel_header.as_bytes())

            && let Ok(metadata) = RequestMetadata::decode(decoded.as_slice())

        {

            let metadata_str = format!("{metadata:?}");

            debug!("Baggage Bazel request metadata: {metadata_str}");

            cx = cx.with_baggage(vec![

                KeyValue::new(BAZEL_METADATA_KEY, request_metadata_to_baggage(&metadata)),

                KeyValue::new(ENDUSER_ID, identity),

            ]);

        }

        let cx = cx.with_value(client_headers);

        // Rate, errors and latency for every gRPC service, taken here because

        // this layer already wraps all of them.

        let route = req.uri().path().to_string();

        let started = std::time::Instant::now();

        Box::pin(async move {

            let result = inner.call(req).with_context(cx).await;

            if let Ok(response) = &result {

                record_rpc_served(

                    &route,

                    grpc_status_of(response.headers()),

                    started.elapsed().as_secs_f64(),

                );

            
}0

            result

        })

    }

}

/// Reads the gRPC status from a response.

///

/// Tonic puts `grpc-status` in the headers when it fails before streaming and

/// in the trailers otherwise. Trailers are not available here without

/// consuming the body, so a call that got far enough to stream is reported as

/// OK. Transport-level failures never reach this point at all.

fn grpc_status_of(headers: &hyper::http::HeaderMap) -> i32 {

    headers

        .get("grpc-status")

        .and_then(|value| 
value0
.
to_str0
().
ok0
())

        .and_then(|value| 
value0
.
parse0
::<i32>().
ok0
())

        .unwrap_or(0)

}

#[derive(Debug, Clone, Copy)]

pub struct OtlpLayer {

    identity_required: bool,

}

impl OtlpLayer {

    pub const fn new(identity_required: bool) -> Self {

        Self { identity_required }

    }

}

impl<S> tower::Layer<S> for OtlpLayer {

    type Service = OtlpMiddleware<S>;

    fn layer(&self, service: S) -> Self::Service {

        OtlpMiddleware::new(service, self.identity_required)

    }

}

#[cfg(test)]

mod tests {

    use nativelink_macro::nativelink_test;

    use serial_test::serial;

    use super::*;

    // ginepro's default resolver (hickory-dns) reads /etc/resolv.conf, which

    // doesn't exist in sandboxed environments (e.g. Nix builds).

    #[cfg(unix)]

    fn dns_configured() -> bool {

        std::path::Path::new("/etc/resolv.conf").exists()

    }

    #[cfg(not(unix))]

    const fn dns_configured() -> bool {

        true

    }

    // Env vars are process-global, so concurrent writes would be a data race

    // on Unix.  `#[serial(env)]` serializes all env-mutating tests.

    #[serial(env)]

    #[nativelink_test("crate")]

    async fn channel_absent_when_env_not_set() {

        // SAFETY: `#[serial(env)]` serializes all env-var writes across tests.

        unsafe { env::remove_var(NL_OTEL_ENDPOINT) };

        assert!(

            maybe_load_balanced_channel().await.is_none(),

            "Expected None when {NL_OTEL_ENDPOINT} is unset"

        );

    }

    #[serial(env)]

    #[nativelink_test("crate")]

    async fn channel_present_when_valid_endpoint_set() {

        if !dns_configured() {

            eprintln!(

                "Skipping channel_present_when_valid_endpoint_set: no DNS configuration \

                 available (e.g. sandboxed Nix build)"

            );

            return;

        }

        // SAFETY: `#[serial(env)]` serializes all env-var writes across tests.

        unsafe { env::set_var(NL_OTEL_ENDPOINT, "http://localhost:4317") };

        let result = maybe_load_balanced_channel().await;

        // SAFETY: `#[serial(env)]` serializes all env-var writes across tests.

        unsafe { env::remove_var(NL_OTEL_ENDPOINT) };

        assert!(

            result.is_some(),

            "Expected Some(channel) when {NL_OTEL_ENDPOINT} points to a valid URL"

        );

    }

}