package cz.moneta.test.harness.connectors.messaging;

import cz.moneta.test.harness.connectors.messaging.kafkautils.CustomKafkaAvroDeserializer;
import cz.moneta.test.harness.connectors.messaging.kafkautils.JsonToAvroConverter;
import cz.moneta.test.harness.support.messaging.exception.MessagingConnectionException;
import cz.moneta.test.harness.support.messaging.exception.MessagingDestinationException;
import cz.moneta.test.harness.support.messaging.exception.MessagingSchemaException;
import cz.moneta.test.harness.support.messaging.exception.MessagingTimeoutException;
import cz.moneta.test.harness.support.messaging.kafka.MessageContentType;
import cz.moneta.test.harness.support.messaging.kafka.ReceivedMessage;
import io.confluent.kafka.schemaregistry.client.CachedSchemaRegistryClient;
import io.confluent.kafka.schemaregistry.client.SchemaRegistryClientConfig;
import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDeConfig;
import io.confluent.kafka.serializers.KafkaAvroSerializer;
import org.apache.avro.generic.GenericRecord;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.function.Predicate;

/**
 * Kafka connector for sending and receiving messages.
 * Supports Avro serialization with Confluent Schema Registry.
 * <p>
 * Uses manual partition assignment (no consumer group) for test isolation.
 * Each receive operation creates a new consumer to avoid offset sharing.
 */
public class KafkaConnector implements cz.moneta.test.harness.connectors.Connector {

    private static final Logger LOG = LogManager.getLogger(KafkaConnector.class);

    private final Properties producerConfig;
    private final Properties consumerConfig;
    private final String schemaRegistryUrl;
    private final CachedSchemaRegistryClient schemaRegistryClient;
    private KafkaProducer<String, GenericRecord> producer;

    /**
     * Creates a new KafkaConnector.
     *
     * @param bootstrapServers        Kafka bootstrap servers
     * @param apiKey                  Kafka API key
     * @param apiSecret               Kafka API secret
     * @param schemaRegistryUrl       Schema Registry URL
     * @param schemaRegistryApiKey    Schema Registry API key
     * @param schemaRegistryApiSecret Schema Registry API secret
     */
    public KafkaConnector(String bootstrapServers,
                          String apiKey,
                          String apiSecret,
                          String schemaRegistryUrl,
                          String schemaRegistryApiKey,
                          String schemaRegistryApiSecret) {
        this.schemaRegistryUrl = schemaRegistryUrl;

        HashMap<String, String> schemaRegistryProps = new HashMap<>();
        schemaRegistryProps.put(SchemaRegistryClientConfig.BASIC_AUTH_CREDENTIALS_SOURCE, "USER_INFO");
        schemaRegistryProps.put(SchemaRegistryClientConfig.USER_INFO_CONFIG, schemaRegistryApiKey + ":" + schemaRegistryApiSecret);
        this.schemaRegistryClient = new CachedSchemaRegistryClient(
                Collections.singletonList(schemaRegistryUrl), 100, schemaRegistryProps);

        this.producerConfig = createProducerConfig(bootstrapServers, apiKey, apiSecret, schemaRegistryApiKey, schemaRegistryApiSecret);
        this.consumerConfig = createConsumerConfig(bootstrapServers, apiKey, apiSecret, schemaRegistryApiKey, schemaRegistryApiSecret);
    }

    /**
     * Creates producer configuration.
     */
    private Properties createProducerConfig(String bootstrapServers, String apiKey, String apiSecret,
                                            String schemaRegistryApiKey, String schemaRegistryApiSecret) {
        Properties config = new Properties();
        config.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
        config.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
        config.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, KafkaAvroSerializer.class);
        config.put(ProducerConfig.ACKS_CONFIG, "all");
        config.put(ProducerConfig.LINGER_MS_CONFIG, 1);
        config.put(ProducerConfig.BATCH_SIZE_CONFIG, 16384);

        config.put(AbstractKafkaSchemaSerDeConfig.AUTO_REGISTER_SCHEMAS, false);
        config.put(AbstractKafkaSchemaSerDeConfig.USE_LATEST_VERSION, true);
        config.put(AbstractKafkaSchemaSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistryUrl);
        config.put(SchemaRegistryClientConfig.BASIC_AUTH_CREDENTIALS_SOURCE, "USER_INFO");
        config.put(SchemaRegistryClientConfig.USER_INFO_CONFIG, schemaRegistryApiKey + ":" + schemaRegistryApiSecret);

//         SASL/PLAIN authentication
        config.put("security.protocol", "SASL_SSL");
        config.put("sasl.mechanism", "PLAIN");
        config.put("sasl.jaas.config",
                "org.apache.kafka.common.security.plain.PlainLoginModule required " +
                        "username=\"" + apiKey + "\" password=\"" + apiSecret + "\";");

//         SSL configuration
        config.put("ssl.endpoint.identification.algorithm", "https");

