# kafka_test.py

from kafka import KafkaProducer
from kafka.admin import KafkaAdminClient, NewTopic
from opentelemetry import trace
from opentelemetry.instrumentation.kafka import KafkaInstrumentor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, ConsoleSpanExporter


provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
trace.set_tracer_provider(provider)

TOPIC = "test-topic"
NUM_PARTITIONS = 12
admin = KafkaAdminClient(bootstrap_servers="localhost:9092")
try:
    admin.create_topics([NewTopic(TOPIC, num_partitions=NUM_PARTITIONS, replication_factor=1)])
except Exception:
    pass

def produce_hook(span, args, kwargs):
    if span and span.is_recording():
        value = kwargs.get("value") if "value" in kwargs else (args[1] if len(args) > 1 else None)
        if value is not None:
            span.set_attribute(
                "messaging.kafka.message.value",
                value.decode() if isinstance(value, (bytes, bytearray)) else str(value),
            )
        key = kwargs.get("key") if "key" in kwargs else (args[2] if len(args) > 2 else None)
        if key is not None:
            span.set_attribute(
                "messaging.kafka.message.key",
                key.decode() if isinstance(key, (bytes, bytearray)) else str(key),
            )


KafkaInstrumentor().instrument(produce_hook=produce_hook)
producer = KafkaProducer(bootstrap_servers="localhost:9092")

N = 10

for i in range(N):
    fut = producer.send(TOPIC, value=f"v{i}".encode())
    actual = fut.get(timeout=10).partition