// Copyright 2025 Jonas Kruckenberg
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use alloc::vec::Vec;
use core::cell::RefCell;
use core::sync::atomic;
use core::sync::atomic::{AtomicUsize, Ordering};
use cpu_local::collection::CpuLocal;
use ksharded_slab::Pool;
use ksharded_slab::pool::Ref;
use tracing_core::field::FieldSet;
use tracing_core::span::{Attributes, Current, Id, Record};
use tracing_core::{Collect, Dispatch, Event, Interest, Metadata, dispatch};

/// A shared, reusable store for spans.
///
/// This registry is implemented using a [lock-free sharded slab][slab], and is
/// highly optimized for concurrent access.
///
/// # Span ID Generation
///
/// Span IDs are not globally unique, but the registry ensures that
/// no two currently active spans have the same ID within a process.
///
/// One of the primary responsibilities of the registry is to generate [span
/// IDs]. Therefore, it's important for other code that interacts with the
/// registry, such as subscribers, to understand the guarantees of the
/// span IDs that are generated.
///
/// The registry's span IDs are guaranteed to be unique **at a given point
/// in time**. This means that an active span will never be assigned the
/// same ID as another **currently active** span. However, the registry
/// **will** eventually reuse the IDs of [closed] spans, although an ID
/// will never be reassigned immediately after a span has closed.
///
/// Spans are not [considered closed] by the `Registry` until *every*
/// [`Span`] reference with that ID has been dropped.
///
/// Thus: span IDs generated by the registry should be considered unique
/// only at a given point in time, and only relative to other spans
/// generated by the same process. Two spans with the same ID will not exist
/// in the same process concurrently. However, if historical span data is
/// being stored, the same ID may occur for multiple spans times in that
/// data. If spans must be uniquely identified in historical data, the user
/// code storing this data must assign its own unique identifiers to those
/// spans. A counter is generally sufficient for this.
///
/// Similarly, span IDs generated by the registry are not unique outside of
/// a given process. Distributed tracing systems may require identifiers
/// that are unique across multiple processes on multiple machines (for
/// example, [OpenTelemetry's `SpanId`s and `TraceId`s][opentelemetry]). `tracing` span
/// IDs generated by the registry should **not** be used for this purpose.
/// Instead, code which integrates with a distributed tracing system should
/// generate and propagate its own IDs according to the rules specified by
/// the distributed tracing system. These IDs can be associated with
/// `tracing` spans using [fields] and/or [stored span data].
///
/// [span IDs]: https://docs.rs/tracing-core/latest/tracing_core/span/struct.Id.html
/// [slab]: https://docs.rs/crate/sharded-slab/
/// [closed]: https://docs.rs/tracing/latest/tracing/span/index.html#closing-spans
/// [considered closed]: https://docs.rs/tracing-core/latest/tracing_core/subscriber/trait.Subscriber.html#method.try_close
/// [`Span`]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html
/// [opentelemetry]: https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/api.md#spancontext
/// [fields]: https://docs.rs/tracing-core/latest/tracing-core/field/index.html
#[derive(Debug)]
pub struct Registry {
    spans: Pool<DataInner>,
    current_spans: CpuLocal<RefCell<SpanStack>>,
}

/// Span data stored in a [`Registry`].
///
/// The registry stores well-known data defined by tracing: span relationships,
/// metadata and reference counts.
#[derive(Debug)]
pub struct Data<'a> {
    /// Immutable reference to the pooled `DataInner` entry.
    inner: Ref<'a, DataInner>,
}

/// Stored data associated with a span.
///
/// This type is pooled using `sharded_slab::Pool`; when a span is dropped, the
/// `DataInner` entry at that span's slab index is cleared in place and reused
/// by a future span. Thus, the `Default` and `sharded_slab::Clear`
/// implementations for this type are load-bearing.
#[derive(Debug)]
struct DataInner {
    metadata: &'static Metadata<'static>,
    parent: Option<Id>,
    ref_count: AtomicUsize,
}

/// A reference to [`Data`] and the associated [`Registry`].
///
/// This type implements all the same methods as [`Data`], and provides
#[derive(Debug)]
pub struct SpanRef<'a> {
    registry: &'a Registry,
    data: Data<'a>,
}

/// An iterator over the parents of a span, ordered from leaf to root.
///
/// This is returned by the [`SpanRef::scope`] method.
#[derive(Debug)]
pub struct Scope<'a> {
    registry: &'a Registry,
    next: Option<Id>,
}

#[derive(Debug)]
struct ContextId {
    id: Id,
    duplicate: bool,
}

/// `SpanStack` tracks what spans are currently executing on a thread-local basis.
///
/// A "separate current span" for each thread is a semantic choice, as each span
/// can be executing in a different thread.
#[derive(Debug, Default)]
pub(crate) struct SpanStack {
    stack: Vec<ContextId>,
}

// === impl Registry ===

