//! Telemetry module: OTel init, context propagation, span helpers.
pub mod baggage_span_processor;
pub mod connection;
pub mod context;
pub mod http_instrumentation;
pub mod json_serializer;
pub mod log_exporter;
pub mod metrics_exporter;
pub mod otel_worker_gauges;
pub mod payload;
pub mod span_exporter;
pub mod span_ops;
pub mod types;
pub mod worker_metrics;

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};

use opentelemetry::propagation::TextMapCompositePropagator;
use opentelemetry::trace::{SpanKind, Status, TraceContextExt, Tracer};
use opentelemetry::{Context as OtelContext, KeyValue};
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::logs::{BatchConfigBuilder, BatchLogProcessor, SdkLoggerProvider};
use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
use opentelemetry_sdk::propagation::{BaggagePropagator, TraceContextPropagator};
use opentelemetry_sdk::trace::{
    BatchConfigBuilder as SpanBatchConfigBuilder, BatchSpanProcessor, SdkTracerProvider,
};
use tokio::sync::Mutex;

use self::connection::SharedEngineConnection;
use self::log_exporter::EngineLogExporter;
use self::metrics_exporter::EngineMetricsExporter;
use self::span_exporter::EngineSpanExporter;
use self::types::OtelConfig;

/// Global OTEL state, initialized once via `init_otel`
struct OtelState {
    tracer_provider: SdkTracerProvider,
    meter_provider: Option<SdkMeterProvider>,
    logger_provider: Option<SdkLoggerProvider>,
    connection: Arc<SharedEngineConnection>,
    shutdown_timeout: std::time::Duration,
}

static OTEL_STATE: OnceLock<Mutex<Option<OtelState>>> = OnceLock::new();
static OTEL_INITIALIZED: AtomicBool = AtomicBool::new(false);
/// Reference count of active workers using OpenTelemetry. `shutdown_otel` only
/// tears down the global state when this reaches zero.
static OTEL_REF_COUNT: AtomicUsize = AtomicUsize::new(0);

fn get_otel_lock() -> &'static Mutex<Option<OtelState>> {
    OTEL_STATE.get_or_init(|| Mutex::new(None))
}

/// Normalize an engine WebSocket URL into the dedicated OTEL endpoint.
///
/// The engine exposes `/otel` for telemetry-only WS connections. Appending
/// the path here means the SDK never shows up in `worker_registry` as a
/// ghost null-metadata worker. Handles trailing slashes, is idempotent
fn detect_service_name_from_argv0() -> Option<String> {
    let argv0 = std::env::args().next()?;
    if argv0.is_empty() {
        return None;
    }
    let path = std::path::Path::new(&argv0);
    path.file_name().and_then(|f| f.to_str()).map(String::from)
}

/// when the URL already ends in `/otel`, and preserves query strings and
/// fragments by inserting `/otel` into the path segment only.
fn append_otel_path(base: &str) -> String {
    let split_idx = base.find(['?', '#']).unwrap_or(base.len());
    let (prefix, suffix) = base.split_at(split_idx);
    let trimmed = prefix.trim_end_matches('/');
    if trimmed.ends_with("/otel") {
        format!("{trimmed}{suffix}")
    } else {
        format!("{trimmed}/otel{suffix}")
    }
}

#[cfg(test)]
mod otel_path_tests {
    use super::append_otel_path;

    #[test]
    fn appends_otel_path() {
        assert_eq!(
            append_otel_path("ws://localhost:49134"),
            "ws://localhost:49134/otel"
        );
    }

    #[test]
    fn strips_trailing_slash_before_appending() {
        assert_eq!(
            append_otel_path("ws://localhost:49134/"),
            "ws://localhost:49134/otel"
        );
    }

    #[test]
    fn is_idempotent() {
        assert_eq!(
            append_otel_path("ws://localhost:49134/otel"),
            "ws://localhost:49134/otel"
        );
    }

    #[test]
    fn preserves_query_string() {
        assert_eq!(
            append_otel_path("ws://localhost:49134?token=abc"),
            "ws://localhost:49134/otel?token=abc"
        );
    }

