package unit;

import java.util.*;
import java.util.stream.*;

/**
 * Standalone unit test for doris-0001:
 * BindExpression.processNonStandardAggregate — List.contains() per projection → O(P×G).
 *
 * Simulates processNonStandardAggregate():
 *   for each projection, check if it's in groupingExprs (passed as List).
 *
 * Compile: javac -d . BindExprGroupingAlgorithm.java
 * Run:     java unit.BindExprGroupingAlgorithm
 */
public class BindExprGroupingAlgorithm {

    // ── Expression stubs ─────────────────────────────────────────────────────

    static class Expression {
        final String id;
        Expression(String id) { this.id = id; }
        @Override public boolean equals(Object o) {
            return o instanceof Expression && id.equals(((Expression) o).id);
        }
        @Override public int hashCode() { return id.hashCode(); }
        @Override public String toString() { return id; }
    }

    static class SlotReference extends Expression {
        SlotReference(String id) { super(id); }
    }

    // ── Result ───────────────────────────────────────────────────────────────

    static class Result {
        final List<String> output;  // projection labels: "alias:X" or "keep:X"
        final long ns;
        Result(List<String> output, long ns) { this.output = output; this.ns = ns; }
    }

    // ── Defective: Collection<Expression> as List → O(G) .contains per proj ─

    static class DefectiveAlgorithm {
        Result processNonStandardAggregate(
                List<Expression> originalProjections,
                Collection<Expression> groupingExprs) {  // List passed → O(G) .contains
            long t0 = System.nanoTime();
            List<String> out = new ArrayList<>();
            for (Expression projection : originalProjections) {
                if (projection instanceof SlotReference
                        && !groupingExprs.contains(projection)) {  // CWE-407 site
                    out.add("alias:" + projection.id);
                } else {
                    out.add("keep:" + projection.id);
                }
            }
            return new Result(out, System.nanoTime() - t0);
        }
    }

    // ── Fixed: convert to Set once → O(1) .contains ──────────────────────────

    static class FixedAlgorithm {
        Result processNonStandardAggregate(
                List<Expression> originalProjections,
                Collection<Expression> groupingExprs) {
            long t0 = System.nanoTime();
            // Ensure O(1) membership
            Set<Expression> groupingSet = (groupingExprs instanceof Set)
                    ? (Set<Expression>) groupingExprs
                    : new HashSet<>(groupingExprs);  // O(G) once
            List<String> out = new ArrayList<>();
            for (Expression projection : originalProjections) {
                if (projection instanceof SlotReference
                        && !groupingSet.contains(projection)) {  // O(1)
                    out.add("alias:" + projection.id);
                } else {
                    out.add("keep:" + projection.id);
                }
            }
            return new Result(out, System.nanoTime() - t0);
        }
    }

    // ── Helpers ──────────────────────────────────────────────────────────────

    static List<Expression> makeProjections(int P, int G) {
        // First G projections are SlotReferences that ARE in groupingExprs → "keep"
        // Next (P-G) are SlotReferences NOT in groupingExprs → "alias"
        List<Expression> list = new ArrayList<>(P);
        for (int i = 0; i < P; i++) {
            list.add(new SlotReference("s" + i));
        }
        return list;
    }

    static List<Expression> makeGroupingList(int G) {
        List<Expression> list = new ArrayList<>(G);
        for (int i = 0; i < G; i++) list.add(new SlotReference("s" + i));
        return list;
    }

    // ── Assertions ───────────────────────────────────────────────────────────

    static void assertEquals(Object expected, Object actual, String msg) {
        if (!expected.equals(actual))
            throw new AssertionError(msg + ": expected=" + expected + " actual=" + actual);
    }

    static void assertTrue(boolean cond, String msg) {
        if (!cond) throw new AssertionError(msg);
    }

    public static void main(String[] args) {
        System.out.println("=== doris-0001: BindExprGroupingAlgorithm ===");

        DefectiveAlgorithm defAlg = new DefectiveAlgorithm();
        FixedAlgorithm fixAlg = new FixedAlgorithm();

        // Correctness: small case
        {
            int P = 10, G = 4;
            List<Expression> projections = makeProjections(P, G);
            List<Expression> grouping = makeGroupingList(G);

            Result dr = defAlg.processNonStandardAggregate(projections, grouping);
            Result fr = fixAlg.processNonStandardAggregate(projections, grouping);

            assertEquals(dr.output, fr.output, "output labels");
            // First G should be "keep:", rest "alias:"
            for (int i = 0; i < G; i++) assertTrue(dr.output.get(i).startsWith("keep:"), "keep at " + i);
            for (int i = G; i < P; i++) assertTrue(dr.output.get(i).startsWith("alias:"), "alias at " + i);
            System.out.println("PASS  correctness (P=10, G=4)");
        }

        // Correctness: Set passed directly (should not double-wrap)
        {
            int P = 8, G = 3;
            List<Expression> projections = makeProjections(P, G);
            Set<Expression> groupingSet = new HashSet<>(makeGroupingList(G));

            Result dr = defAlg.processNonStandardAggregate(projections, new ArrayList<>(groupingSet));
            Result fr = fixAlg.processNonStandardAggregate(projections, groupingSet);
            assertEquals(dr.output, fr.output, "Set input output labels");
            System.out.println("PASS  correctness Set input (P=8, G=3)");
        }

        // Benchmark at P=400, G=200
        {
            int P = 400, G = 200;
            List<Expression> projections = makeProjections(P, G);
            List<Expression> grouping = makeGroupingList(G);

            // warm up
            for (int i = 0; i < 5; i++) {
                defAlg.processNonStandardAggregate(projections, new ArrayList<>(grouping));
                fixAlg.processNonStandardAggregate(projections, new ArrayList<>(grouping));
            }

            long defNs = 0, fixNs = 0;
            int reps = 50;
            for (int i = 0; i < reps; i++) {
                defNs += defAlg.processNonStandardAggregate(projections, new ArrayList<>(grouping)).ns;
                fixNs += fixAlg.processNonStandardAggregate(projections, new ArrayList<>(grouping)).ns;
            }
            defNs /= reps; fixNs /= reps;
            double ratio = (double) defNs / Math.max(1, fixNs);

            System.out.printf("BENCH P=%d G=%d reps=%d%n", P, G, reps);
            System.out.printf("  defective avg: %,d ns%n", defNs);
            System.out.printf("  fixed     avg: %,d ns%n", fixNs);
            System.out.printf("  speedup:       %.1fx%n", ratio);

            assertTrue(ratio >= 2.0, "Expected fixed >= 2x faster, got " + ratio + "x");
            System.out.println("PASS  speedup >= 2x");
        }

        System.out.println("=== ALL PASS ===");
    }
}