impl Default for Registry {
    fn default() -> Self {
        Self {
            spans: Pool::new(),
            current_spans: CpuLocal::new(),
        }
    }
}

impl Registry {
    fn get(&self, id: &Id) -> Option<Ref<'_, DataInner>> {
        self.spans.get(id_to_idx(id))
    }

    fn span_data(&self, id: &Id) -> Option<Data> {
        let inner = self.get(id)?;
        Some(Data { inner })
    }

    fn span(&self, id: &Id) -> Option<SpanRef>
    where
        Self: Sized,
    {
        let data = self.span_data(id)?;
        Some(SpanRef {
            registry: self,
            data,
        })
    }
}

impl Collect for Registry {
    fn register_callsite(&self, _metadata: &'static Metadata<'static>) -> Interest {
        unreachable!("Registry::register_callsite should never be called")
    }

    fn enabled(&self, _metadata: &Metadata<'_>) -> bool {
        unreachable!("Registry::enabled should never be called")
    }

    fn new_span(&self, attrs: &Attributes<'_>) -> Id {
        let parent = if attrs.is_root() {
            None
        } else if attrs.is_contextual() {
            self.current_span().id().map(|id| self.clone_span(id))
        } else {
            attrs.parent().map(|id| self.clone_span(id))
        };

        let id = self
            .spans
            // Check out a `DataInner` entry from the pool for the new span. If
            // there are free entries already allocated in the pool, this will
            // preferentially reuse one; otherwise, a new `DataInner` is
            // allocated and added to the pool.
            .create_with(|data| {
                data.metadata = attrs.metadata();
                data.parent = parent;

                let refs = data.ref_count.get_mut();
                debug_assert_eq!(*refs, 0);
                *refs = 1;
            })
            .expect("Unable to allocate another span");
        idx_to_id(id)
    }

    fn record(&self, _span: &Id, _values: &Record<'_>) {}

    fn record_follows_from(&self, _span: &Id, _follows: &Id) {}

    fn event(&self, _event: &Event<'_>) {}

    fn enter(&self, id: &Id) {
        if self
            .current_spans
            .get_or_default()
            .borrow_mut()
            .push(id.clone())
        {
            self.clone_span(id);
        }
    }

    fn exit(&self, id: &Id) {
        if let Some(spans) = self.current_spans.get() {
            if spans.borrow_mut().pop(id) {
                dispatch::get_default(|dispatch| dispatch.try_close(id.clone()));
            }
        }
    }

    fn clone_span(&self, id: &Id) -> Id {
        let span = self
            .get(id)
            .unwrap_or_else(|| panic!("tried to clone {:?}, but no span exists with that ID", id));
        // Like `std::sync::Arc`, adds to the ref count (on clone) don't require
        // a strong ordering; if we call` clone_span`, the reference count must
        // always at least 1. The only synchronization necessary is between
        // calls to `try_close`: we have to ensure that all threads have
        // dropped their refs to the span before the span is closed.
        let refs = span.ref_count.fetch_add(1, Ordering::Relaxed);
        assert_ne!(
            refs, 0,
            "tried to clone a span ({:?}) that already closed",
            id
        );
        id.clone()
    }

    fn try_close(&self, id: Id) -> bool {
        let Some(span) = self.get(&id) else {
            panic!("tried to drop a ref to {:?}, but no such span exists!", id);
        };

        let refs = span.ref_count.fetch_sub(1, Ordering::Release);
        assert!(refs < usize::MAX, "reference count overflow!");
        if refs > 1 {
            return false;
        }

        // Synchronize if we are actually removing the span (stolen
        // from std::Arc); this ensures that all other `try_close` calls on
        // other threads happen-before we actually remove the span.
        atomic::fence(Ordering::Acquire);
        true
    }

    fn current_span(&self) -> Current {
        self.current_spans
            .get()
            .and_then(|spans| {
                let spans = spans.borrow();
                let id = spans.current()?;
                let span = self.get(id)?;
                Some(Current::new(id.clone(), span.metadata))
            })
            .unwrap_or_else(Current::none)
    }
}

// === impl DataInner ===

impl Default for DataInner {
    fn default() -> Self {
        // Since `DataInner` owns a `&'static Callsite` pointer, we need
        // something to use as the initial default value for that callsite.
        // Since we can't access a `DataInner` until it has had actual span data
        // inserted into it, the null metadata will never actually be accessed.
        struct NullCallsite;
        impl tracing_core::callsite::Callsite for NullCallsite {
            fn set_interest(&self, _: Interest) {
                unreachable!(
                    "/!\\ Tried to register the null callsite /!\\\n \
                    This should never have happened and is definitely a bug. \
                    A `tracing` bug report would be appreciated."
                )
            }

            fn metadata(&self) -> &Metadata<'_> {
                unreachable!(
                    "/!\\ Tried to access the null callsite's metadata /!\\\n \
                    This should never have happened and is definitely a bug. \
                    A `tracing` bug report would be appreciated."
                )
            }
        }

        static NULL_CALLSITE: NullCallsite = NullCallsite;
        static NULL_METADATA: Metadata<'static> = tracing_core::metadata! {
            name: "",
            target: "",
            level: tracing_core::Level::TRACE,
            fields: &[],
            callsite: &NULL_CALLSITE,
            kind: tracing_core::metadata::Kind::SPAN,
        };

        Self {
            metadata: &NULL_METADATA,
            parent: None,
            ref_count: AtomicUsize::new(0),
        }
    }
}

