import java.util.*;

/**
 * CWE-407 unit tests for Apache Kafka consumer rebalance defects.
 *
 * kafka-0001: AbstractStickyAssignor.isBalanced() — currentAssignment values are
 *   ArrayList<TopicPartition>; .contains() is O(P/C) per call inside triple loop → O(P³/C).
 *   Fix: parallel Map<String, Set<TopicPartition>> currentAssignmentSet for O(1) contains.
 *
 * kafka-0002: AbstractStickyAssignor — consumer2AllPotentialTopics values are ArrayList<String>;
 *   .contains(topic) called O(P×C) times → O(P×C×T) per rebalance.
 *   Fix: values as HashSet<String> for O(1) per lookup.
 *
 * kafka-0003: RoundRobinAssignor.assign() — Subscription.topics().contains(topic)
 *   is O(T) List scan in while-loop; up to C consumers skipped per partition → O(P×C×T).
 *   Fix: pre-compute Map<memberId, Set<String>> subscribedTopicsSet; O(1) per lookup.
 */
public class KafkaTest {

    // --- kafka-0001 ---

    static boolean isBalanced_list(Map<String, List<Integer>> assignment, int totalPartitions) {
        for (Map.Entry<String, List<Integer>> e : assignment.entrySet()) {
            List<Integer> consumerParts = e.getValue();
            for (int p = 0; p < totalPartitions; p++) {
                if (!consumerParts.contains(p)) {  // O(P/C) — defect
                    // check balance
                }
            }
        }
        return true;
    }

    static boolean isBalanced_set(Map<String, Set<Integer>> assignmentSet, int totalPartitions) {
        for (Map.Entry<String, Set<Integer>> e : assignmentSet.entrySet()) {
            Set<Integer> consumerParts = e.getValue();
            for (int p = 0; p < totalPartitions; p++) {
                if (!consumerParts.contains(p)) {  // O(1) — fix
                    // check balance
                }
            }
        }
        return true;
    }

    static void testKafka0001() throws Exception {
        int C = 5;     // consumers
        int P = 5000;  // partitions — P/C=1000 items per list, O(1000) contains vs O(1)

        Map<String, List<Integer>> assignList = new HashMap<>();
        Map<String, Set<Integer>> assignSet = new HashMap<>();
        for (int c = 0; c < C; c++) {
            List<Integer> parts = new ArrayList<>();
            Set<Integer> partsSet = new HashSet<>();
            for (int p = c; p < P; p += C) { parts.add(p); partsSet.add(p); }
            assignList.put("consumer-" + c, parts);
            assignSet.put("consumer-" + c, partsSet);
        }

        // correctness
        assert isBalanced_list(assignList, P) == isBalanced_set(assignSet, P);

        int REPS = 20;
        long t0 = System.nanoTime();
        for (int r = 0; r < REPS; r++) isBalanced_list(assignList, P);
        long tList = System.nanoTime() - t0;

        t0 = System.nanoTime();
        for (int r = 0; r < REPS; r++) isBalanced_set(assignSet, P);
        long tSet = System.nanoTime() - t0;

        double ratio = (double) tList / tSet;
        System.out.printf("kafka-0001: list=%.3fs set=%.3fs ratio=%.1f×%n",
                tList / 1e9, tSet / 1e9, ratio);
        assert ratio > 5 : "Expected >5× speedup, got " + ratio;
        System.out.println("PASS kafka-0001");
    }

    // --- kafka-0002 ---

    static boolean topicSubscribed_list(Map<String, List<String>> consumer2Topics,
                                        String consumer, String topic) {
        return consumer2Topics.get(consumer).contains(topic);  // O(T) — defect
    }

    static boolean topicSubscribed_set(Map<String, Set<String>> consumer2TopicsSet,
                                       String consumer, String topic) {
        return consumer2TopicsSet.get(consumer).contains(topic);  // O(1) — fix
    }