    #[test]
    fn preserves_query_string_when_already_otel() {
        assert_eq!(
            append_otel_path("ws://localhost:49134/otel?token=abc"),
            "ws://localhost:49134/otel?token=abc"
        );
    }

    #[test]
    fn preserves_fragment() {
        assert_eq!(
            append_otel_path("ws://localhost:49134#frag"),
            "ws://localhost:49134/otel#frag"
        );
    }

    #[test]
    fn strips_trailing_slash_before_query() {
        assert_eq!(
            append_otel_path("ws://localhost:49134/otel/?token=abc"),
            "ws://localhost:49134/otel?token=abc"
        );
    }
}

#[cfg(test)]
mod service_name_tests {
    //! `detect_service_name_from_argv0` reads `std::env::args()` directly,
    //! so these tests exercise the extraction logic via a local shim.

    use std::path::Path;

    fn name_from_path(argv0: &str) -> Option<String> {
        if argv0.is_empty() {
            return None;
        }
        Path::new(argv0)
            .file_name()
            .and_then(|f| f.to_str())
            .map(String::from)
    }

    #[test]
    fn extracts_binary_name_from_release_path() {
        assert_eq!(
            name_from_path("/path/to/workers/harness/target/release/harness").as_deref(),
            Some("harness")
        );
        assert_eq!(
            name_from_path("/path/to/workers/turn-orchestrator/target/release/turn-orchestrator")
                .as_deref(),
            Some("turn-orchestrator")
        );
    }

    #[test]
    fn extracts_bare_binary_name() {
        assert_eq!(name_from_path("iii").as_deref(), Some("iii"));
        assert_eq!(name_from_path("harness").as_deref(), Some("harness"));
    }

    #[test]
    fn returns_none_for_empty_argv0() {
        assert_eq!(name_from_path(""), None);
    }
}

