package unit;

import java.util.*;

/**
 * Unit test for presto-0001/0002/0003/0004: CWE-407 O(D×V) contains() defects
 * in PushDownDereferences and PayloadJoinOptimizer.
 *
 * Defects:
 *   presto-0001 (PushDownDereferences.ExtractFromJoin, line 206):
 *     joinNode.getLeft().getOutputVariables().contains(baseVariable)
 *     called inside a for-each over D dereferences. ImmutableList.contains() is O(V).
 *     Total: O(D × V).
 *
 *   presto-0002 (PushDownDereferences.PushDownDereferenceThroughJoin, line 369):
 *     Same pattern — joinNode.getLeft().getOutputVariables().contains(baseVariable)
 *     inside dereference loop. O(D × V).
 *
 *   presto-0003 (PushDownDereferences.PushDownDereferenceThroughSemiJoin, line 414):
 *     semiJoinNode.getFilteringSource().getOutputVariables().contains(baseVariable)
 *     inside dereference loop. O(D × V).
 *
 *   presto-0004 (PayloadJoinOptimizer.visitJoin, line 208):
 *     inputJoinKeys.stream().filter(key -> rightNode.getOutputVariables().contains(key))
 *     ImmutableList.contains() called per join key. O(K × V).
 *
 * Fix for all: snapshot getOutputVariables() to ImmutableSet.copyOf() before the
 * loop/stream. VariableReferenceExpression has correct equals()/hashCode().
 * ImmutableSet.contains() is O(1) average.
 *
 * This test models the defective and fixed strategies using instrumented
 * List/Set membership, counts element-level comparisons, and proves:
 *   - Defective: O(D × V) comparisons
 *   - Fixed:     O(D + V) — V to build set, D × O(1) lookups
 *   - Ratio > 10x at scale (D=100, V=100)
 */
public class PrestoDerefPushdownTest {

    // ── Minimal VariableReferenceExpression stub ─────────────────────────────

    static class VarRef {
        final String name;
        VarRef(String name) { this.name = name; }

        @Override
        public boolean equals(Object o) {
            return o instanceof VarRef && ((VarRef) o).name.equals(name);
        }

        @Override
        public int hashCode() { return name.hashCode(); }

        @Override
        public String toString() { return "Var(" + name + ")"; }
    }

    // ── Instrumented membership containers ───────────────────────────────────

    /**
     * Wraps an ArrayList<VarRef> and counts every element-level equality check
     * performed during contains(). Models the defective ImmutableList.contains().
     */
    static class InstrumentedList {
        final List<VarRef> backing;
        long comparisons;

        InstrumentedList(List<VarRef> vars) {
            this.backing = new ArrayList<>(vars);
            this.comparisons = 0;
        }

        boolean contains(VarRef target) {
            for (VarRef v : backing) {
                comparisons++;
                if (v.equals(target)) return true;
            }
            return false;
        }
    }

    /**
     * Wraps a HashSet<VarRef> and counts each contains() call as 1 operation
     * (O(1) hash lookup). Models the fixed ImmutableSet.copyOf().contains().
     */
    static class InstrumentedSet {
        final Set<VarRef> backing;
        long comparisons;

        InstrumentedSet(List<VarRef> vars) {
            this.backing = new HashSet<>(vars);
            // Building the set costs V insertions — counted separately as buildCost
            this.comparisons = 0;
        }

        boolean contains(VarRef target) {
            comparisons++;   // O(1) hash lookup counts as 1
            return backing.contains(target);
        }
    }

    // ── Test fixture builders ─────────────────────────────────────────────────

    /** Build V output variables for a plan node side. */
    static List<VarRef> outputVars(int v, String prefix) {
        List<VarRef> list = new ArrayList<>();
        for (int i = 0; i < v; i++) list.add(new VarRef(prefix + i));
        return list;
    }

    /**
     * Build D dereference base-variables. Half are drawn from leftVars (will be
     * found), half are novel (will not be found). This exercises both the
     * early-exit and the full-scan branches.
     */
    static List<VarRef> derefBases(int d, List<VarRef> leftVars) {
        List<VarRef> bases = new ArrayList<>();
        for (int i = 0; i < d; i++) {
            if (i % 2 == 0 && i / 2 < leftVars.size()) {
                bases.add(leftVars.get(i / 2));   // found in left
            } else {
                bases.add(new VarRef("right_base_" + i));  // not found in left
            }
        }
        return bases;
    }

