mod xray_propagator;

use crate::xray_propagator::XrayPropagator;
use lambda_runtime::{
    Error, LambdaEvent, Runtime, layers::OpenTelemetryLayer as OTELLayer, service_fn,
};
use opentelemetry::{
    Context, global,
    trace::{TraceContextExt, TracerProvider},
};
use opentelemetry_http::HeaderExtractor;
use serde_json::Value;
use std::env;
use tracing_subscriber::prelude::*;

async fn handler(_: LambdaEvent<Value>) -> Result<(), Error> {
    let context = global::get_text_map_propagator(|propagator| {
        println!("{:?}", propagator);
        let carrier = http::HeaderMap::new();
        let extractor = HeaderExtractor(&carrier);

        propagator.extract_with_context(&Context::current(), &extractor)
    });

    let x_amzn_trace_id = env::var("_X_AMZN_TRACE_ID").unwrap();
    println!("x_amzn_trace_id: {}", x_amzn_trace_id);

    println!(
        "trace_id pre attach: {}",
        Context::current().span().span_context().trace_id()
    );

    let _guard = context.attach();

    println!(
        "trace_id post attach: {}",
        Context::current().span().span_context().trace_id()
    );

    Ok(())
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    global::set_text_map_propagator(XrayPropagator::new());
    let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build();
    let tracer = tracer_provider.tracer("OpenTelemetryRustLambda");
    tracing_subscriber::registry()
        .with(tracing_opentelemetry::OpenTelemetryLayer::new(tracer))
        .init();

    Runtime::new(service_fn(handler))
        .layer(OTELLayer::new(|| {}))
        .run()
        .await
}

    fn extract_span_context(&self, extractor: &dyn Extractor) -> Option<SpanContext> {
        match extractor.get(AWS_XRAY_TRACE_HEADER) {
            None => {
                match env::var(AWS_XRAY_ENV) {
                    Ok(env) => span_context_from_str(env.as_str())
                    Err(_) => None,
                }
            }
            Some(header) => span_context_from_str(header.trim()),
        }
    }