/// Initialize OpenTelemetry with the given configuration.
///
/// Sets up distributed tracing, optional metrics, and optional log export
/// over a shared WebSocket connection to the III Engine.
///
/// This should be called once at startup. Subsequent calls are no-ops.
///
/// Returns `true` if this caller should call `shutdown_otel` when done
/// (i.e. otel was either freshly initialized or was already active).
/// Returns `false` when otel is disabled and no ref count was taken.
pub async fn init_otel(config: OtelConfig) -> bool {
    let lock = get_otel_lock();
    let mut state = lock.lock().await;

    if state.is_some() {
        OTEL_REF_COUNT.fetch_add(1, Ordering::SeqCst);
        tracing::debug!("OpenTelemetry already initialized, incrementing ref count");
        return true;
    }

    let enabled = config.enabled.unwrap_or_else(|| {
        std::env::var("OTEL_ENABLED")
            .map(|v| v == "true" || v == "1")
            .unwrap_or(true)
    });

    if !enabled {
        tracing::debug!("OpenTelemetry disabled, skipping initialization");
        return false;
    }

    // service.name resolution order (highest priority first):
    //   1. explicit `OtelConfig::service_name`
    //   2. `OTEL_SERVICE_NAME` env var
    //   3. argv[0] basename — the binary's own filename
    //      (`harness`, `turn-orchestrator`, `provider-anthropic`, ...).
    //      We deliberately DON'T fall back to `detect_project_name`
    //      (Cargo.toml-from-cwd) here because shared-cwd spawning
    //      patterns (the harness Makefile spawns every worker from
    //      `workers/harness/`) would make every worker report the
    //      same service name. argv[0] is unique per binary regardless
    //      of where it was launched from.
    //   4. fallback literal `"iii-rust-sdk"` for cases where argv[0]
    //      isn't a recognizable executable name (rare).
    let service_name = config
        .service_name
        .or_else(|| std::env::var("OTEL_SERVICE_NAME").ok())
        .or_else(detect_service_name_from_argv0)
        .unwrap_or_else(|| "iii-rust-sdk".to_string());

    let service_version = config
        .service_version
        .or_else(|| std::env::var("SERVICE_VERSION").ok())
        .unwrap_or_else(|| env!("CARGO_PKG_VERSION").to_string());

    let service_instance_id = config
        .service_instance_id
        .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());

    let ws_url = {
        let base = config
            .engine_ws_url
            .or_else(|| std::env::var("III_URL").ok())
            .unwrap_or_else(|| "ws://localhost:49134".to_string());
        // Route OTEL to the dedicated `/otel` endpoint so the engine
        // doesn't register this connection in `worker_registry`. The
        // telemetry socket would otherwise show up as a ghost worker
        // (null name/os/pid/runtime) alongside every real worker.
        append_otel_path(&base)
    };

    let reconnection_config = config.reconnection_config.unwrap_or_default();

    // Build resource
    let mut resource_attrs = vec![
        KeyValue::new("service.name", service_name.clone()),
        KeyValue::new("service.version", service_version),
        KeyValue::new("service.instance.id", service_instance_id),
        KeyValue::new("telemetry.sdk.name", "iii-rust-sdk"),
        KeyValue::new("telemetry.sdk.language", "rust"),
        KeyValue::new("telemetry.sdk.version", env!("CARGO_PKG_VERSION")),
    ];

    if let Some(ns) = config.service_namespace {
        resource_attrs.push(KeyValue::new("service.namespace", ns));
    }

    let resource = Resource::builder().with_attributes(resource_attrs).build();

    // Create shared WebSocket connection
    let channel_capacity = config.channel_capacity.unwrap_or(10_000);
    let connection = Arc::new(SharedEngineConnection::with_channel_capacity(
        ws_url,
        reconnection_config,
        channel_capacity,
    ));

    // Set up global propagator (W3C Trace Context + Baggage)
    let propagator = TextMapCompositePropagator::new(vec![
        Box::new(TraceContextPropagator::new()),
        Box::new(BaggagePropagator::new()),
    ]);
    opentelemetry::global::set_text_map_propagator(propagator);

    // BaggageSpanProcessor must register first: on_start fires in
    // registration order, so baggage entries are materialized as span
    // attributes before the batch exporter reads them.
    let span_exporter = EngineSpanExporter::new(connection.clone());

    // Without an explicit scheduled delay, the batch span processor inherits the
    // OpenTelemetry default of ~5000ms, so an ended span sits in the buffer up to
    // 5s before reaching the engine — the dominant reason traces appear seconds
    // late. Mirror the logs path: resolve config > custom env var > 100ms default.
    let spans_flush_ms = config
        .spans_flush_interval_ms
        .or_else(|| {
            std::env::var("OTEL_SPANS_FLUSH_INTERVAL_MS")
                .ok()
                .and_then(|v| v.parse::<u64>().ok())
        })
        .unwrap_or(100);

    let span_batch_config = SpanBatchConfigBuilder::default()
        .with_scheduled_delay(std::time::Duration::from_millis(spans_flush_ms))
        .build();
    let span_processor = BatchSpanProcessor::builder(span_exporter)
        .with_batch_config(span_batch_config)
        .build();

    let tracer_provider = SdkTracerProvider::builder()
        .with_resource(resource.clone())
        .with_span_processor(baggage_span_processor::BaggageSpanProcessor::default())
        .with_span_processor(span_processor)
        .build();

    tracing::debug!(
        flush_interval_ms = spans_flush_ms,
        "Span provider configured"
    );

    opentelemetry::global::set_tracer_provider(tracer_provider.clone());

    // Set up metrics if enabled
    let meter_provider = if config.metrics_enabled.unwrap_or(true) {
        let metrics_exporter = EngineMetricsExporter::new(connection.clone());
        let interval_ms = config.metrics_export_interval_ms.unwrap_or(60_000);

        let reader = PeriodicReader::builder(metrics_exporter)
            .with_interval(std::time::Duration::from_millis(interval_ms))
            .build();

        let provider = SdkMeterProvider::builder()
            .with_resource(resource.clone())
            .with_reader(reader)
            .build();

        opentelemetry::global::set_meter_provider(provider.clone());
        Some(provider)
    } else {
        None
    };

    // Set up logger provider if enabled
    let (logger_provider, resolved_flush_ms, resolved_batch_size) =
        if config.logs_enabled.unwrap_or(true) {
            let log_exporter = EngineLogExporter::new(connection.clone());

            // Resolve: config > custom env var > default (100ms flush, batch size 1).
            // Deliberately overrides standard OTEL_BLRP_* env vars for cross-SDK consistency.
            let flush_ms = config
                .logs_flush_interval_ms
                .or_else(|| {
                    std::env::var("OTEL_LOGS_FLUSH_INTERVAL_MS")
                        .ok()
                        .and_then(|v| v.parse::<u64>().ok())
                })
                .unwrap_or(100);

            let batch_size = config
                .logs_batch_size
                .or_else(|| {
                    std::env::var("OTEL_LOGS_BATCH_SIZE")
                        .ok()
                        .and_then(|v| v.parse::<usize>().ok())
                        .filter(|&v| v >= 1)
                })
                .unwrap_or(1);

            let batch_config = BatchConfigBuilder::default()
                .with_scheduled_delay(std::time::Duration::from_millis(flush_ms))
                .with_max_export_batch_size(batch_size)
                .build();

            let log_processor = BatchLogProcessor::builder(log_exporter)
                .with_batch_config(batch_config)
                .build();

            tracing::debug!(
                flush_interval_ms = flush_ms,
                batch_size = batch_size,
                "Log provider configured"
            );

            let provider = SdkLoggerProvider::builder()
                .with_resource(resource)
                .with_log_processor(log_processor)
                .build();

            (Some(provider), flush_ms, batch_size)
        } else {
            (None, 0, 0)
        };

    let shutdown_timeout =
        std::time::Duration::from_millis(config.shutdown_timeout_ms.unwrap_or(10_000));

    let otel_state = OtelState {
        tracer_provider,
        meter_provider,
        logger_provider,
        connection,
        shutdown_timeout,
    };

    *state = Some(otel_state);
    OTEL_INITIALIZED.store(true, Ordering::Release);
    OTEL_REF_COUNT.fetch_add(1, Ordering::SeqCst);

    tracing::info!(
        service = %service_name,
        logs_flush_ms = resolved_flush_ms,
        logs_batch = resolved_batch_size,
        "OpenTelemetry initialized"
    );

    true
}

