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

use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome};
use tracing::{Id, Subscriber, field::Visit};
use tracing_subscriber::{Layer, layer::SubscriberExt, registry::LookupSpan};

#[test]
fn stage_names_and_attributes_are_closed() {
    let captured = Arc::new(Mutex::new(Vec::new()));
    let subscriber = tracing_subscriber::registry().with(CaptureLayer(Arc::clone(&captured)));

    tracing::subscriber::with_default(subscriber, || {
        let span = Stage::RuntimeExecute.span();
        StageOutcome::Success.record(&span);
        ErrorCategory::Mapping.record(&span);
        drop(span);

        let db_span = DbOperation::InvocationHistoryWrite.span();
        StageOutcome::Error.record(&db_span);
        drop(db_span);
    });

    let spans = captured.lock().expect("captured spans");
    assert_eq!(spans[0].name, "runtime.execute");
    assert_eq!(spans[0].fields["outcome"], "success");
    assert_eq!(spans[0].fields["error.category"], "mapping");
    assert_eq!(spans[1].name, "db.query");
    assert_eq!(spans[1].fields["db.system"], "postgresql");
    assert_eq!(spans[1].fields["db.operation"], "invocation_history.write");
}

#[test]
fn database_operation_names_are_closed() {
    assert_eq!(DbOperation::CatalogLoad.as_str(), "catalog.load");
}

#[derive(Clone)]
struct CaptureLayer(Arc<Mutex<Vec<CapturedSpan>>>);

#[derive(Debug)]
struct CapturedSpan {
    name: &'static str,
    fields: std::collections::BTreeMap<String, String>,
}

impl<S> Layer<S> for CaptureLayer
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    fn on_new_span(
        &self,
        attributes: &tracing::span::Attributes<'_>,
        id: &Id,
        context: tracing_subscriber::layer::Context<'_, S>,
    ) {
        let mut visitor = FieldVisitor::default();
        attributes.record(&mut visitor);
        context
            .span(id)
            .expect("span exists")
            .extensions_mut()
            .insert(self.0.lock().expect("capture lock").len());
        self.0.lock().expect("capture lock").push(CapturedSpan {
            name: attributes.metadata().name(),
            fields: visitor.fields,
        });
    }

    fn on_record(
        &self,
        id: &Id,
        values: &tracing::span::Record<'_>,
        context: tracing_subscriber::layer::Context<'_, S>,
    ) {
        let span = context.span(id).expect("span exists");
        let index = *span.extensions().get::<usize>().expect("capture index");
        let mut visitor = FieldVisitor::default();
        values.record(&mut visitor);
        self.0.lock().expect("capture lock")[index]
            .fields
            .extend(visitor.fields);
    }
}

#[derive(Default)]
struct FieldVisitor {
    fields: std::collections::BTreeMap<String, String>,
}

impl Visit for FieldVisitor {
    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
        self.fields
            .insert(field.name().to_owned(), value.to_owned());
    }

    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        self.fields
            .insert(field.name().to_owned(), format!("{value:?}"));
    }
}