package unit;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * CWE-407 unit test: kafka-0010
 *
 * Models RoundRobinAssignor.assign() topic membership check.
 *
 * DEFECT: The inner while loop calls subscription.topics().contains(topic)
 *         where topics() returns List<String>. For each partition being
 *         assigned, the assignor may scan multiple consumers' topic lists.
 *         Total: O(P x C x T) where P=partitions, C=consumers, T=topics/consumer.
 *
 * FIX:    Pre-convert each consumer's topic list to a HashSet<String> before
 *         the assignment loop. O(1) membership per check.
 *         Total: O(P x C) with O(1) set lookup.
 *
 * Asserts: slowOps > fastOps * 10 at C=100, T=50, P=500 (actual ratio >= 50x).
 */
public class Kafka0010RoundRobinAssignorTest {

    /**
     * Simulates the defective round-robin assignment inner loop.
     * For each partition, scans consumers until one subscribed to the topic is found.
     * topics is a List<String> — contains() is O(T).
     */
    static long slow(int numConsumers, int topicsPerConsumer, int numPartitions) {
        // Each consumer subscribes to topicsPerConsumer topics
        // Topic "topic-0" is subscribed by all consumers, others by one each
        List<List<String>> consumerTopics = new ArrayList<>();
        for (int c = 0; c < numConsumers; c++) {
            List<String> topics = new ArrayList<>();
            // Put consumer-specific topics first, shared-topic last
            // so contains("shared-topic") must scan the entire list
            for (int t = 1; t < topicsPerConsumer; t++) {
                topics.add("topic-" + c + "-" + t);
            }
            topics.add("shared-topic"); // last — worst case for linear scan
            consumerTopics.add(topics);
        }

        long ops = 0;
        // For each partition (all on "shared-topic"), scan consumers
        for (int p = 0; p < numPartitions; p++) {
            String topic = "shared-topic";
            // Round-robin: start from consumer p % numConsumers, scan until subscribed
            int startIdx = p % numConsumers;
            int idx = startIdx;
            do {
                List<String> topicList = consumerTopics.get(idx);
                // O(T) scan — the defect
                for (String t : topicList) {
                    ops++;
                    if (t.equals(topic)) break;
                }
                // This consumer is always subscribed to shared-topic (it's first in the list)
                break;
            } while (true);
        }
        return ops;
    }

    /**
     * Simulates the patched assignment loop.
     * Pre-converts each consumer's topic list to HashSet; contains() is O(1).
     */
    static long fast(int numConsumers, int topicsPerConsumer, int numPartitions) {
        // Same setup as slow()
        List<Set<String>> consumerTopicSets = new ArrayList<>();
        for (int c = 0; c < numConsumers; c++) {
            Set<String> topics = new HashSet<>();
            for (int t = 1; t < topicsPerConsumer; t++) {
                topics.add("topic-" + c + "-" + t);
            }
            topics.add("shared-topic");
            consumerTopicSets.add(topics);
        }

        long ops = 0;
        for (int p = 0; p < numPartitions; p++) {
            String topic = "shared-topic";
            int startIdx = p % numConsumers;
            int idx = startIdx;
            do {
                Set<String> topicSet = consumerTopicSets.get(idx);
                ops++; // O(1) hash probe
                if (topicSet.contains(topic)) break;
            } while (true);
        }
        return ops;
    }

    /**
     * Worst-case: topic is subscribed only by the last consumer in each scan.
     * slow() must traverse all topics for each consumer checked.
     */
    static long slowWorstCase(int numConsumers, int topicsPerConsumer, int numPartitions) {
        // Topic subscribed by last consumer only — forces full scan per partition
        List<List<String>> consumerTopics = new ArrayList<>();
        for (int c = 0; c < numConsumers; c++) {
            List<String> topics = new ArrayList<>();
            for (int t = 0; t < topicsPerConsumer; t++) {
                topics.add("topic-c" + c + "-t" + t);
            }
            if (c == numConsumers - 1) {
                topics.add("target-topic"); // only last consumer has it, placed at end
            }
            consumerTopics.add(topics);
        }

        long ops = 0;
        for (int p = 0; p < numPartitions; p++) {
            String topic = "target-topic";
            // Must scan all consumers until last one
            for (int c = 0; c < numConsumers; c++) {
                List<String> topicList = consumerTopics.get(c);
                for (String t : topicList) {
                    ops++;
                    if (t.equals(topic)) break;
                }
            }
        }
        return ops;
    }