/// Shutdown OpenTelemetry gracefully, flushing all pending data.
///
/// The shutdown sequence is bounded by the configured `shutdown_timeout_ms`
/// (default 10 seconds). If the timeout is exceeded, a warning is logged and
/// the function returns without waiting further.
pub async fn shutdown_otel() {
    let prev = OTEL_REF_COUNT.fetch_sub(1, Ordering::SeqCst);
    if prev > 1 {
        tracing::debug!(
            remaining = prev - 1,
            "OpenTelemetry still in use, skipping shutdown"
        );
        return;
    }

    let lock = get_otel_lock();
    let mut state = lock.lock().await;

    if let Some(otel) = state.take() {
        OTEL_INITIALIZED.store(false, Ordering::Release);

        let timeout_duration = otel.shutdown_timeout;

        match tokio::time::timeout(timeout_duration, async {
            // Shutdown tracer provider (flushes pending spans)
            if let Err(e) = otel.tracer_provider.shutdown() {
                tracing::warn!(error = %e, "Error shutting down tracer provider");
            }

            // Shutdown meter provider
            if let Some(meter) = otel.meter_provider {
                if let Err(e) = meter.shutdown() {
                    tracing::warn!(error = %e, "Error shutting down meter provider");
                }
            }

            // Shutdown logger provider
            if let Some(logger) = otel.logger_provider {
                if let Err(e) = logger.shutdown() {
                    tracing::warn!(error = %e, "Error shutting down logger provider");
                }
            }

            // Shutdown shared connection
            otel.connection.shutdown().await;
        })
        .await
        {
            Ok(()) => {
                tracing::info!("OpenTelemetry shut down");
            }
            Err(_) => {
                tracing::warn!(
                    timeout_ms = timeout_duration.as_millis() as u64,
                    "OpenTelemetry shutdown timed out"
                );
            }
        }
    }
}

