package unit;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * CWE-407 unit test: opensearch-004
 * TransportSegmentReplicationStatsAction uses a List<Integer> (shardsToFetch) and calls
 * .contains() on it inside a for loop over all shard responses.
 *
 * Defect:  O(S × F) — List<Integer>.contains() is O(F) per response (S = responses, F = fetched shard IDs)
 * Fix:     O(S + F) — HashSet<Integer>.contains() is O(1)
 *
 * No JUnit. Run with: javac SegmentReplicationShardsToFetchContains.java && java -cp . unit.SegmentReplicationShardsToFetchContains
 */
public class SegmentReplicationShardsToFetchContains {

    // Minimal stand-in for a shard response with an integer shard ID
    static class FakeShardResponse {
        final int shardId;
        FakeShardResponse(int shardId) { this.shardId = shardId; }
    }

    // ---- SLOW: mirrors the defective action logic ----
    static List<Integer> filterResponsesSlow(List<FakeShardResponse> responses, List<Integer> shardsToFetch) {
        List<Integer> matched = new ArrayList<>();
        for (FakeShardResponse r : responses) {
            if (shardsToFetch.isEmpty() || shardsToFetch.contains(r.shardId)) {  // O(F) per response
                matched.add(r.shardId);
            }
        }
        return matched;
    }

    // ---- FAST: use a HashSet for O(1) lookup ----
    static List<Integer> filterResponsesFast(List<FakeShardResponse> responses, Set<Integer> shardsToFetch) {
        List<Integer> matched = new ArrayList<>();
        for (FakeShardResponse r : responses) {
            if (shardsToFetch.isEmpty() || shardsToFetch.contains(r.shardId)) {  // O(1) per response
                matched.add(r.shardId);
            }
        }
        return matched;
    }

    // ---- count total .contains() probes ----
    static long countSlowProbes(int numResponses, int numShardsToFetch) {
        return (long) numResponses * numShardsToFetch;  // worst case: never found early
    }

    static long countFastProbes(int numResponses) {
        return numResponses;  // one O(1) lookup per response
    }

    public static void main(String[] args) {
        System.out.println("=== opensearch-004: TransportSegmentReplicationStatsAction CWE-407 ===\n");

        // --- Correctness check ---
        // 20 shard responses; we want shards 5, 10, 15
        List<FakeShardResponse> responses = new ArrayList<>();
        for (int i = 0; i < 20; i++) responses.add(new FakeShardResponse(i));

        List<Integer> fetchList = Arrays.asList(5, 10, 15);
        Set<Integer> fetchSet = new HashSet<>(fetchList);

        List<Integer> slowResult = filterResponsesSlow(responses, fetchList);
        List<Integer> fastResult = filterResponsesFast(responses, fetchSet);

        if (!slowResult.equals(fastResult)) {
            System.out.println("FAIL correctness: slow=" + slowResult + " fast=" + fastResult);
            System.exit(1);
        }
        if (slowResult.size() != 3) {
            System.out.println("FAIL: expected 3 matches, got " + slowResult.size() + " => " + slowResult);
            System.exit(1);
        }
        System.out.println("correctness OK  matched=" + slowResult);

        // --- Empty shardsToFetch (fetch all) should also work ----
        List<Integer> allSlow = filterResponsesSlow(responses, new ArrayList<>());
        List<Integer> allFast = filterResponsesFast(responses, new HashSet<>());
        if (!allSlow.equals(allFast) || allSlow.size() != 20) {
            System.out.println("FAIL: empty-fetch mismatch or wrong count");
            System.exit(1);
        }
        System.out.println("empty-fetch (all shards) OK");

        // --- Op-count comparison at scale ---
        System.out.println("\n=== Op-count (S=responses, F=shardsToFetch) ===");
        System.out.printf("%-8s  %-8s  %-16s  %-14s  %-10s%n", "S", "F", "slow_probes", "fast_probes", "ratio");
        System.out.println("-".repeat(62));

        int[][] cases = {{100, 10}, {500, 50}, {1000, 50}, {2000, 100}, {5000, 200}};
        for (int[] c : cases) {
            int S = c[0], F = c[1];
            long slowOps = countSlowProbes(S, F);
            long fastOps = countFastProbes(S);
            double ratio = (double) slowOps / fastOps;
            System.out.printf("%-8d  %-8d  %-16d  %-14d  %-10.1f%n", S, F, slowOps, fastOps, ratio);

            if (slowOps <= fastOps) {
                System.out.println("FAIL: slow was not worse than fast at S=" + S + " F=" + F);
                System.exit(1);
            }
        }

        // --- Verify S=1000, F=50: should be exactly 50× worse ----
        long slowOps = countSlowProbes(1000, 50);
        long fastOps = countFastProbes(1000);
        double ratio = (double) slowOps / fastOps;
        if (Math.abs(ratio - 50.0) > 0.01) {
            System.out.printf("FAIL: expected ratio 50x at S=1000 F=50, got %.2fx%n", ratio);
            System.exit(1);
        }
        System.out.printf("%nspeedup at S=1000, F=50: %.0fx  PASS%n", ratio);

        System.out.println("\nALL PASS");
    }
}