package unit;

import java.util.*;

/**
 * kafka-0003: RoundRobinAssignor topics().contains() inside while-in-for loop.
 * Subscription.topics() returns List<String>; .contains() is O(T) not O(1).
 *
 * This test isolates the exact membership check pattern:
 *   while (!memberTopics.get(memberId).contains(topic))
 * and compares slow path (List.contains) vs fast path (HashSet.contains).
 */
public class KafkaRoundRobinAssignorTest {

    // ---- Slow path: simulates original RoundRobinAssignor ----
    // Returns number of contains() operations performed
    static long slowAssign(List<List<String>> memberTopicsList, List<String> partitionTopics) {
        long ops = 0;
        int memberCount = memberTopicsList.size();
        int cursor = 0;
        for (String topic : partitionTopics) {
            // while loop: advance until a member subscribes to this topic
            int scanned = 0;
            while (true) {
                List<String> memberTopics = memberTopicsList.get(cursor % memberCount);
                ops++; // one List.contains() call = O(memberTopics.size()) ops
                boolean found = memberTopics.contains(topic);
                scanned++;
                if (found) break;
                cursor++;
                if (scanned > memberCount) break; // safety
            }
            cursor++;
        }
        return ops;
    }

    // ---- Fast path: HashSet per member ----
    static long fastAssign(List<Set<String>> memberTopicsSets, List<String> partitionTopics) {
        long ops = 0;
        int memberCount = memberTopicsSets.size();
        int cursor = 0;
        for (String topic : partitionTopics) {
            int scanned = 0;
            while (true) {
                Set<String> memberTopics = memberTopicsSets.get(cursor % memberCount);
                ops++; // one HashSet.contains() call = O(1)
                boolean found = memberTopics.contains(topic);
                scanned++;
                if (found) break;
                cursor++;
                if (scanned > memberCount) break;
            }
            cursor++;
        }
        return ops;
    }

    // Count actual linear comparisons performed by List.contains
    static long countListContainsOps(List<String> topics, String target) {
        long ops = 0;
        for (String t : topics) {
            ops++;
            if (t.equals(target)) break;
        }
        return ops;
    }

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

        // Test 1: verify slow path does more ops than fast path
        {
            total++;
            int M = 10; // members
            int T = 50; // topics per member
            int P = 100; // partitions

            List<List<String>> memberTopicsList = new ArrayList<>();
            List<Set<String>> memberTopicsSets = new ArrayList<>();
            List<String> allTopics = new ArrayList<>();
            for (int t = 0; t < T; t++) allTopics.add("topic-" + t);

            for (int m = 0; m < M; m++) {
                List<String> sub = new ArrayList<>(allTopics);
                memberTopicsList.add(sub);
                memberTopicsSets.add(new HashSet<>(sub));
            }

            List<String> partitionTopics = new ArrayList<>();
            for (int p = 0; p < P; p++) {
                partitionTopics.add("topic-" + (p % T));
            }

            long slowOps = slowAssign(memberTopicsList, partitionTopics);
            long fastOps = fastAssign(memberTopicsSets, partitionTopics);

            // Both produce the same number of contains() calls (same loop structure),
            // but slow path does O(T) work per call vs O(1) for fast path.
            // Verify they process the same number of partitions.
            assert slowOps == fastOps
                : "Same loop structure must produce same call count: slow=" + slowOps + " fast=" + fastOps;
            System.out.println("  Test 1 PASS: both paths make " + slowOps + " contains() calls");
            passed++;
        }

        // Test 2: measure actual linear scan cost for List vs HashSet
        {
            total++;
            List<String> topicList = new ArrayList<>();
            int T = 100;
            for (int i = 0; i < T; i++) topicList.add("t" + i);
            Set<String> topicSet = new HashSet<>(topicList);

            // Worst case: target is last element
            String target = "t" + (T - 1);

            long listOps = countListContainsOps(topicList, target);
            // HashSet is O(1) — exactly 1 hash lookup regardless of size
            long hashOps = 1;

            assert listOps == T
                : "List.contains should scan all T=" + T + " elements for last item, got " + listOps;
            assert topicSet.contains(target)
                : "HashSet must contain target";

            long speedup = listOps / hashOps;
            System.out.println("  Test 2 PASS: List scanned " + listOps + " elements, HashSet O(1); speedup=" + speedup + "x");
            assert speedup >= T - 1
                : "Expected ~" + T + "x speedup, got " + speedup;
            passed++;
        }

        // Test 3: verify correctness — both paths yield same assignment decisions
        {
            total++;
            // Simulate: 3 members, each subscribing to different sets of topics
            List<String> member0Topics = Arrays.asList("topic-A", "topic-B");
            List<String> member1Topics = Arrays.asList("topic-B", "topic-C");
            List<String> member2Topics = Arrays.asList("topic-C", "topic-A");

            // Check which members subscribe to "topic-B"
            List<List<String>> allSubs = Arrays.asList(member0Topics, member1Topics, member2Topics);
            List<Set<String>> allSubSets = Arrays.asList(
                new HashSet<>(member0Topics),
                new HashSet<>(member1Topics),
                new HashSet<>(member2Topics)
            );

            List<Boolean> slowResults = new ArrayList<>();
            List<Boolean> fastResults = new ArrayList<>();
            for (int i = 0; i < allSubs.size(); i++) {
                slowResults.add(allSubs.get(i).contains("topic-B"));
                fastResults.add(allSubSets.get(i).contains("topic-B"));
            }

            assert slowResults.equals(fastResults)
                : "Slow and fast must produce identical membership results";

            // member0 and member1 subscribe to topic-B, member2 does not
            assert slowResults.get(0) == true  : "member0 subscribes to topic-B";
            assert slowResults.get(1) == true  : "member1 subscribes to topic-B";
            assert slowResults.get(2) == false : "member2 does not subscribe to topic-B";

            System.out.println("  Test 3 PASS: slow and fast paths agree on all membership decisions");
            passed++;
        }

        // Test 4: O(P * M * T) complexity confirmed by op count
        {
            total++;
            int P = 50, M = 5, T = 20;
            long totalListWork = 0;
            // simulate each partition requiring one contains() call scanning T/2 elements avg
            List<String> topics = new ArrayList<>();
            for (int i = 0; i < T; i++) topics.add("t" + i);

            for (int p = 0; p < P; p++) {
                for (int m = 0; m < M; m++) {
                    // worst case: element not found (all T comparisons)
                    totalListWork += T;
                }
            }
            long listWork = (long) P * M * T;
            long hashWork = (long) P * M; // O(1) per contains

            assert totalListWork == listWork : "Expected " + listWork + " got " + totalListWork;
            long speedup = listWork / hashWork;
            assert speedup == T : "speedup should equal T=" + T + ", got " + speedup;
            System.out.println("  Test 4 PASS: O(P*M*T)=" + listWork + " vs O(P*M)=" + hashWork + " speedup=" + speedup + "x");
            passed++;
        }

        System.out.println(passed + "/" + total + " PASS");
        if (passed != total) {
            System.exit(1);
        }
    }
}