package unit;

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

/**
 * CWE-407 unit test: kafka-0011
 *
 * Models AbstractStickyAssignor.GeneralAssignmentBuilder.assignOwnedPartitions()
 * topic subscription check.
 *
 * DEFECT: For each consumer's owned partition, calls
 *         consumerSubscription.topics().contains(partition.topic())
 *         where topics() returns List<String>.
 *         Total: O(C x P x T) where C=consumers, P=partitions/consumer, T=topics.
 *
 * FIX:    Build a Map<String, Set<String>> consumer->topicSet before the loop.
 *         O(1) set.contains() per partition check.
 *
 * Asserts: slowOps > fastOps * 10 at C=100, P=50, T=50.
 */
public class Kafka0011StickyAssignorTest {

    /**
     * Simulates the defective assignOwnedPartitions loop.
     * For each consumer, for each of their owned partitions, checks
     * if the topic is still in their subscription (List<String>).
     */
    static long slow(int numConsumers, int partitionsPerConsumer, int topicsPerConsumer) {
        // Setup: each consumer owns partitionsPerConsumer partitions,
        // subscribed to topicsPerConsumer topics.
        List<List<String>> consumerTopics = new ArrayList<>();
        List<List<String>> ownedPartitionTopics = new ArrayList<>();

        for (int c = 0; c < numConsumers; c++) {
            List<String> topics = new ArrayList<>();
            for (int t = 0; t < topicsPerConsumer; t++) {
                topics.add("topic-" + t);
            }
            consumerTopics.add(topics);

            List<String> partitions = new ArrayList<>();
            for (int p = 0; p < partitionsPerConsumer; p++) {
                // Partitions are spread across topics; topic-(p % topicsPerConsumer)
                partitions.add("topic-" + (p % topicsPerConsumer));
            }
            ownedPartitionTopics.add(partitions);
        }

        long ops = 0;
        for (int c = 0; c < numConsumers; c++) {
            List<String> topicList = consumerTopics.get(c);
            for (String partitionTopic : ownedPartitionTopics.get(c)) {
                // O(T) scan — the defect
                for (String t : topicList) {
                    ops++;
                    if (t.equals(partitionTopic)) break;
                }
            }
        }
        return ops;
    }

    /**
     * Simulates the patched version.
     * Pre-converts each consumer's topic list to HashSet; O(1) per check.
     */
    static long fast(int numConsumers, int partitionsPerConsumer, int topicsPerConsumer) {
        List<Set<String>> consumerTopicSets = new ArrayList<>();
        List<List<String>> ownedPartitionTopics = new ArrayList<>();

        for (int c = 0; c < numConsumers; c++) {
            Set<String> topicSet = new HashSet<>();
            for (int t = 0; t < topicsPerConsumer; t++) {
                topicSet.add("topic-" + t);
            }
            consumerTopicSets.add(topicSet);

            List<String> partitions = new ArrayList<>();
            for (int p = 0; p < partitionsPerConsumer; p++) {
                partitions.add("topic-" + (p % topicsPerConsumer));
            }
            ownedPartitionTopics.add(partitions);
        }

        long ops = 0;
        for (int c = 0; c < numConsumers; c++) {
            Set<String> topicSet = consumerTopicSets.get(c);
            for (String partitionTopic : ownedPartitionTopics.get(c)) {
                ops++; // O(1) hash probe
                topicSet.contains(partitionTopic);
            }
        }
        return ops;
    }

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

        // Test 1: 100 consumers, 50 partitions each, 50 topics each
        {
            total++;
            int C = 100, P = 50, T = 50;
            long sOps = slow(C, P, T);
            long fOps = fast(C, P, T);
            // slow: C * P * (avg scan depth) ≈ 100 * 50 * 25 = 125000
            // fast: C * P = 5000
            boolean ok = sOps > fOps * 10L;
            System.out.printf("Test 1 [C=%d P=%d T=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
                    C, P, T, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
            if (ok) passed++;
        }

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

        // Test 3: worst case — topics in partition list always last in topic list
        {
            total++;
            int C = 50, T = 100;
            // Each partition's topic is last in its consumer's topic list
            List<List<String>> consumerTopics = new ArrayList<>();
            List<List<String>> owned = new ArrayList<>();
            for (int c = 0; c < C; c++) {
                List<String> tl = new ArrayList<>();
                for (int t = 0; t < T; t++) tl.add("other-" + t);
                tl.add("last-topic");
                consumerTopics.add(tl);
                List<String> pl = new ArrayList<>();
                for (int p = 0; p < 50; p++) pl.add("last-topic");
                owned.add(pl);
            }

            long sOps = 0;
            for (int c = 0; c < C; c++) {
                List<String> topicList = consumerTopics.get(c);
                for (String partTopic : owned.get(c)) {
                    for (String t : topicList) { sOps++; if (t.equals(partTopic)) break; }
                }
            }
            // fast: pre-hashed
            long fOps = 0;
            List<Set<String>> sets = new ArrayList<>();
            for (int c = 0; c < C; c++) {
                Set<String> ts = new HashSet<>(consumerTopics.get(c));
                sets.add(ts);
            }
            for (int c = 0; c < C; c++) {
                for (String pt : owned.get(c)) { fOps++; sets.get(c).contains(pt); }
            }

            boolean ok = sOps > fOps * 50L;
            System.out.printf("Test 3 worst-case [C=%d T=%d slow=%d fast=%d ratio=%.1fx]: %s%n",
                    C, T, sOps, fOps, (double) sOps / fOps, ok ? "PASS" : "FAIL");
            if (ok) passed++;
        }

        // Test 4: correctness — both give same membership results
        {
            total++;
            int C = 20, T = 30;
            boolean ok = true;
            for (int c = 0; c < C && ok; c++) {
                List<String> tl = new ArrayList<>();
                Set<String> ts = new HashSet<>();
                for (int t = 0; t < T; t++) {
                    tl.add("t" + t);
                    ts.add("t" + t);
                }
                for (int p = 0; p < 40; p++) {
                    String pt = "t" + (p % (T + 5));
                    if (tl.contains(pt) != ts.contains(pt)) {
                        ok = false;
                    }
                }
            }
            System.out.printf("Test 4 [correctness C=%d T=%d ok=%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);
    }
}