// Copyright (c) 2026 Yingjie Shang
// agent-telemetry is licensed under Mulan PSL v2.

///| Types

///|
/// Configuration for initializing OpenTelemetry in an agent application.
pub struct TelemetryConfig {
  otel_stdout : Bool
  otel_endpoint : String
  service_name : String
} derive(Debug)

///|
/// Create a new telemetry configuration.
pub fn TelemetryConfig::new(
  otel_stdout? : Bool = false,
  otel_endpoint? : String = "http://localhost:4318",
  service_name? : String = "agent-telemetry",
) -> TelemetryConfig {
  { otel_stdout, otel_endpoint, service_name }
}

///| OpenTelemetry provider initialization

///|
/// Initialize the global OpenTelemetry tracer provider.
///
/// If `exporter` is provided, attaches it to a batch span processor.
/// Otherwise builds a no-op provider that discards all spans.
pub fn init_telemetry(
  config : TelemetryConfig,
  exporter? : @sdktrace.SpanExporter,
  id_generator : @sdktrace.IdGenerator,
) -> @sdktrace.SdkTracerProvider {
  let resource = build_resource(config.service_name)
  let batch_config = @sdktrace.BatchConfig::new(
    max_queue_size=64,
    max_export_batch_size=16,
    scheduled_delay_millis=1000,
    export_timeout_millis=5000,
  )
  let builder = @sdk.tracer_provider_builder()
    .with_resource(resource)
    .with_id_generator(id_generator)
  let provider = match exporter {
    Some(exporter) =>
      builder.with_batch_exporter(exporter, config=batch_config).build()
    None => builder.build()
  }
  @sdk.set_tracer_provider(provider)
  provider
}

///| Tracer cache

///|
/// Cache of tracers keyed by scope name.
let tracer_cache : Ref[Map[String, @trace.Tracer]] = Ref({})

///|
/// Get or create a tracer for the given scope name.
///
/// Tracers are cached so repeated calls return the same instance.
pub fn tracer(
  scope_name : String,
  version? : String = "0.1.0",
) -> @trace.Tracer {
  match tracer_cache.val.get(scope_name) {
    Some(t) => t
    None => {
      let t = @otel.tracer(scope_name, version=Some(version))
      tracer_cache.val.set(scope_name, t)
      t
    }
  }
}

///| Span lifecycle helpers

///|
/// Start a new span with the given tracer.
pub fn start_span(
  tracer : @trace.Tracer,
  name : String,
  kind? : @trace.SpanKind = @trace.Internal,
  attributes? : Array[@otel.KeyValue] = [],
  parent_context? : @context.Context = @context.Context::empty(),
) -> @trace.Span {
  let builder = tracer
    .span_builder(name)
    .with_kind(kind)
    .with_attributes(attributes)
  tracer.build_with_context(builder, parent_context)
}

///|
/// End a span, optionally setting a final status.
pub async fn end_span(span : @trace.Span, status? : @trace.Status) -> Unit {
  match status {
    Some(s) => span.set_status(s)
    None => ()
  }
  span.end()
}

///|
/// End a span marking it as successful.
pub async fn end_span_ok(span : @trace.Span) -> Unit {
  end_span(span, status=@trace.Status::ok())
}

///|
/// End a span marking it as failed.
pub async fn end_span_error(span : @trace.Span, description~ : String) -> Unit {
  end_span(span, status=@trace.Status::error(description=Some(description)))
}

///|
/// Set multiple attributes on a span.
pub fn set_attributes(
  span : @trace.Span,
  attributes : Array[@otel.KeyValue],
) -> Unit {
  for attr in attributes {
    span.set_attribute(attr)
  }
}

///| Internal helpers

///|
/// Build a default resource, ensuring service.name is set.
fn build_resource(service_name : String) -> @resource.Resource {
  let base = @resource.Resource::builder().build()
  match @env.get_env_var("OTEL_SERVICE_NAME") {
    Some(_) => base
    None =>
      base.merge(
        @resource.Resource::new([
          @common.KeyValue::new("service.name", String(service_name)),
        ]),
      )
  }
}