from kafka import KafkaConsumer
from kafka import TopicPartition
import boto3
import json
from datetime import datetime

# Configuratie
KAFKA_BROKER = "10.0.21.36:9092"
S3_ENDPOINT = "http://10.0.20.111:9020"
S3_BUCKET = "data"
S3_ACCESS_KEY = "REDACTED"
S3_SECRET_KEY = "REDACTED"

print("Starting S3 consumer for ALL topics...")
print(f"Kafka: {KAFKA_BROKER}")
print(f"S3: {S3_ENDPOINT}, Bucket: {S3_BUCKET}")

# S3 client
s3 = boto3.client(
    's3',
    endpoint_url=S3_ENDPOINT,
    aws_access_key_id=S3_ACCESS_KEY,
    aws_secret_access_key=S3_SECRET_KEY,
    use_ssl=False,
    verify=False
)

# Consumer zonder specifiek topic
consumer = KafkaConsumer(bootstrap_servers=[KAFKA_BROKER])

# Haal alle topics op
all_topics = consumer.topics()
print(f"Found topics: {list(all_topics)}")

# Wijs alle partitions van alle topics toe
for topic in all_topics:
    partitions = consumer.partitions_for_topic(topic)
    for partition in partitions:
        tp = TopicPartition(topic, partition)
        consumer.assign([tp])
        consumer.seek_to_beginning()
        print(f"Assigned: {topic} - partition {partition}")

print("\nReading all messages from all topics...")
print("-" * 60)

msg_count = 0
topic_count = {}

for msg in consumer:
    msg_count += 1
    topic = msg.topic
    topic_count[topic] = topic_count.get(topic, 0) + 1
    
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
    key = f"kafka/{topic}/message_{timestamp}.json"
    
    try:
        data = json.loads(msg.value.decode('utf-8'))
    except:
        data = {"raw": msg.value.decode('utf-8')}
    
    s3.put_object(
        Bucket=S3_BUCKET,
        Key=key,
        Body=json.dumps(data, indent=2)
    )
    
    print(f"[{msg_count}] {topic} -> s3://{S3_BUCKET}/{key}")