package com.dexorder.flink.kafka;

import org.apache.kafka.clients.admin.*;
import org.apache.kafka.common.config.TopicConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.yaml.snakeyaml.Yaml;
import org.yaml.snakeyaml.constructor.Constructor;

import java.io.InputStream;
import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

/**
 * Manages Kafka topics based on a YAML configuration file.
 * Creates topics if they don't exist, updates configuration if topics exist.
 */
public class TopicManager implements AutoCloseable {
    private static final Logger LOG = LoggerFactory.getLogger(TopicManager.class);

    private final AdminClient adminClient;
    private final List<com.dexorder.flink.kafka.TopicConfig> topicConfigs;

    /**
     * Creates a TopicManager with the specified Kafka bootstrap servers.
     * Loads topic configuration from the classpath resource "topics.yaml".
     *
     * @param bootstrapServers Kafka bootstrap servers
     */
    public TopicManager(String bootstrapServers) {
        this(bootstrapServers, "/topics.yaml");
    }

    /**
     * Creates a TopicManager with the specified Kafka bootstrap servers and config file.
     *
     * @param bootstrapServers Kafka bootstrap servers
     * @param configResourcePath Path to the topics YAML file in classpath
     */
    public TopicManager(String bootstrapServers, String configResourcePath) {
        Properties props = new Properties();
        props.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
        props.put(AdminClientConfig.REQUEST_TIMEOUT_MS_CONFIG, "10000");
        props.put(AdminClientConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "10000");

        this.adminClient = AdminClient.create(props);
        this.topicConfigs = loadTopicConfigs(configResourcePath);
    }

    /**
     * Loads topic configurations from a YAML file.
     */
    private List<com.dexorder.flink.kafka.TopicConfig> loadTopicConfigs(String resourcePath) {
        try (InputStream inputStream = getClass().getResourceAsStream(resourcePath)) {
            if (inputStream == null) {
                LOG.error("Topic configuration file not found: {}", resourcePath);
                return Collections.emptyList();
            }

            Yaml yaml = new Yaml(new Constructor(TopicsWrapper.class, new org.yaml.snakeyaml.LoaderOptions()));
            TopicsWrapper wrapper = yaml.load(inputStream);

            if (wrapper == null || wrapper.getTopics() == null) {
                LOG.warn("No topics defined in configuration file");
                return Collections.emptyList();
            }

            LOG.info("Loaded {} topic configurations from {}", wrapper.getTopics().size(), resourcePath);
            return wrapper.getTopics();

        } catch (Exception e) {
            LOG.error("Failed to load topic configurations from {}", resourcePath, e);
            return Collections.emptyList();
        }
    }

    /**
     * Ensures all configured topics exist with the correct configuration.
     * Creates topics that don't exist, logs warnings for topics that exist with different config.
     */
    public void ensureTopicsExist() throws ExecutionException, InterruptedException {
        if (topicConfigs.isEmpty()) {
            LOG.warn("No topics to create");
            return;
        }

        // Get existing topics
        Set<String> existingTopics = adminClient.listTopics().names().get();
        LOG.info("Found {} existing topics in Kafka", existingTopics.size());

        // Separate topics to create vs topics to check
        List<com.dexorder.flink.kafka.TopicConfig> topicsToCreate = topicConfigs.stream()
                .filter(tc -> !existingTopics.contains(tc.getName()))
                .collect(Collectors.toList());

        List<com.dexorder.flink.kafka.TopicConfig> existingConfiguredTopics = topicConfigs.stream()
                .filter(tc -> existingTopics.contains(tc.getName()))
                .collect(Collectors.toList());

        // Create new topics
        if (!topicsToCreate.isEmpty()) {
            createTopics(topicsToCreate);
        }

        // Verify existing topics
        if (!existingConfiguredTopics.isEmpty()) {
            verifyTopicConfigurations(existingConfiguredTopics);
        }

        LOG.info("Topic management complete");
    }

    /**
     * Creates the specified topics.
     */
    private void createTopics(List<com.dexorder.flink.kafka.TopicConfig> topics)
            throws ExecutionException, InterruptedException {

        List<NewTopic> newTopics = topics.stream()
                .map(tc -> {
                    NewTopic newTopic = new NewTopic(
                            tc.getName(),
                            tc.getPartitions(),
                            (short) tc.getReplication()
                    );

                    if (tc.getConfig() != null && !tc.getConfig().isEmpty()) {
                        newTopic.configs(tc.getConfig());
                    }

                    return newTopic;
                })
                .collect(Collectors.toList());

        LOG.info("Creating {} topics", newTopics.size());

        CreateTopicsResult result = adminClient.createTopics(newTopics);

        // Wait for all topics to be created and log results
        for (Map.Entry<String, org.apache.kafka.common.KafkaFuture<Void>> entry : result.values().entrySet()) {
            try {
                entry.getValue().get();
                LOG.info("Successfully created topic: {}", entry.getKey());
            } catch (ExecutionException e) {
                LOG.error("Failed to create topic: {}", entry.getKey(), e);
                throw e;
            }
        }
    }

    /**
     * Verifies that existing topics have the expected configuration.
     * Logs warnings if configuration differs.
     */
    private void verifyTopicConfigurations(List<com.dexorder.flink.kafka.TopicConfig> topics)
            throws ExecutionException, InterruptedException {

        List<String> topicNames = topics.stream()
                .map(com.dexorder.flink.kafka.TopicConfig::getName)
                .collect(Collectors.toList());

        // Describe topics to get their configurations
        DescribeTopicsResult describeResult = adminClient.describeTopics(topicNames);
        Map<String, TopicDescription> descriptions = describeResult.all().get();

        for (com.dexorder.flink.kafka.TopicConfig tc : topics) {
            TopicDescription desc = descriptions.get(tc.getName());
            if (desc == null) {
                continue;
            }

            // Check partition count
            if (desc.partitions().size() != tc.getPartitions()) {
                LOG.warn("Topic {} has {} partitions, expected {}. " +
                        "Partition count cannot be changed automatically.",
                        tc.getName(), desc.partitions().size(), tc.getPartitions());
            }

            // Check replication factor
            if (!desc.partitions().isEmpty()) {
                int actualReplication = desc.partitions().get(0).replicas().size();
                if (actualReplication != tc.getReplication()) {
                    LOG.warn("Topic {} has replication factor {}, expected {}. " +
                            "Replication factor cannot be changed automatically.",
                            tc.getName(), actualReplication, tc.getReplication());
                }
            }

            LOG.info("Verified existing topic: {}", tc.getName());
        }
    }

    /**
     * Gets the list of configured topic names.
     */
    public List<String> getTopicNames() {
        return topicConfigs.stream()
                .map(com.dexorder.flink.kafka.TopicConfig::getName)
                .collect(Collectors.toList());
    }

    @Override
    public void close() {
        if (adminClient != null) {
            adminClient.close();
        }
    }

    /**
     * Wrapper class for YAML deserialization.
     */
    public static class TopicsWrapper {
        private List<com.dexorder.flink.kafka.TopicConfig> topics;

        public List<com.dexorder.flink.kafka.TopicConfig> getTopics() {
            return topics;
        }

        public void setTopics(List<com.dexorder.flink.kafka.TopicConfig> topics) {
            this.topics = topics;
        }
    }
}