    // ── Defective simulation: O(D × V) ───────────────────────────────────────

    static class DefectiveResult {
        final long comparisons;
        final List<VarRef> leftMatched;
        DefectiveResult(long c, List<VarRef> m) { comparisons = c; leftMatched = m; }
    }

    /**
     * Models presto-0001/0002/0003: loop over D dereference bases, call
     * list.contains() (O(V)) for each — total O(D × V).
     */
    static DefectiveResult runDefective(List<VarRef> derefBases, List<VarRef> leftOutputVars) {
        InstrumentedList leftList = new InstrumentedList(leftOutputVars);
        List<VarRef> leftMatched = new ArrayList<>();

        for (VarRef base : derefBases) {
            if (leftList.contains(base)) {
                leftMatched.add(base);
                // (else goes to rightSideDereferences — not relevant to comparison count)
            }
        }
        return new DefectiveResult(leftList.comparisons, leftMatched);
    }

    /**
     * Models presto-0004: stream over K join keys, call list.contains() (O(V))
     * for each — total O(K × V).
     */
    static DefectiveResult runDefectiveStream(List<VarRef> joinKeys, List<VarRef> rightOutputVars) {
        InstrumentedList rightList = new InstrumentedList(rightOutputVars);
        List<VarRef> rightMatched = new ArrayList<>();

        for (VarRef key : joinKeys) {
            if (rightList.contains(key)) {
                rightMatched.add(key);
            }
        }
        return new DefectiveResult(rightList.comparisons, rightMatched);
    }

    // ── Fixed simulation: O(D + V) ────────────────────────────────────────────

    static class FixedResult {
        final long buildCost;     // V hash insertions to build the set
        final long lookupCost;    // D × O(1) lookups
        final List<VarRef> leftMatched;
        FixedResult(long build, long lookup, List<VarRef> m) {
            buildCost = build;
            lookupCost = lookup;
            leftMatched = m;
        }
        long totalCost() { return buildCost + lookupCost; }
    }

    /**
     * Models the fix for presto-0001/0002/0003: snapshot to ImmutableSet before
     * the loop, then O(1) contains() per dereference.
     */
    static FixedResult runFixed(List<VarRef> derefBases, List<VarRef> leftOutputVars) {
        long buildCost = leftOutputVars.size();   // V insertions into HashSet
        InstrumentedSet leftSet = new InstrumentedSet(leftOutputVars);
        List<VarRef> leftMatched = new ArrayList<>();

        for (VarRef base : derefBases) {
            if (leftSet.contains(base)) {
                leftMatched.add(base);
            }
        }
        return new FixedResult(buildCost, leftSet.comparisons, leftMatched);
    }

    /**
     * Models the fix for presto-0004: snapshot rightNode.getOutputVariables()
     * to ImmutableSet before the stream, then O(1) contains() per key.
     */
    static FixedResult runFixedStream(List<VarRef> joinKeys, List<VarRef> rightOutputVars) {
        long buildCost = rightOutputVars.size();
        InstrumentedSet rightSet = new InstrumentedSet(rightOutputVars);
        List<VarRef> rightMatched = new ArrayList<>();

        for (VarRef key : joinKeys) {
            if (rightSet.contains(key)) {
                rightMatched.add(key);
            }
        }
        return new FixedResult(buildCost, rightSet.comparisons, rightMatched);
    }

    // ── Tests ─────────────────────────────────────────────────────────────────

    /**
     * Test 1: correctness — defective and fixed must identify the same
     * left-side dereference bases. Models presto-0001/0002/0003.
     */
    static void testCorrectnessPushdown() {
        List<VarRef> leftVars = outputVars(20, "lv");
        List<VarRef> bases = derefBases(30, leftVars);

        DefectiveResult def = runDefective(bases, leftVars);
        FixedResult     fix = runFixed(bases, leftVars);

        assert new HashSet<>(def.leftMatched).equals(new HashSet<>(fix.leftMatched))
            : "defective and fixed must classify the same dereference bases as left-side";
        System.out.println("PASS testCorrectnessPushdown");
    }

