package unit;

import java.util.*;

/**
 * kafka-0007: StreamsPartitionAssignor.assignTasksToThreads()
 *             PriorityQueue.contains() O(T²) → HashSet O(T)
 *
 * Demonstrates that PriorityQueue.contains() is O(N) (linear scan of heap array),
 * making the nested task-assignment loop O(C × T²) instead of O(C × T).
 *
 * Compile: javac -d . KafkaStreamsAssignTasksThreadsTest.java
 * Run:     java unit.KafkaStreamsAssignTasksThreadsTest
 */
public class KafkaStreamsAssignTasksThreadsTest {

    static class TaskId implements Comparable<TaskId> {
        final int id;
        TaskId(int id) { this.id = id; }

        @Override public int compareTo(TaskId o) { return Integer.compare(this.id, o.id); }
        @Override public boolean equals(Object o) { return o instanceof TaskId && ((TaskId)o).id == id; }
        @Override public int hashCode() { return Integer.hashCode(id); }
        @Override public String toString() { return "T" + id; }
    }

    // -------------------------------------------------------------------
    // DEFECTIVE: simulates the double-loop pattern in assignTasksToThreads
    // PriorityQueue.contains() is O(T) — linear scan of backing heap array
    // consumersToFill (LinkedList) .contains is O(C)
    // -------------------------------------------------------------------
    static long defectiveAssign(List<String> consumers, List<TaskId> prevTasksPerConsumer,
                                 List<TaskId> allTasks) {
        long ops = 0;
        final PriorityQueue<TaskId> unassigned = new PriorityQueue<>(allTasks);
        final Queue<String> consumersToFill = new LinkedList<>();
        final Map<TaskId, String> skipped = new TreeMap<>();

        // First pass: assign to previous owners
        for (String consumer : consumers) {
            for (TaskId task : prevTasksPerConsumer) {
                ops++;  // loop iteration
                // PriorityQueue.contains is O(T) — count as T ops
                ops += unassigned.size();  // simulate O(T) linear scan
                if (unassigned.contains(task)) {
                    unassigned.remove(task);
                } else {
                    skipped.put(task, consumer);
                }
            }
            consumersToFill.offer(consumer);
        }

        // Second pass: skipped tasks
        for (Map.Entry<TaskId, String> e : skipped.entrySet()) {
            TaskId task = e.getKey();
            String consumer = e.getValue();
            // LinkedList.contains is O(C) — count as C ops
            ops += consumersToFill.size();  // simulate O(C) scan
            ops += unassigned.size();        // simulate O(T) scan
            if (consumersToFill.contains(consumer) && unassigned.contains(task)) {
                unassigned.remove(task);
                consumersToFill.remove(consumer);
            }
        }
        return ops;
    }

    // -------------------------------------------------------------------
    // FIXED: HashSet for O(1) membership, LinkedHashSet for consumer fill
    // -------------------------------------------------------------------
    static long fixedAssign(List<String> consumers, List<TaskId> prevTasksPerConsumer,
                             List<TaskId> allTasks) {
        long ops = 0;
        final PriorityQueue<TaskId> unassignedOrdered = new PriorityQueue<>(allTasks);
        final Set<TaskId> unassignedSet = new HashSet<>(allTasks);
        final Set<String> consumersToFill = new LinkedHashSet<>();
        final Map<TaskId, String> skipped = new TreeMap<>();

        // First pass: assign to previous owners
        for (String consumer : consumers) {
            for (TaskId task : prevTasksPerConsumer) {
                ops++;  // loop iteration
                ops++;  // O(1) HashSet.contains
                if (unassignedSet.contains(task)) {
                    unassignedSet.remove(task);
                    unassignedOrdered.remove(task);
                } else {
                    skipped.put(task, consumer);
                }
            }
            consumersToFill.add(consumer);
        }

        // Second pass: skipped tasks
        for (Map.Entry<TaskId, String> e : skipped.entrySet()) {
            TaskId task = e.getKey();
            String consumer = e.getValue();
            ops++;  // O(1) LinkedHashSet.contains
            ops++;  // O(1) HashSet.contains
            if (consumersToFill.contains(consumer) && unassignedSet.contains(task)) {
                unassignedSet.remove(task);
                unassignedOrdered.remove(task);
                consumersToFill.remove(consumer);
            }
        }
        return ops;
    }

    static void test(String name, boolean condition) {
        System.out.println((condition ? "PASS" : "FAIL") + ": " + name);
        if (!condition) throw new AssertionError("FAIL: " + name);
    }

