package org.zalando.nakadi.repository.kafka;

import org.apache.curator.framework.CuratorFramework;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.json.JSONException;
import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Scheduled;
import org.zalando.nakadi.repository.zookeeper.ZooKeeperHolder;

import java.io.UnsupportedEncodingException;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

public class KafkaLocationManager {

    private static final Logger LOG = LoggerFactory.getLogger(KafkaLocationManager.class);
    private static final String BROKERS_IDS_PATH = "/brokers/ids";

    private final ZooKeeperHolder zkFactory;
    private final Properties kafkaProperties;
    private final KafkaSettings kafkaSettings;

    public KafkaLocationManager(final ZooKeeperHolder zkFactory, final KafkaSettings kafkaSettings) {
        this.zkFactory = zkFactory;
        this.kafkaProperties = buildKafkaProperties(fetchBrokers());
        this.kafkaSettings = kafkaSettings;
    }

    static class Broker {
        final String host;
        final Integer port;

        private Broker(final String host, final Integer port) {
            this.host = host;
            this.port = port;
        }

        static Broker fromByteJson(final byte[] data) throws JSONException, UnsupportedEncodingException {
            final JSONObject json = new JSONObject(new String(data, "UTF-8"));
            final String host = json.getString("host");
            final Integer port = json.getInt("port");
            return new Broker(host, port);
        }

        public String toString() {
            return this.host + ":" + this.port;
        }
    }

    private List<Broker> fetchBrokers() {
        final List<Broker> brokers = new ArrayList<>();
        try {
            final CuratorFramework curator = zkFactory.get();
            for (final String brokerId : curator.getChildren().forPath(BROKERS_IDS_PATH)) {
                try {
                    final byte[] brokerData = curator.getData().forPath(BROKERS_IDS_PATH + "/" + brokerId);
                    brokers.add(Broker.fromByteJson(brokerData));
                } catch (final Exception e) {
                    LOG.info(String.format("Failed to fetch connection string for broker %s", brokerId), e);
                }
            }
        } catch (final Exception e) {
            LOG.error("Failed to fetch list of brokers from ZooKeeper", e);
        }

        return brokers;
    }

    private static String buildBootstrapServers(final List<Broker> brokers) {
        final StringBuilder builder = new StringBuilder();
        brokers.stream().forEach(broker -> builder.append(broker).append(","));
        return builder.deleteCharAt(builder.length() - 1).toString();
    }

    private Properties buildKafkaProperties(final List<Broker> brokers) {
        final Properties props = new Properties();
        props.put("bootstrap.servers", buildBootstrapServers(brokers));
        props.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        props.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
        return props;
    }

    @Scheduled(fixedDelay = 30000)
    private void updateBrokers() {
        if (kafkaProperties != null) {
            final List<Broker> brokers = fetchBrokers();
            if (!brokers.isEmpty()) {
                kafkaProperties.setProperty("bootstrap.servers", buildBootstrapServers(brokers));
            }
        }
    }

    public Properties getKafkaConsumerProperties() {
        final Properties properties = (Properties) kafkaProperties.clone();
        properties.put("enable.auto.commit", kafkaSettings.getEnableAutoCommit());
        properties.put("key.deserializer", "org.apache.kafka.common.serialization.ByteArrayDeserializer");
        properties.put("value.deserializer", "org.apache.kafka.common.serialization.ByteArrayDeserializer");
        return properties;
    }

    public Properties getKafkaProducerProperties() {
        final Properties producerProps = getKafkaConsumerProperties();
        producerProps.put("value.serializer", "org.apache.kafka.common.serialization.StringSerializer");
        producerProps.put("key.serializer", "org.apache.kafka.common.serialization.StringSerializer");
        producerProps.put("acks", "all");
        producerProps.put("request.timeout.ms", kafkaSettings.getRequestTimeoutMs());
        producerProps.put("batch.size", kafkaSettings.getBatchSize());
        producerProps.put("linger.ms", kafkaSettings.getLingerMs());
        producerProps.put(ProducerConfig.COMPRESSION_TYPE_CONFIG, "lz4");
        return producerProps;
    }
}

@@ -93,16 +93,126 @@
 mework;%0A
+import org.apache.kafka.clients.CommonClientConfigs;%0Aimport org.apache.kafka.clients.consumer.ConsumerConfig;%0A
 import o
@@ -3201,37 +3201,77 @@
 ops.put(
-%22value.deserializer%22,
+ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,%0A               
  %22org.ap
@@ -3294,38 +3294,41 @@
 n.serialization.
-String
+ByteArray
 Deserializer%22);%0A
@@ -3345,35 +3345,79 @@
 ops.put(
-%22key.deserializer%22,
+ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,%0A               
  %22org.ap
@@ -3448,22 +3448,25 @@
 ization.
-String
+ByteArray
 Deserial
@@ -3750,36 +3750,85 @@
 roperty(
-%22bootstrap.servers%22,
+CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,%0A                       
  buildBo
@@ -4040,28 +4040,48 @@
 put(
-%22enable.auto.commit%22
+ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG
 , ka
@@ -4143,27 +4143,69 @@
 put(
-%22key.deserializer%22,
+ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,%0A               
  %22or
@@ -4292,29 +4292,71 @@
 put(
-%22value.deserializer%22,
+ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,%0A               
  %22or
@@ -4604,27 +4604,69 @@
 put(
-%22value.serializer%22,
+ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,%0A               
  %22or
@@ -4751,25 +4751,67 @@
 put(
-%22key.serializer%22,
+ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,%0A               
  %22or
@@ -4896,14 +4896,34 @@
 put(
-%22acks%22
+ProducerConfig.ACKS_CONFIG
 , %22a
@@ -4958,28 +4958,48 @@
 put(
-%22request.timeout.ms%22
+ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG
 , ka
@@ -5064,20 +5064,40 @@
 put(
-%22batch.size%22
+ProducerConfig.BATCH_SIZE_CONFIG
 , ka
@@ -5155,19 +5155,39 @@
 put(
-%22linger.ms%22
+ProducerConfig.LINGER_MS_CONFIG
 , ka

replace str with constant from kafka client