    /**
     * Test 2: correctness — defective and fixed must identify the same
     * right-join-key matches. Models presto-0004.
     */
    static void testCorrectnessPayloadJoin() {
        List<VarRef> rightVars = outputVars(20, "rv");
        // Half of join keys are in rightVars
        List<VarRef> joinKeys = new ArrayList<>();
        for (int i = 0; i < 20; i++) {
            joinKeys.add(i % 2 == 0 ? rightVars.get(i / 2) : new VarRef("lk_" + i));
        }

        DefectiveResult def = runDefectiveStream(joinKeys, rightVars);
        FixedResult     fix = runFixedStream(joinKeys, rightVars);

        assert new HashSet<>(def.leftMatched).equals(new HashSet<>(fix.leftMatched))
            : "defective and fixed must find the same right-side join keys";
        System.out.println("PASS testCorrectnessPayloadJoin");
    }

    /**
     * Test 3: defective grows quadratically — doubling D and V together should
     * increase comparisons by > 3x (approaching 4x). Models presto-0001/0002/0003.
     */
    static void testDefectiveGrowsQuadratically() {
        long prev = -1;
        for (int n : new int[]{20, 40, 80}) {
            List<VarRef> leftVars = outputVars(n, "lv");
            List<VarRef> bases    = derefBases(n, leftVars);
            long c = runDefective(bases, leftVars).comparisons;
            if (prev > 0) {
                double ratio = (double) c / prev;
                assert ratio > 3.0
                    : "defective comparisons should grow >3x when D and V double; got " + ratio + " at n=" + n;
            }
            prev = c;
        }
        System.out.println("PASS testDefectiveGrowsQuadratically");
    }

    /**
     * Test 4: ratio at scale — at D=V=100, the fixed strategy (including set
     * build cost) must be >10x fewer operations than the defective strategy.
     * Covers all four defects.
     */
    static void testRatioAtScaleExceedsTenX() {
        int d = 100, v = 100;

        // presto-0001/0002/0003 shape
        List<VarRef> leftVars  = outputVars(v, "lv");
        List<VarRef> bases     = derefBases(d, leftVars);
        long defComparisons    = runDefective(bases, leftVars).comparisons;
        FixedResult fixResult  = runFixed(bases, leftVars);
        long fixTotal          = fixResult.totalCost();   // build + lookup

        double ratio = (double) defComparisons / fixTotal;
        assert ratio > 10.0
            : "at D=V=100, defective should be >10x worse than fixed (incl. build cost); ratio=" + ratio;
        System.out.printf(
            "PASS testRatioAtScaleExceedsTenX  [presto-0001/0002/0003] "
            + "defective=%d, fixed_total=%d (build=%d + lookup=%d), ratio=%.1fx%n",
            defComparisons, fixTotal, fixResult.buildCost, fixResult.lookupCost, ratio);

        // presto-0004 shape (K join keys, V right output vars)
        int k = 100;
        List<VarRef> rightVars  = outputVars(v, "rv");
        List<VarRef> joinKeys   = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            joinKeys.add(i % 2 == 0 && i / 2 < rightVars.size()
                ? rightVars.get(i / 2) : new VarRef("lk_" + i));
        }
        long defStream          = runDefectiveStream(joinKeys, rightVars).comparisons;
        FixedResult fixStream   = runFixedStream(joinKeys, rightVars);
        long fixStreamTotal     = fixStream.totalCost();

        double ratio4 = (double) defStream / fixStreamTotal;
        assert ratio4 > 10.0
            : "at K=V=100, defective presto-0004 should be >10x worse; ratio=" + ratio4;
        System.out.printf(
            "PASS testRatioAtScaleExceedsTenX  [presto-0004] "
            + "defective=%d, fixed_total=%d (build=%d + lookup=%d), ratio=%.1fx%n",
            defStream, fixStreamTotal, fixStream.buildCost, fixStream.lookupCost, ratio4);
    }

    // ── Entry point ───────────────────────────────────────────────────────────

    public static void main(String[] args) {
        testCorrectnessPushdown();
        testCorrectnessPayloadJoin();
        testDefectiveGrowsQuadratically();
        testRatioAtScaleExceedsTenX();
        System.out.println("All presto-0001/0002/0003/0004 tests passed.");
    }
}