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

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import cz.moneta.test.harness.connectors.Connector;
import cz.moneta.test.harness.exception.MessagingTimeoutException;
import cz.moneta.test.harness.messaging.model.MessageContentType;
import cz.moneta.test.harness.messaging.model.ReceivedMessage;
import io.confluent.kafka.schemaregistry.client.CachedSchemaRegistryClient;
import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDeConfig;
import io.confluent.kafka.serializers.KafkaAvroDeserializer;
import io.confluent.kafka.serializers.KafkaAvroDeserializerConfig;
import io.confluent.kafka.serializers.KafkaAvroSerializer;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericData;
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.Header;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Properties;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.function.Predicate;

public class KafkaConnector implements Connector {

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    private final Properties producerProps = new Properties();
    private final Properties consumerProps = new Properties();
    private final CachedSchemaRegistryClient schemaRegistryClient;
    private volatile KafkaProducer<String, GenericRecord> producer;
    private final Object producerLock = new Object();

    public KafkaConnector(String bootstrapServers,
                          String apiKey,
                          String apiSecret,
                          String schemaRegistryUrl,
                          String schemaRegistryApiKey,
                          String schemaRegistryApiSecret) {
        String jaasConfig = String.format("org.apache.kafka.common.security.plain.PlainLoginModule required username=\"%s\" password=\"%s\";",
                apiKey,
                apiSecret);
        String schemaRegistryAuth = schemaRegistryApiKey + ":" + schemaRegistryApiSecret;

        producerProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
        producerProps.put("security.protocol", "SASL_SSL");
        producerProps.put("sasl.mechanism", "PLAIN");
        producerProps.put("sasl.jaas.config", jaasConfig);
        producerProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
        producerProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, KafkaAvroSerializer.class.getName());
        producerProps.put(AbstractKafkaSchemaSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistryUrl);
        producerProps.put(AbstractKafkaSchemaSerDeConfig.BASIC_AUTH_CREDENTIALS_SOURCE, "USER_INFO");
        producerProps.put(AbstractKafkaSchemaSerDeConfig.USER_INFO_CONFIG, schemaRegistryAuth);
        producerProps.put("auto.register.schemas", "false");