    public static void main(String[] args) {
        System.out.println("kafka-0007: StreamsPartitionAssignor.assignTasksToThreads()");
        System.out.println("          PriorityQueue.contains() O(T²) → HashSet O(T)");
        System.out.println("=".repeat(65));

        int[] taskSizes = {50, 100, 200, 400};
        int consumers = 5;
        boolean allPass = true;

        System.out.printf("%-8s  %-12s  %-12s  %-10s%n",
            "Tasks(T)", "Defective-ops", "Fixed-ops", "Ratio");
        System.out.println("-".repeat(50));

        long prevSlowOps = -1;
        long prevFastOps = -1;

        for (int T : taskSizes) {
            // Build T tasks
            List<TaskId> allTasks = new ArrayList<>();
            for (int i = 0; i < T; i++) allTasks.add(new TaskId(i));

            // Each consumer has T/2 previous tasks (worst case: most tasks are "known")
            List<TaskId> prev = allTasks.subList(0, T / 2);
            List<String> consumerList = new ArrayList<>();
            for (int c = 0; c < consumers; c++) consumerList.add("consumer-" + c);

            long slowOps = defectiveAssign(consumerList, prev, allTasks);
            long fastOps = fixedAssign(consumerList, prev, allTasks);
            double ratio = (double) slowOps / Math.max(fastOps, 1);

            System.out.printf("T=%-6d  %-12d  %-12d  %.1fx%n", T, slowOps, fastOps, ratio);

            boolean ratioOk = ratio >= 5.0;
            if (!ratioOk) allPass = false;

            // O(T²) scaling: doubling T should roughly 4x slow ops
            if (prevSlowOps > 0) {
                double slowScale = (double) slowOps / prevSlowOps;
                double fastScale = (double) fastOps / prevFastOps;
                boolean superlinear = slowScale >= 3.0;
                boolean linear = fastScale <= 2.5;
                if (!superlinear || !linear) allPass = false;
            }

            prevSlowOps = slowOps;
            prevFastOps = fastOps;
        }

        System.out.println("=".repeat(65));

        // Explicit PASS/FAIL tests
        {
            int T = 200;
            List<TaskId> allTasks = new ArrayList<>();
            for (int i = 0; i < T; i++) allTasks.add(new TaskId(i));
            List<TaskId> prev = allTasks.subList(0, T / 2);
            List<String> consumerList = new ArrayList<>();
            for (int c = 0; c < consumers; c++) consumerList.add("c" + c);

            long slowOps = defectiveAssign(consumerList, prev, allTasks);
            long fastOps = fixedAssign(consumerList, prev, allTasks);
            double ratio = (double) slowOps / Math.max(fastOps, 1);

            try {
                test("T=200: defective ops >> fixed ops (ratio >= 5x)", ratio >= 5.0);
                test("T=200: defective ops are super-linear (>= T*T/4)",
                    slowOps >= (long) T * T / 4);
                test("T=200: fixed ops are linear (< T*T/10)",
                    fastOps < (long) T * T / 10);
            } catch (AssertionError e) {
                allPass = false;
                System.out.println(e.getMessage());
            }
        }

        // Doubling test
        {
            int T1 = 100, T2 = 200;
            List<String> cs = Arrays.asList("c0","c1","c2","c3","c4");
            List<TaskId> tasks1 = new ArrayList<>();
            for (int i = 0; i < T1; i++) tasks1.add(new TaskId(i));
            List<TaskId> tasks2 = new ArrayList<>();
            for (int i = 0; i < T2; i++) tasks2.add(new TaskId(i));

            long s1 = defectiveAssign(cs, tasks1.subList(0, T1/2), tasks1);
            long s2 = defectiveAssign(cs, tasks2.subList(0, T2/2), tasks2);
            long f1 = fixedAssign(cs, tasks1.subList(0, T1/2), tasks1);
            long f2 = fixedAssign(cs, tasks2.subList(0, T2/2), tasks2);

            double slowScale = (double) s2 / s1;
            double fastScale = (double) f2 / f1;

            System.out.printf("Doubling T: defective-ops scale=%.2fx (expect ~4x), fixed-ops scale=%.2fx (expect ~2x)%n",
                slowScale, fastScale);

            try {
                test("Doubling T: defective ops scale >= 3x (super-linear, O(T²))", slowScale >= 3.0);
                test("Doubling T: fixed ops scale <= 2.5x (linear, O(T))", fastScale <= 2.5);
            } catch (AssertionError e) {
                allPass = false;
                System.out.println(e.getMessage());
            }
        }

        System.out.println("=".repeat(65));
        System.out.println(allPass ? "ALL PASS" : "SOME FAILED");
        if (!allPass) System.exit(1);
    }
}