package org.mockserver.async.publish;

import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.header.Header;
import org.junit.Before;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;

import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.Map;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.*;
import static org.mockito.Mockito.verify;
import static org.mockito.MockitoAnnotations.openMocks;

/**
 * Tests for the enhanced {@link KafkaMessagePublisher} with key and header support.
 */
public class KafkaMessagePublisherKeyHeadersTest {

    @Mock
    private KafkaProducer<String, String> mockProducer;

    private KafkaMessagePublisher publisher;

    @Before
    public void setUp() {
        openMocks(this);
        publisher = new KafkaMessagePublisher(mockProducer);
    }

    @Test
    public void shouldSendRecordWithKey() {
        publisher.publish("my-topic", "my-key", "{\"data\":1}", null);

        @SuppressWarnings("unchecked")
        ArgumentCaptor<ProducerRecord<String, String>> captor =
            ArgumentCaptor.forClass(ProducerRecord.class);
        verify(mockProducer).send(captor.capture());

        ProducerRecord<String, String> record = captor.getValue();
        assertThat(record.topic(), is("my-topic"));
        assertThat(record.key(), is("my-key"));
        assertThat(record.value(), is("{\"data\":1}"));
    }

    @Test
    public void shouldSendRecordWithNullKey() {
        publisher.publish("my-topic", null, "{\"data\":1}", null);

        @SuppressWarnings("unchecked")
        ArgumentCaptor<ProducerRecord<String, String>> captor =
            ArgumentCaptor.forClass(ProducerRecord.class);
        verify(mockProducer).send(captor.capture());

        ProducerRecord<String, String> record = captor.getValue();
        assertThat(record.key(), is(nullValue()));
    }

    @Test
    public void shouldSendRecordWithHeaders() {
        Map<String, String> headers = new LinkedHashMap<>();
        headers.put("trace-id", "abc123");
        headers.put("content-type", "application/json");

        publisher.publish("my-topic", "key1", "{\"data\":1}", headers);

        @SuppressWarnings("unchecked")
        ArgumentCaptor<ProducerRecord<String, String>> captor =
            ArgumentCaptor.forClass(ProducerRecord.class);
        verify(mockProducer).send(captor.capture());

        ProducerRecord<String, String> record = captor.getValue();
        Header traceHeader = record.headers().lastHeader("trace-id");
        assertThat(traceHeader, is(notNullValue()));
        assertThat(new String(traceHeader.value(), StandardCharsets.UTF_8), is("abc123"));

        Header contentTypeHeader = record.headers().lastHeader("content-type");
        assertThat(contentTypeHeader, is(notNullValue()));
        assertThat(new String(contentTypeHeader.value(), StandardCharsets.UTF_8), is("application/json"));
    }

    @Test
    public void shouldSendWithoutHeadersWhenNull() {
        publisher.publish("my-topic", "key1", "{\"data\":1}", null);

        @SuppressWarnings("unchecked")
        ArgumentCaptor<ProducerRecord<String, String>> captor =
            ArgumentCaptor.forClass(ProducerRecord.class);
        verify(mockProducer).send(captor.capture());

        ProducerRecord<String, String> record = captor.getValue();
        assertThat(record.headers().toArray().length, is(0));
    }

    @Test
    public void shouldSendViaDefaultInterfaceMethod() {
        // The simple publish(channel, payload) should still work
        publisher.publish("my-topic", "{\"simple\":true}");

        @SuppressWarnings("unchecked")
        ArgumentCaptor<ProducerRecord<String, String>> captor =
            ArgumentCaptor.forClass(ProducerRecord.class);
        verify(mockProducer).send(captor.capture());

        ProducerRecord<String, String> record = captor.getValue();
        assertThat(record.topic(), is("my-topic"));
        assertThat(record.key(), is(nullValue()));
        assertThat(record.value(), is("{\"simple\":true}"));
    }
}
