package com.lingniu.ingest.sink.kafka;

import com.lingniu.ingest.api.consumer.EnvelopeConsumerProcessor;
import com.lingniu.ingest.api.consumer.EnvelopeConsumerRecord;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;

import java.time.Duration;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

public final class KafkaEnvelopeConsumerWorker implements AutoCloseable {

    private final Consumer<String, byte[]> consumer;
    private final Map<String, EnvelopeConsumerProcessor> processorsByTopic;

    public KafkaEnvelopeConsumerWorker(Consumer<String, byte[]> consumer,
                                       Map<String, EnvelopeConsumerProcessor> processorsByTopic) {
        if (consumer == null) {
            throw new IllegalArgumentException("consumer must not be null");
        }
        if (processorsByTopic == null || processorsByTopic.isEmpty()) {
            throw new IllegalArgumentException("processorsByTopic must not be empty");
        }
        this.consumer = consumer;
        this.processorsByTopic = Map.copyOf(processorsByTopic);
    }

    public void subscribe(Collection<String> topics) {
        consumer.subscribe(topics);
    }

    public int pollOnce(Duration timeout) {
        ConsumerRecords<String, byte[]> records = consumer.poll(timeout == null ? Duration.ZERO : timeout);
        Map<EnvelopeConsumerProcessor, List<EnvelopeConsumerRecord>> byProcessor = new LinkedHashMap<>();
        int processed = 0;
        for (ConsumerRecord<String, byte[]> record : records) {
            EnvelopeConsumerProcessor processor = processorsByTopic.get(record.topic());
            if (processor == null) {
                // worker 可能订阅多个 topic；没有显式绑定处理器的 topic 不参与提交语义。
                continue;
            }
            byProcessor.computeIfAbsent(processor, ignored -> new ArrayList<>()).add(new EnvelopeConsumerRecord(
                    record.topic(),
                    record.partition(),
                    record.offset(),
                    record.key(),
                    record.value()));
            processed++;
        }
        for (Map.Entry<EnvelopeConsumerProcessor, List<EnvelopeConsumerRecord>> entry : byProcessor.entrySet()) {
            // EnvelopeConsumerProcessor 内部会把解析或业务错误转成 DLQ 记录，
            // 这里保持 Kafka worker 的职责单一：轮询、分发、成功后提交 offset。
            entry.getKey().processBatch(entry.getValue());
        }
        if (processed > 0) {
            // commitSync 放在批次末尾，保证同一个 poll 批次内的消息按 Kafka offset 一起确认。
            consumer.commitSync();
        }
        return processed;
    }

    public void wakeup() {
        consumer.wakeup();
    }

    @Override
    public void close() {
        consumer.close();
    }
}
