use std::collections::HashMap;

use opentelemetry::propagation::{
    Extractor, Injector, TextMapCompositePropagator, TextMapPropagator,
};
use opentelemetry::trace::{SpanContext, TraceContextExt};
use opentelemetry::{Context, global};
use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};
use tracing_opentelemetry::OpenTelemetrySpanExt;

pub const TRACEPARENT: &str = "traceparent";
pub const TRACESTATE: &str = "tracestate";
pub const BAGGAGE: &str = "baggage";

#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct TextMap(pub HashMap<String, String>);

impl TextMap {
    pub fn new() -> Self {
        Self(HashMap::new())
    }

    pub fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).map(String::as_str)
    }

    pub fn insert(&mut self, key: impl Into<String>, value: impl Into<String>) {
        self.0.insert(key.into(), value.into());
    }
}

impl Injector for TextMap {
    fn set(&mut self, key: &str, value: String) {
        self.0.insert(key.to_owned(), value);
    }
}

impl Extractor for TextMap {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).map(String::as_str)
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(String::as_str).collect()
    }
}

pub fn install() {
    global::set_text_map_propagator(TextMapCompositePropagator::new(vec![
        Box::new(TraceContextPropagator::new()),
        Box::new(BaggagePropagator::new()),
    ]));
}

pub fn current_context() -> Context {
    let context = tracing::Span::current().context();
    if context.span().span_context().is_valid() {
        return context;
    }
    Context::current()
}

pub fn inject(into: &mut dyn Injector) {
    let context = current_context();
    global::get_text_map_propagator(|propagator| {
        propagator.inject_context(&context, &mut *into);
    });
}

pub fn inject_map(into: &mut TextMap) {
    inject(into);
}

pub fn extract(from: &dyn Extractor) -> Context {
    global::get_text_map_propagator(|propagator| propagator.extract(from))
}

pub fn traceparent() -> Option<String> {
    let mut map = TextMap::new();
    inject_map(&mut map);
    map.0.get(TRACEPARENT).cloned()
}

pub fn context_from_traceparent(traceparent: &str) -> Context {
    let mut map = TextMap::new();
    map.insert(TRACEPARENT, traceparent);
    TraceContextPropagator::new().extract(&map)
}

pub fn current() -> Option<SpanContext> {
    let span_context = current_context().span().span_context().clone();
    span_context.is_valid().then_some(span_context)
}

pub fn set_parent(span: &tracing::Span, parent: &Context) -> bool {
    span.set_parent(parent.clone()).is_ok()
}

pub fn set_parent_from_traceparent(span: &tracing::Span, traceparent: Option<&str>) -> bool {
    match traceparent.filter(|value| !value.is_empty()) {
        Some(value) => set_parent(span, &context_from_traceparent(value)),
        None => false,
    }
}