/// Flush all pending telemetry data (spans, metrics, logs) through the connection.
///
/// This forces the OTEL providers to export any buffered data and then drains
/// the WebSocket connection's message queue. This is a no-op if OpenTelemetry
/// has not been initialized.
pub async fn flush_otel() {
    let lock = get_otel_lock();
    let state = lock.lock().await;

    if let Some(otel) = state.as_ref() {
        if let Err(e) = otel.tracer_provider.force_flush() {
            tracing::warn!(error = %e, "Error flushing tracer provider");
        }

        if let Some(meter) = &otel.meter_provider {
            if let Err(e) = meter.force_flush() {
                tracing::warn!(error = %e, "Error flushing meter provider");
            }
        }

        // Drain buffered messages to WebSocket
        otel.connection.flush().await;
    }
}

/// Execute a function within a new span, automatically handling errors and status.
///
/// If OpenTelemetry is not initialized, the function is executed without tracing.
///
/// # Arguments
/// * `name` - The span name
/// * `traceparent` - Optional W3C traceparent header to set parent context
/// * `kind` - Optional span kind (defaults to Internal)
/// * `f` - The async function to execute within the span
pub async fn with_span<F, Fut, T>(
    name: &str,
    traceparent: Option<&str>,
    kind: Option<SpanKind>,
    f: F,
) -> Result<T, Box<dyn std::error::Error + Send + Sync>>
where
    F: FnOnce() -> Fut,
    Fut: std::future::Future<Output = Result<T, Box<dyn std::error::Error + Send + Sync>>>,
{
    use opentelemetry::trace::FutureExt as OtelFutureExt;

    let tracer = opentelemetry::global::tracer("iii-rust-sdk");

    let parent_cx = match traceparent {
        Some(tp) => context::extract_traceparent(tp),
        None => OtelContext::current(),
    };

    let span_kind = kind.unwrap_or(SpanKind::Internal);

    let span = tracer
        .span_builder(name.to_string())
        .with_kind(span_kind)
        .start_with_context(&tracer, &parent_cx);

    let cx = parent_cx.with_span(span);

    match f().with_context(cx.clone()).await {
        Ok(result) => {
            cx.span().set_status(Status::Ok);
            Ok(result)
        }
        Err(err) => {
            let span = cx.span();
            span.set_status(Status::error(err.to_string()));

            let backtrace = std::backtrace::Backtrace::force_capture();
            span.add_event(
                "exception",
                vec![
                    KeyValue::new("exception.type", "Error"),
                    KeyValue::new("exception.message", err.to_string()),
                    KeyValue::new("exception.stacktrace", backtrace.to_string()),
                ],
            );

            Err(err)
        }
    }
}

/// Run a future inside a new span. Returns the future's output unchanged.
///
/// Does NOT set span status automatically — use [`with_span`] for
/// automatic OK/Error status on `Result` returns, or call
/// [`set_current_span_error`] from inside `f` to mark errors on typed
/// `thiserror` paths.
pub async fn run_in_span<F, Fut, T>(name: &str, kind: Option<SpanKind>, f: F) -> T
where
    F: FnOnce() -> Fut,
    Fut: std::future::Future<Output = T>,
{
    use opentelemetry::trace::FutureExt as OtelFutureExt;

    let tracer = opentelemetry::global::tracer("iii-rust-sdk");
    let parent_cx = OtelContext::current();
    let span = tracer
        .span_builder(name.to_string())
        .with_kind(kind.unwrap_or(SpanKind::Internal))
        .start_with_context(&tracer, &parent_cx);
    let cx = parent_cx.with_span(span);
    f().with_context(cx).await
}

/// Get a tracer instance for creating spans manually.
pub fn get_tracer() -> opentelemetry::global::BoxedTracer {
    opentelemetry::global::tracer("iii-rust-sdk")
}

/// Get a meter instance for creating metrics.
pub fn get_meter() -> opentelemetry::metrics::Meter {
    opentelemetry::global::meter("iii-rust-sdk")
}

/// Check whether OpenTelemetry has been initialized.
pub fn is_initialized() -> bool {
    OTEL_INITIALIZED.load(Ordering::Acquire)
}

