import os

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult


class CountingExporter(SpanExporter):
    def __init__(self):
        self.names = []

    def export(self, spans):
        self.names.extend(span.name for span in spans)
        return SpanExportResult.SUCCESS


exporter = CountingExporter()
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
trace.set_tracer_provider(provider)

tracer = trace.get_tracer("sampling-check")
for index in range(40):
    with tracer.start_as_current_span(f"checkout-{index:02d}"):
        with tracer.start_as_current_span(f"checkout-db-{index:02d}"):
            pass

provider.force_flush()

root_spans = [name for name in exporter.names if not name.startswith("checkout-db-")]
child_spans = [name for name in exporter.names if name.startswith("checkout-db-")]

print(f"sampler={os.getenv('OTEL_TRACES_SAMPLER', 'unset')}")
print(f"sampler_arg={os.getenv('OTEL_TRACES_SAMPLER_ARG', 'unset')}")
print("root_spans_started=40")
print(f"exported_root_spans={len(root_spans)}")
print(f"exported_child_spans={len(child_spans)}")
print("first_exported_roots=" + ",".join(root_spans[:8]))