        return config;
    }

    /**
     * Creates consumer configuration.
     */
    private Properties createConsumerConfig(String bootstrapServers, String apiKey, String apiSecret,
                                            String schemaRegistryApiKey, String schemaRegistryApiSecret) {
        Properties config = new Properties();
        config.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
        config.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        config.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, CustomKafkaAvroDeserializer.class);
        config.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none");
        config.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);


        // SASL/PLAIN authentication
        config.put("security.protocol", "SASL_SSL");
        config.put("sasl.mechanism", "PLAIN");
        config.put("sasl.jaas.config",
                "org.apache.kafka.common.security.plain.PlainLoginModule required " +
                        "username=\"" + apiKey + "\" password=\"" + apiSecret + "\";");

        // Schema Registry for deserialization

        config.put(AbstractKafkaSchemaSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistryUrl);
        config.put(AbstractKafkaSchemaSerDeConfig.AUTO_REGISTER_SCHEMAS, false);
        config.put(AbstractKafkaSchemaSerDeConfig.USE_LATEST_VERSION, true);
        config.put(SchemaRegistryClientConfig.BASIC_AUTH_CREDENTIALS_SOURCE, "USER_INFO");
        config.put(SchemaRegistryClientConfig.USER_INFO_CONFIG, schemaRegistryApiKey + ":" + schemaRegistryApiSecret);
        config.put("specific.avro.reader", false);

        // SSL configuration
        config.put("ssl.endpoint.identification.algorithm", "https");

        return config;
    }

    /**
     * Sends a message to a Kafka topic.
     */
    public void send(String topic, String key, String jsonPayload, Map<String, String> headers) {
        try {
            org.apache.avro.Schema schema = getSchemaForTopic(topic);
            GenericRecord record = jsonToAvro(jsonPayload, schema);

            ProducerRecord<String, GenericRecord> producerRecord =
                    new ProducerRecord<>(topic, key, record);

            // Add headers
            if (headers != null) {
                Headers kafkaHeaders = producerRecord.headers();
                headers.forEach((k, v) ->
                        kafkaHeaders.add(k, v.getBytes(StandardCharsets.UTF_8)));
            }

            // Send and wait for confirmation
            getProducer().send(producerRecord, (metadata, exception) -> {
                if (exception != null) {
                    LOG.error("Failed to send message", exception);
                } else {
                    LOG.debug("Message sent to topic {} partition {} offset {}",
                            metadata.topic(), metadata.partition(), metadata.offset());
                }
            }).get(10, java.util.concurrent.TimeUnit.SECONDS);

        } catch (ExecutionException e) {
            throw new MessagingConnectionException(
                    "Failed to send message to topic " + topic, e.getCause());
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new MessagingConnectionException(
                    "Interrupted while sending message to topic " + topic, e);
        } catch (MessagingSchemaException e) {
            throw e;
        } catch (Exception e) {
            throw new MessagingSchemaException(
                    "Failed to convert JSON to Avro for topic " + topic, e);
        }
    }

    /**
     * Receives a message from a Kafka topic matching the filter.
     */
    public List<ReceivedMessage> receive(String topic,
                                         Predicate<ReceivedMessage> filter,
                                         Duration timeout, boolean startFromBeginning) {
        KafkaConsumer<String, GenericRecord> consumer = null;
        try {
            consumer = new KafkaConsumer<>(consumerConfig);

            // Get partitions for the topic
            List<TopicPartition> partitions = getPartitionsForTopic(consumer, topic);
            if (partitions.isEmpty()) {
                throw new MessagingDestinationException(
                        "Topic '" + topic + "' does not exist or has no partitions");
            }

            consumer.assign(partitions);
            if (startFromBeginning) {
                consumer.seekToBeginning(partitions);
            } else {
                consumer.seekToEnd(partitions);
            }

            // Poll loop with exponential backoff
            long startTime = System.currentTimeMillis();
            Duration pollInterval = Duration.ofMillis(100);
            Duration maxPollInterval = Duration.ofSeconds(1);

            while (Duration.ofMillis(System.currentTimeMillis() - startTime).compareTo(timeout) < 0) {
                ConsumerRecords<String, GenericRecord> records = consumer.poll(pollInterval);

                for (ConsumerRecord<String, GenericRecord> record : records) {
                    ReceivedMessage message = convertToReceivedMessage(record, topic);
                    if (filter.test(message)) {
                        LOG.debug("Found matching message on topic {} partition {} offset {}",
                                record.topic(), record.partition(), record.offset());
                        return Collections.singletonList(message);
                    }
                }

                // Exponential backoff
                pollInterval = Duration.ofMillis(
                        Math.min(pollInterval.toMillis() * 2, maxPollInterval.toMillis()));
            }

            throw new MessagingTimeoutException(
                    "No message matching filter found on topic '" + topic + "' within " + timeout);

        } finally {
            if (consumer != null) {
                consumer.close();
            }
        }
    }

    /**
     * Gets partitions for a topic.
     */
    private List<TopicPartition> getPartitionsForTopic(KafkaConsumer<?, ?> consumer, String topic) {
        List<TopicPartition> partitions = new ArrayList<>();
        List<org.apache.kafka.common.PartitionInfo> partitionInfos = consumer.partitionsFor(topic);
        if (partitionInfos != null) {
            for (org.apache.kafka.common.PartitionInfo partitionInfo : partitionInfos) {
                partitions.add(new TopicPartition(topic, partitionInfo.partition()));
            }
        }
        return partitions;
    }

    /**
     * Saves current offsets for a topic.
     */
    public Map<TopicPartition, Long> saveOffsets(String topic) {
        return new HashMap<>();
    }

    /**
     * Closes the connector and releases resources.
     */
    @Override
    public void close() {
        if (producer != null) {
            producer.close();
        }
    }

    /**
     * Gets or creates the producer (singleton, thread-safe).
     */
    private KafkaProducer<String, GenericRecord> getProducer() {
        if (producer == null) {
            synchronized (this) {
                if (producer == null) {
                    producer = new KafkaProducer<>(producerConfig);
                }
            }
        }
        return producer;
    }

    /**
     * Retrieves schema from Schema Registry based on topic name.
     */
    private org.apache.avro.Schema getSchemaForTopic(String topic) {
        String subject = topic + "-value";
        try {
            io.confluent.kafka.schemaregistry.client.SchemaMetadata metadata =
                    schemaRegistryClient.getLatestSchemaMetadata(subject);
            return new org.apache.avro.Schema.Parser().parse(metadata.getSchema());
        } catch (Exception e) {
            if (e.getMessage() != null && e.getMessage().contains("404")) {
                throw new MessagingSchemaException(
                        "Schema not found for subject '" + subject + "' in Schema Registry. " +
                                "Make sure the topic exists and schema is registered.");
            }
            throw new MessagingSchemaException(
                    "Failed to retrieve schema for topic '" + topic + "'", e);
        }
    }

    /**
     * Converts JSON string to Avro GenericRecord.
     */
    private GenericRecord jsonToAvro(String json, org.apache.avro.Schema schema) {
        try {
            GenericRecord genericRecord = JsonToAvroConverter.processJson(json, schema);
            return genericRecord;
        } catch (Exception e) {
            throw new MessagingSchemaException("Failed to convert JSON to Avro: " + e.getMessage(), e);
        }
    }

    /**
     * Converts Kafka ConsumerRecord to ReceivedMessage.
     */
    private ReceivedMessage convertToReceivedMessage(ConsumerRecord<String, GenericRecord> record, String topic) {
        try {
            String jsonBody;
            if (null != record.value()) {
                jsonBody = avroToJson(record.value());
            } else {
                jsonBody = "";
            }
            Map<String, String> headers = new HashMap<>();
            Headers kafkaHeaders = record.headers();
            for (org.apache.kafka.common.header.Header header : kafkaHeaders) {
                if (header.value() != null) {
                    headers.put(header.key(), new String(header.value(), StandardCharsets.UTF_8));
                }
            }

            return ReceivedMessage.builder()
                    .body(jsonBody)
                    .contentType(MessageContentType.JSON)
                    .headers(headers)
                    .timestamp(record.timestamp())
                    .source(topic)
                    .key(record.key())
                    .build();

        } catch (Exception e) {
            LOG.error("Failed to convert Avro record to ReceivedMessage", e);
            throw new RuntimeException("Failed to convert message", e);
        }
    }

    /**
     * Converts Avro GenericRecord to JSON string.
     */
    private String avroToJson(GenericRecord record) {
        try {

            return record.toString();
        } catch (Exception e) {
            throw new RuntimeException("Failed to convert Avro to JSON: " + e.getMessage(), e);
        }
    }
}
