import java.util.*;

/**
 * CWE-407 unit test for Apache Storm deeper scan (storm-0001).
 *
 * storm-0001: WorkerState.refreshConnections()
 *   storm-client/src/jvm/org/apache/storm/daemon/worker/WorkerState.java
 *
 *   localTaskIds is declared as ArrayList<Integer> and populated from the worker's executors.
 *   In refreshConnections() (called on a recurring timer in every Storm worker process),
 *   the code iterates over all outbound tasks and calls localTaskIds.contains(task):
 *
 *     for (Map.Entry<Integer, NodeInfo> entry : taskToNodePort.entrySet()) {  // O(T_out)
 *         Integer task = entry.getKey();
 *         if (outboundTasks.contains(task)) {         // O(1) — Set
 *             ...
 *             if (!localTaskIds.contains(task)) {     // O(L) — ArrayList linear scan — defect
 *                 neededConnections.add(...);
 *             }
 *
 *   Total per refresh: O(T_out * L) where T_out = outbound task count, L = local task count.
 *   Since refreshConnections() is called every second (TOPOLOGY_WORKER_RECEIVER_THREAD_COUNT
 *   timer), this is a sustained hot path in every running Storm worker.
 *
 *   Fix: change localTaskIds from ArrayList<Integer> to HashSet<Integer>.
 *   The Collections.sort(localTaskIds) call (used for deterministic ordering) is removed;
 *   callers that require ordered output (e.g. getLocalTaskIds()) can return a sorted copy
 *   on demand when needed.
 */
public class StormTest {

    /** Simulate defect: ArrayList.contains in refreshConnections loop */
    static int simulateRefreshConnections_defect(
            List<Integer> localTaskIds,      // ArrayList — defect
            Set<Integer> outboundTasks,
            Map<Integer, String> taskToNodePort) {
        Set<String> neededConnections = new HashSet<>();
        for (Map.Entry<Integer, String> entry : taskToNodePort.entrySet()) {
            Integer task = entry.getKey();
            if (outboundTasks.contains(task)) {          // O(1) — HashSet
                if (!localTaskIds.contains(task)) {       // O(L) — defect
                    neededConnections.add(entry.getValue());
                }
            }
        }
        return neededConnections.size();
    }

    /** Simulate fix: HashSet.contains is O(1) */
    static int simulateRefreshConnections_fixed(
            Set<Integer> localTaskIds,       // HashSet — fix
            Set<Integer> outboundTasks,
            Map<Integer, String> taskToNodePort) {
        Set<String> neededConnections = new HashSet<>();
        for (Map.Entry<Integer, String> entry : taskToNodePort.entrySet()) {
            Integer task = entry.getKey();
            if (outboundTasks.contains(task)) {          // O(1)
                if (!localTaskIds.contains(task)) {       // O(1) — fix
                    neededConnections.add(entry.getValue());
                }
            }
        }
        return neededConnections.size();
    }

    static void testStorm0001() throws Exception {
        // T_out = outbound tasks (topology-wide tasks visible to worker)
        // L     = local tasks per worker
        int T_out = 1000;
        int L = 20;

        // Populate local task IDs
        List<Integer> localTaskIdsList = new ArrayList<>();
        Set<Integer> localTaskIdsSet = new HashSet<>();
        for (int i = 0; i < L; i++) {
            localTaskIdsList.add(i);
            localTaskIdsSet.add(i);
        }

        // outboundTasks is a HashSet of all tasks this worker routes to
        Set<Integer> outboundTasks = new HashSet<>();
        for (int i = 0; i < T_out; i++) outboundTasks.add(i);

        // taskToNodePort: topology assignment map (T_out entries)
        Map<Integer, String> taskToNodePort = new LinkedHashMap<>();
        for (int i = 0; i < T_out; i++) {
            taskToNodePort.put(i, "node" + (i % 50) + ":6700");
        }

        // Warm up
        simulateRefreshConnections_defect(localTaskIdsList, outboundTasks, taskToNodePort);
        simulateRefreshConnections_fixed(localTaskIdsSet, outboundTasks, taskToNodePort);

        // Simulate timer-based refresh: many iterations
        int RUNS = 5000;

        long t0 = System.nanoTime();
        for (int r = 0; r < RUNS; r++) {
            simulateRefreshConnections_defect(localTaskIdsList, outboundTasks, taskToNodePort);
        }
        long defectNs = System.nanoTime() - t0;

        long t1 = System.nanoTime();
        for (int r = 0; r < RUNS; r++) {
            simulateRefreshConnections_fixed(localTaskIdsSet, outboundTasks, taskToNodePort);
        }
        long fixedNs = System.nanoTime() - t1;

        double ratio = (double) defectNs / Math.max(fixedNs, 1);
        System.out.printf("storm-0001 [T_out=%d L=%d RUNS=%d]: defect=%.1fms fixed=%.1fms ratio=%.1fx%n",
                T_out, L, RUNS, defectNs / 1e6, fixedNs / 1e6, ratio);

        if (ratio < 1.5) {
            throw new AssertionError("Expected defect to be slower; ratio=" + ratio);
        }
        System.out.println("storm-0001 PASS");
    }

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