impl ksharded_slab::Clear for DataInner {
    /// Clears the span's data in place, dropping the parent's reference count.
    fn clear(&mut self) {
        // A span is not considered closed until all of its children have closed.
        // Therefore, each span's `DataInner` holds a "reference" to the parent
        // span, keeping the parent span open until all its children have closed.
        // When we close a span, we must then decrement the parent's ref count
        // (potentially, allowing it to close, if this child is the last reference
        // to that span).
        // We have to actually unpack the option inside the `get_default`
        // closure, since it is a `FnMut`, but testing that there _is_ a value
        // here lets us avoid the thread-local access if we don't need the
        // dispatcher at all.
        if self.parent.is_some() {
            // Note that --- because `Layered::try_close` works by calling
            // `try_close` on the inner subscriber and using the return value to
            // determine whether to call the subscriber's `on_close` callback ---
            // we must call `try_close` on the entire subscriber stack, rather
            // than just on the registry. If the registry called `try_close` on
            // itself directly, the subscribers wouldn't see the close notification.
            let subscriber = dispatch::get_default(Dispatch::clone);
            if let Some(parent) = self.parent.take() {
                let _ = subscriber.try_close(parent);
            }
        }
    }
}

impl Data<'_> {
    fn id(&self) -> Id {
        idx_to_id(self.inner.key())
    }

    fn metadata(&self) -> &'static Metadata<'static> {
        self.inner.metadata
    }

    fn parent(&self) -> Option<&Id> {
        self.inner.parent.as_ref()
    }
}

impl<'a> SpanRef<'a> {
    /// Returns this span's ID.
    pub fn id(&self) -> Id {
        self.data.id()
    }

    /// Returns a static reference to the span's metadata.
    pub fn metadata(&self) -> &'static Metadata<'static> {
        self.data.metadata()
    }

    /// Returns the span's name,
    pub fn name(&self) -> &'static str {
        self.data.metadata().name()
    }

    /// Returns a list of [fields] defined by the span.
    ///
    /// [fields]: tracing_core::field
    pub fn fields(&self) -> &FieldSet {
        self.data.metadata().fields()
    }

    /// Returns the ID of this span's parent, or `None` if this span is the root
    /// of its trace tree.
    pub fn parent_id(&self) -> Option<&Id> {
        self.data.parent()
    }

    /// Returns a `SpanRef` describing this span's parent, or `None` if this
    /// span is the root of its trace tree.
    pub fn parent(&self) -> Option<Self> {
        let id = self.data.parent()?;
        let data = self.registry.span_data(id)?;

        Some(Self {
            registry: self.registry,
            data,
        })
    }

    /// Returns an iterator over all parents of this span, starting with this span,
    /// ordered from leaf to root.
    ///
    /// The iterator will first return the span, then the span's immediate parent,
    /// followed by that span's parent, and so on, until it reaches a root span.
    pub fn scope(&self) -> Scope<'a> {
        Scope {
            registry: self.registry,
            next: Some(self.id()),
        }
    }
}

// === impl SpanStack ===

impl SpanStack {
    #[inline]
    pub(crate) fn push(&mut self, id: Id) -> bool {
        let duplicate = self.stack.iter().any(|i| i.id == id);
        self.stack.push(ContextId { id, duplicate });
        !duplicate
    }

    #[inline]
    pub(crate) fn pop(&mut self, expected_id: &Id) -> bool {
        if let Some((idx, _)) = self
            .stack
            .iter()
            .enumerate()
            .rev()
            .find(|(_, ctx_id)| ctx_id.id == *expected_id)
        {
            let ContextId { id: _, duplicate } = self.stack.remove(idx);
            return !duplicate;
        }
        false
    }

    #[inline]
    pub(crate) fn iter(&self) -> impl Iterator<Item = &Id> {
        self.stack
            .iter()
            .rev()
            .filter_map(|ContextId { id, duplicate }| if !*duplicate { Some(id) } else { None })
    }

    #[inline]
    pub(crate) fn current(&self) -> Option<&Id> {
        self.iter().next()
    }
}

#[inline]
fn idx_to_id(idx: usize) -> Id {
    Id::from_u64(idx as u64 + 1)
}

#[inline]
fn id_to_idx(id: &Id) -> usize {
    usize::try_from(id.into_u64()).unwrap() - 1
}