        consumerProps.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
        consumerProps.put("security.protocol", "SASL_SSL");
        consumerProps.put("sasl.mechanism", "PLAIN");
        consumerProps.put("sasl.jaas.config", jaasConfig);
        consumerProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        consumerProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, KafkaAvroDeserializer.class.getName());
        consumerProps.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "latest");
        consumerProps.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
        consumerProps.put(KafkaAvroDeserializerConfig.SPECIFIC_AVRO_READER_CONFIG, "false");
        consumerProps.put(AbstractKafkaSchemaSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistryUrl);
        consumerProps.put(AbstractKafkaSchemaSerDeConfig.BASIC_AUTH_CREDENTIALS_SOURCE, "USER_INFO");
        consumerProps.put(AbstractKafkaSchemaSerDeConfig.USER_INFO_CONFIG, schemaRegistryAuth);

        this.schemaRegistryClient = new CachedSchemaRegistryClient(schemaRegistryUrl, 128);
    }

    public void send(String topic, String key, String jsonPayload, Map<String, String> headers) {
        Objects.requireNonNull(topic, "topic");
        Schema schema = getSchemaForTopic(topic);
        GenericRecord record = jsonToAvro(jsonPayload, schema);
        ProducerRecord<String, GenericRecord> producerRecord = new ProducerRecord<>(topic, key, record);
        Optional.ofNullable(headers).orElseGet(HashMap::new)
                .forEach((headerKey, headerValue) -> producerRecord.headers()
                        .add(headerKey, String.valueOf(headerValue).getBytes(StandardCharsets.UTF_8)));

        try {
            getProducer().send(producerRecord).get(30, TimeUnit.SECONDS);
        } catch (Exception e) {
            throw new IllegalStateException("Failed to send Kafka message to topic: " + topic, e);
        }
    }

    public List<ReceivedMessage> receive(String topic,
                                         Predicate<ReceivedMessage> filter,
                                         Duration timeout) {
        long timeoutMillis = Optional.ofNullable(timeout).orElse(Duration.ofSeconds(30)).toMillis();
        long deadline = System.currentTimeMillis() + timeoutMillis;
        long backoff = 100;

        try (KafkaConsumer<String, Object> consumer = createConsumer()) {
            List<TopicPartition> partitions = consumer.partitionsFor(topic).stream()
                    .map(info -> new TopicPartition(topic, info.partition()))
                    .toList();
            consumer.assign(partitions);
            consumer.seekToEnd(partitions);

            while (System.currentTimeMillis() < deadline) {
                ConsumerRecords<String, Object> records = consumer.poll(Duration.ofMillis(100));
                if (!records.isEmpty()) {
                    for (ConsumerRecord<String, Object> record : records) {
                        ReceivedMessage message = toReceivedMessage(record);
                        if (filter == null || filter.test(message)) {
                            return List.of(message);
                        }
                    }
                    backoff = 100;
                    continue;
                }
                TimeUnit.MILLISECONDS.sleep(backoff);
                backoff = Math.min(backoff * 2, 1000);
            }
        } catch (MessagingTimeoutException e) {
            throw e;
        } catch (Exception e) {
            throw new IllegalStateException("Failed to receive Kafka message from topic: " + topic, e);
        }
        throw new MessagingTimeoutException("Timeout waiting for Kafka message from topic: " + topic);
    }

    public Map<TopicPartition, Long> saveOffsets(String topic) {
        try (KafkaConsumer<String, Object> consumer = createConsumer()) {
            List<TopicPartition> partitions = consumer.partitionsFor(topic).stream()
                    .map(info -> new TopicPartition(topic, info.partition()))
                    .toList();
            consumer.assign(partitions);
            consumer.seekToEnd(partitions);
            Map<TopicPartition, Long> offsets = new HashMap<>();
            partitions.forEach(partition -> offsets.put(partition, consumer.position(partition)));
            return offsets;
        }
    }

    @Override
    public void close() {
        KafkaProducer<String, GenericRecord> current = producer;
        if (current != null) {
            current.close(Duration.ofSeconds(5));
        }
    }

    private KafkaProducer<String, GenericRecord> getProducer() {
        KafkaProducer<String, GenericRecord> current = producer;
        if (current == null) {
            synchronized (producerLock) {
                current = producer;
                if (current == null) {
                    producer = current = new KafkaProducer<>(producerProps);
                }
            }
        }
        return current;
    }

    private KafkaConsumer<String, Object> createConsumer() {
        Properties properties = new Properties();
        properties.putAll(consumerProps);
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, "harness-" + UUID.randomUUID());
        return new KafkaConsumer<>(properties);
    }

    private ReceivedMessage toReceivedMessage(ConsumerRecord<String, Object> record) {
        String body = convertValueToJson(record.value());
        Map<String, String> headers = new LinkedHashMap<>();
        for (Header header : record.headers()) {
            headers.put(header.key(), new String(header.value(), StandardCharsets.UTF_8));
        }

        return new ReceivedMessage(body,
                MessageContentType.JSON,
                headers,
                record.timestamp(),
                record.topic());
    }

    private String convertValueToJson(Object value) {
        try {
            if (value instanceof GenericRecord genericRecord) {
                return avroToJson(genericRecord);
            }
            if (value instanceof CharSequence) {
                return value.toString();
            }
            return OBJECT_MAPPER.writeValueAsString(value);
        } catch (Exception e) {
            throw new IllegalStateException("Failed to convert Kafka payload to JSON", e);
        }
    }

    private Schema getSchemaForTopic(String topic) {
        String subject = topic + "-value";
        try {
            // Get all versions and use the latest one
            java.util.List<Integer> versions = schemaRegistryClient.getAllVersions(subject);
            int latestVersion = versions.get(versions.size() - 1);
            io.confluent.kafka.schemaregistry.client.rest.entities.Schema confluentSchema = 
                schemaRegistryClient.getByVersion(subject, latestVersion, false);
            String schemaString = confluentSchema.getSchema();
            return new Schema.Parser().parse(schemaString);
        } catch (Exception e) {
            throw new IllegalStateException("Failed to get schema for subject: " + subject, e);
        }
    }

    private String avroToJson(GenericRecord record) {
        try {
            return OBJECT_MAPPER.writeValueAsString(convertAvroObject(record));
        } catch (Exception e) {
            throw new IllegalStateException("Failed to convert Avro record to JSON", e);
        }
    }

    private GenericRecord jsonToAvro(String jsonPayload, Schema schema) {
        try {
            JsonNode root = OBJECT_MAPPER.readTree(jsonPayload);
            Object converted = convertJsonNode(root, schema);
            return (GenericRecord) converted;
        } catch (Exception e) {
            throw new IllegalStateException("Failed to convert JSON payload to Avro", e);
        }
    }

    private Object convertJsonNode(JsonNode node, Schema schema) {
        return switch (schema.getType()) {
            case RECORD -> {
                GenericData.Record record = new GenericData.Record(schema);
                schema.getFields().forEach(field -> record.put(field.name(),
                        convertJsonNode(node.path(field.name()), field.schema())));
                yield record;
            }
            case ARRAY -> {
                List<Object> values = new ArrayList<>();
                node.forEach(item -> values.add(convertJsonNode(item, schema.getElementType())));
                yield values;
            }
            case MAP -> {
                Map<String, Object> map = new HashMap<>();
                node.fields().forEachRemaining(entry -> map.put(entry.getKey(),
                        convertJsonNode(entry.getValue(), schema.getValueType())));
                yield map;
            }
            case UNION -> resolveUnion(node, schema);
            case ENUM -> new GenericData.EnumSymbol(schema, node.asText());
            case FIXED -> {
                byte[] fixedBytes = toBytes(node);
                yield new GenericData.Fixed(schema, fixedBytes);
            }
            case STRING -> node.isNull() ? null : node.asText();
            case INT -> node.isNull() ? null : node.asInt();
            case LONG -> node.isNull() ? null : node.asLong();
            case FLOAT -> node.isNull() ? null : (float) node.asDouble();
            case DOUBLE -> node.isNull() ? null : node.asDouble();
            case BOOLEAN -> node.isNull() ? null : node.asBoolean();
            case BYTES -> ByteBuffer.wrap(toBytes(node));
            case NULL -> null;
        };
    }

    private Object resolveUnion(JsonNode node, Schema unionSchema) {
        if (node == null || node.isNull()) {
            return null;
        }

        IllegalStateException lastException = null;
        for (Schema candidate : unionSchema.getTypes()) {
            if (candidate.getType() == Schema.Type.NULL) {
                continue;
            }
            try {
                Object value = convertJsonNode(node, candidate);
                if (GenericData.get().validate(candidate, value)) {
                    return value;
                }
            } catch (Exception e) {
                lastException = new IllegalStateException("Failed to resolve union type", e);
            }
        }

        if (lastException != null) {
            throw lastException;
        }
        throw new IllegalStateException("Cannot resolve union for node: " + node);
    }

    private byte[] toBytes(JsonNode node) {
        if (node.isBinary()) {
            try {
                return node.binaryValue();
            } catch (Exception ignored) {
                // fallback to textual representation
            }
        }
        String text = node.asText();
        try {
            return Base64.getDecoder().decode(text);
        } catch (Exception ignored) {
            return text.getBytes(StandardCharsets.UTF_8);
        }
    }

    private Object convertAvroObject(Object value) {
        if (value == null) {
            return null;
        }
        if (value instanceof GenericRecord record) {
            Map<String, Object> jsonObject = new LinkedHashMap<>();
            record.getSchema().getFields().forEach(field -> jsonObject.put(field.name(), convertAvroObject(record.get(field.name()))));
            return jsonObject;
        }
        if (value instanceof GenericData.Array<?> array) {
            List<Object> values = new ArrayList<>(array.size());
            array.forEach(item -> values.add(convertAvroObject(item)));
            return values;
        }
        if (value instanceof Map<?, ?> map) {
            Map<String, Object> values = new LinkedHashMap<>();
            map.forEach((key, item) -> values.put(String.valueOf(key), convertAvroObject(item)));
            return values;
        }
        if (value instanceof ByteBuffer byteBuffer) {
            ByteBuffer duplicate = byteBuffer.duplicate();
            byte[] bytes = new byte[duplicate.remaining()];
            duplicate.get(bytes);
            return Base64.getEncoder().encodeToString(bytes);
        }
        return value;
    }
}
