import type { Context } from '@opentelemetry/api'
import { trace } from '@opentelemetry/api'
import type { BufferConfig, ReadableSpan, Span, SpanExporter } from '@opentelemetry/sdk-trace-base'
import { BatchSpanProcessor } from '@opentelemetry/sdk-trace-base'

import type { TraceStore } from '../core/traceStore'

export class CacheBatchSpanProcessor extends BatchSpanProcessor {
  private cache: TraceStore

  constructor(_exporter: SpanExporter, cache: TraceStore, config?: BufferConfig) {
    super(_exporter, config)
    this.cache = cache
  }

  override onEnd(span: ReadableSpan): void {
    super.onEnd(span)
    this.cache.endSpan(span)
  }

  override onStart(span: Span, parentContext: Context): void {
    super.onStart(span, parentContext)
    this.cache.createSpan({
      name: span.name,
      kind: span.kind,
      spanContext: () => span.spanContext(),
      parentSpanContext: trace.getSpanContext(parentContext),
      startTime: span.startTime,
      status: span.status,
      attributes: span.attributes,
      links: span.links,
      events: span.events,
      duration: span.duration,
      ended: span.ended,
      resource: span.resource,
      instrumentationScope: span.instrumentationScope,
      droppedAttributesCount: span.droppedAttributesCount,
      droppedEventsCount: span.droppedEventsCount,
      droppedLinksCount: span.droppedLinksCount
    } as ReadableSpan)
  }
}