    static long fastWorstCase(int numConsumers, int topicsPerConsumer, int numPartitions) {
        List<Set<String>> consumerTopicSets = new ArrayList<>();
        for (int c = 0; c < numConsumers; c++) {
            Set<String> topics = new HashSet<>();
            for (int t = 0; t < topicsPerConsumer; t++) {
                topics.add("topic-c" + c + "-t" + t);
            }
            if (c == numConsumers - 1) {
                topics.add("target-topic");
            }
            consumerTopicSets.add(topics);
        }

        long ops = 0;
        for (int p = 0; p < numPartitions; p++) {
            String topic = "target-topic";
            for (int c = 0; c < numConsumers; c++) {
                ops++; // O(1) hash probe per consumer
                if (consumerTopicSets.get(c).contains(topic)) break;
            }
        }
        return ops;
    }

    public static void main(String[] args) {
        int passed = 0;
        int total = 0;

        // Test 1: shared-topic, 100 consumers, 50 topics each, 500 partitions
        {
            total++;
            int C = 100, T = 50, P = 500;
            long sOps = slow(C, T, P);
            long fOps = fast(C, T, P);
            // slow: each partition scans T topics for consumer[0] = 500 * 50 = 25000
            // fast: each partition does 1 hash probe = 500
            boolean ok = sOps > fOps * 10L;
            System.out.printf("Test 1 [C=%d T=%d P=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
                    C, T, P, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
            if (ok) passed++;
        }

        // Test 2: worst case — last consumer owns target topic
        {
            total++;
            int C = 50, T = 20, P = 100;
            long sOps = slowWorstCase(C, T, P);
            long fOps = fastWorstCase(C, T, P);
            // slow: P * C * T = 100 * 50 * 20 = 100000
            // fast: P * C = 100 * 50 = 5000
            boolean ok = sOps > fOps * 10L;
            System.out.printf("Test 2 worst-case [C=%d T=%d P=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
                    C, T, P, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
            if (ok) passed++;
        }

        // Test 3: large scale worst case
        {
            total++;
            int C = 200, T = 100, P = 1000;
            long sOps = slowWorstCase(C, T, P);
            long fOps = fastWorstCase(C, T, P);
            boolean ok = sOps > fOps * 50L;
            System.out.printf("Test 3 large-scale [C=%d T=%d P=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
                    C, T, P, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
            if (ok) passed++;
        }

        // Test 4: correctness — same assignment result
        {
            total++;
            int C = 10, T = 5;
            List<List<String>> listTopics = new ArrayList<>();
            List<Set<String>> setTopics = new ArrayList<>();
            for (int c = 0; c < C; c++) {
                List<String> tl = new ArrayList<>();
                Set<String> ts = new HashSet<>();
                for (int t = 0; t < T; t++) {
                    tl.add("topic-" + t);
                    ts.add("topic-" + t);
                }
                listTopics.add(tl);
                setTopics.add(ts);
            }
            // Check contains gives same result
            boolean ok = true;
            for (int c = 0; c < C; c++) {
                for (int t = 0; t < T + 2; t++) {
                    String topic = "topic-" + t;
                    if (listTopics.get(c).contains(topic) != setTopics.get(c).contains(topic)) {
                        ok = false;
                        break;
                    }
                }
            }
            System.out.printf("Test 4 [correctness C=%d T=%d equal=%b]: %s%n",
                    C, T, ok, ok ? "PASS" : "FAIL");
            if (ok) passed++;
        }

        System.out.printf("%d/%d PASS%n", passed, total);
        if (passed != total) System.exit(1);
    }
}