    static void testKafka0002() throws Exception {
        int C = 50;    // consumers
        int T = 50;    // topics per consumer
        int P = 1000;  // partitions to assign

        Map<String, List<String>> c2t_list = new HashMap<>();
        Map<String, Set<String>> c2t_set = new HashMap<>();
        for (int c = 0; c < C; c++) {
            List<String> topics = new ArrayList<>(T);
            Set<String> topicSet = new HashSet<>(T);
            for (int t = 0; t < T; t++) { topics.add("topic-" + t); topicSet.add("topic-" + t); }
            c2t_list.put("consumer-" + c, topics);
            c2t_set.put("consumer-" + c, topicSet);
        }

        String targetConsumer = "consumer-0";
        String targetTopic = "topic-" + (T - 1);

        assert topicSubscribed_list(c2t_list, targetConsumer, targetTopic) ==
               topicSubscribed_set(c2t_set, targetConsumer, targetTopic);

        // Simulate assignment loop: O(P × C) lookups
        int REPS = 200;
        long t0 = System.nanoTime();
        for (int r = 0; r < REPS; r++) {
            for (int p = 0; p < P; p++) {
                for (int c = 0; c < C; c++) {
                    topicSubscribed_list(c2t_list, "consumer-" + c, targetTopic);
                }
            }
        }
        long tList = System.nanoTime() - t0;

        t0 = System.nanoTime();
        for (int r = 0; r < REPS; r++) {
            for (int p = 0; p < P; p++) {
                for (int c = 0; c < C; c++) {
                    topicSubscribed_set(c2t_set, "consumer-" + c, targetTopic);
                }
            }
        }
        long tSet = System.nanoTime() - t0;

        double ratio = (double) tList / tSet;
        System.out.printf("kafka-0002: list=%.3fs set=%.3fs ratio=%.1f×%n",
                tList / 1e9, tSet / 1e9, ratio);
        assert ratio > 5 : "Expected >5× speedup, got " + ratio;
        System.out.println("PASS kafka-0002");
    }

    // --- kafka-0003 ---

    static boolean topicsContains_list(List<String> topics, String topic) {
        return topics.contains(topic);   // O(T) — defect
    }

    static boolean topicsContains_set(Set<String> topicSet, String topic) {
        return topicSet.contains(topic); // O(1) — fix
    }

    static void testKafka0003() throws Exception {
        int C = 50;    // consumers in round-robin
        int T = 50;    // topics per consumer subscription
        int P = 5000;  // partitions to assign

        List<List<String>> allLists = new ArrayList<>(C);
        List<Set<String>> allSets = new ArrayList<>(C);
        for (int c = 0; c < C; c++) {
            List<String> lst = new ArrayList<>(T);
            Set<String> set = new HashSet<>(T);
            for (int t = 0; t < T; t++) { lst.add("topic-" + t); set.add("topic-" + t); }
            allLists.add(lst);
            allSets.add(set);
        }

        String lastTopic = "topic-" + (T - 1);

        // correctness
        assert topicsContains_list(allLists.get(0), lastTopic) ==
               topicsContains_set(allSets.get(0), lastTopic);

        // Simulate round-robin loop: O(P × C) checks
        int REPS = 50;
        long t0 = System.nanoTime();
        for (int r = 0; r < REPS; r++) {
            for (int p = 0; p < P; p++) {
                for (int c = 0; c < C; c++) topicsContains_list(allLists.get(c), lastTopic);
            }
        }
        long tList = System.nanoTime() - t0;

        t0 = System.nanoTime();
        for (int r = 0; r < REPS; r++) {
            for (int p = 0; p < P; p++) {
                for (int c = 0; c < C; c++) topicsContains_set(allSets.get(c), lastTopic);
            }
        }
        long tSet = System.nanoTime() - t0;

        double ratio = (double) tList / tSet;
        System.out.printf("kafka-0003: list=%.3fs set=%.3fs ratio=%.1f×%n",
                tList / 1e9, tSet / 1e9, ratio);
        assert ratio > 5 : "Expected >5× speedup, got " + ratio;
        System.out.println("PASS kafka-0003");
    }

    public static void main(String[] args) throws Exception {
        testKafka0001();
        testKafka0002();
        testKafka0003();
        System.out.println("ALL PASS");
    }
}