/// Get a clone of the logger provider, if OTel has been initialized and logs are enabled.
///
/// Returns `None` if OTel is not initialized or logs are disabled.
/// This is used by the SDK logger to emit OTel LogRecords alongside the
/// engine invoker calls.
pub fn get_logger_provider() -> Option<SdkLoggerProvider> {
    if !is_initialized() {
        return None;
    }
    // Try to acquire the lock without blocking. If contended, skip emission.
    let lock = get_otel_lock();
    let state = lock.try_lock().ok()?;
    state.as_ref()?.logger_provider.clone()
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_logs_env_var_parsing_valid() {
        // Verify the parsing logic: valid u64 string parses correctly
        let val: Option<u64> = "500".parse::<u64>().ok();
        assert_eq!(val, Some(500));
    }

    #[test]
    fn test_logs_env_var_parsing_invalid() {
        // Verify the parsing logic: non-numeric string returns None
        let val: Option<u64> = "not-a-number".parse::<u64>().ok();
        assert!(val.is_none());
    }

    #[test]
    fn test_logs_batch_size_minimum_filter() {
        // Verify filter logic: batch size 0 is rejected
        let val: Option<usize> = "0".parse::<usize>().ok().filter(|&v| v >= 1);
        assert!(val.is_none());

        // Verify filter logic: batch size 1 is accepted
        let val: Option<usize> = "1".parse::<usize>().ok().filter(|&v| v >= 1);
        assert_eq!(val, Some(1));
    }

    #[tokio::test]
    #[serial_test::serial]
    async fn test_with_span_error_records_exception_event() {
        use opentelemetry_sdk::trace::{InMemorySpanExporter, SdkTracerProvider};

        let exporter = InMemorySpanExporter::default();
        let provider = SdkTracerProvider::builder()
            .with_simple_exporter(exporter.clone())
            .build();

        opentelemetry::global::set_tracer_provider(provider.clone());

        let result = super::with_span("test-error-span", None, None, || async {
            Err::<(), _>("test error".into())
        })
        .await;
        assert!(result.is_err());

        let _ = provider.force_flush();

        let spans = exporter.get_finished_spans().unwrap();
        assert!(!spans.is_empty(), "expected at least 1 span");

        let span = spans
            .iter()
            .find(|s| s.name == "test-error-span")
            .expect("should find test-error-span");
        let exc_event = span
            .events
            .iter()
            .find(|e| e.name == "exception")
            .expect("span should have an 'exception' event");

        let has_type = exc_event
            .attributes
            .iter()
            .any(|kv| kv.key.as_str() == "exception.type");
        let has_message = exc_event
            .attributes
            .iter()
            .any(|kv| kv.key.as_str() == "exception.message");
        let has_stacktrace = exc_event
            .attributes
            .iter()
            .any(|kv| kv.key.as_str() == "exception.stacktrace");

        assert!(has_type, "exception event should have exception.type");
        assert!(has_message, "exception event should have exception.message");
        assert!(
            has_stacktrace,
            "exception event should have exception.stacktrace"
        );

        let message_val = exc_event
            .attributes
            .iter()
            .find(|kv| kv.key.as_str() == "exception.message")
            .map(|kv| kv.value.to_string())
            .unwrap();
        assert!(
            message_val.contains("test error"),
            "exception.message should contain error text"
        );
    }

    #[tokio::test]
    #[serial_test::serial]
    async fn test_with_span_success_no_exception_event() {
        use opentelemetry_sdk::trace::{InMemorySpanExporter, SdkTracerProvider};

        let exporter = InMemorySpanExporter::default();
        let provider = SdkTracerProvider::builder()
            .with_simple_exporter(exporter.clone())
            .build();

        opentelemetry::global::set_tracer_provider(provider.clone());

        let result = super::with_span("test-ok-span", None, None, || async {
            Ok::<_, Box<dyn std::error::Error + Send + Sync>>(42)
        })
        .await;
        assert!(result.is_ok());

        let _ = provider.force_flush();

        let spans = exporter.get_finished_spans().unwrap();
        assert!(!spans.is_empty(), "expected at least 1 span");

        let span = spans
            .iter()
            .find(|s| s.name == "test-ok-span")
            .expect("should find test-ok-span");
        let exc_event = span.events.iter().find(|e| e.name == "exception");
        assert!(
            exc_event.is_none(),
            "successful span should not have an exception event"
        );
    }
}
