package support;

import java.util.*;

/**
 * MongodbPipelineAlgorithm — models two CWE-407 sites in MongoDB aggregation pipeline.
 *
 * ── mongodb-0001  document_source_streaming_group.cpp:142  (MEDIUM) ──────────
 *
 * DEFECT: During $group stage initialization, a list of id field names is built
 * from the BSON spec (idFieldNames vector). For each element in monotonicIdFields,
 * the code calls std::find(idFieldNames.begin(), idFieldNames.end(), fieldName)
 * to resolve the index position. With N id fields and M monotonic fields, this
 * is O(N × M) — quadratic when N ≈ M.
 *
 * Site: document_source_streaming_group.cpp:142
 *   for (const auto& fieldNameElem : monotonicIdFields) {
 *       auto it = std::find(idFieldNames.begin(), idFieldNames.end(), fieldName);
 *       // O(N) per iteration
 *   }
 *
 * Fix: Build an unordered_map<StringData, size_t> from idFieldNames once, then
 * resolve each monotonicIdField in O(1). Total: O(N + M) instead of O(N × M).
 *
 * ── mongodb-0002  document_source_internal_unpack_bucket.cpp:1076  (HIGH) ────
 *
 * DEFECT: During pipeline optimization, pushDownComputedMetaProjection() tracks
 * which DocumentSource* stages it has already processed to prevent infinite loops.
 * The tracker is a std::vector<DocumentSource*> (_triedComputedMetaPushDownFor).
 * Each call does std::find(...) — O(n) pointer comparison — before pushing back.
 * As the optimization loop processes more stages, the vector grows and each check
 * becomes more expensive. Total: O(S²) where S = stages processed.
 *
 * Site: document_source_internal_unpack_bucket.h:457
 *   std::vector<DocumentSource*> _triedComputedMetaPushDownFor;
 *
 * Site: document_source_internal_unpack_bucket.cpp:1076
 *   if (std::find(_triedComputedMetaPushDownFor.begin(),
 *                 _triedComputedMetaPushDownFor.end(),
 *                 nextTransform) != _triedComputedMetaPushDownFor.end())
 *       return boost::none;
 *   _triedComputedMetaPushDownFor.push_back(nextTransform);
 *
 * Fix: Replace vector with absl::flat_hash_set<DocumentSource*> (or
 * std::unordered_set<DocumentSource*>). Pointer identity is correct here —
 * each DocumentSource* is a unique stage object. No custom hash required.
 * O(1) lookup replaces O(n) scan.
 *
 * Comparison counts:
 *   0001 defective: N * M  (N idFields × M monotonicFields)
 *   0001 fixed:     N + M  (build map once, then M O(1) lookups)
 *   0002 defective: 0+1+2+...+(S-1) = S*(S-1)/2  (triangular — each new stage
 *                   scans all previously processed stages)
 *   0002 fixed:     S  (one O(1) hash lookup per stage)
 */
public class MongodbPipelineAlgorithm {

    // ── 0001: Streaming group field index resolution ──────────────────────────

    public static final class FieldIndexResult {
        public final List<Integer> indexes;    // resolved positions
        public final long comparisons;

        public FieldIndexResult(List<Integer> indexes, long comparisons) {
            this.indexes = indexes;
            this.comparisons = comparisons;
        }
    }

    /** Defective: O(N × M) — std::find scan per monotonic field. */
    public static FieldIndexResult defectiveResolveIndexes(
            List<String> idFieldNames, List<String> monotonicIdFields) {
        List<Integer> indexes = new ArrayList<>();
        long comparisons = 0;

        for (String fieldName : monotonicIdFields) {
            // std::find: scan idFieldNames from beginning
            for (int i = 0; i < idFieldNames.size(); i++) {
                comparisons++;
                if (idFieldNames.get(i).equals(fieldName)) {
                    indexes.add(i);
                    break;
                }
            }
        }

        Collections.sort(indexes);
        return new FieldIndexResult(indexes, comparisons);
    }

    /** Fixed: O(N + M) — HashMap built once, O(1) per monotonic field lookup. */
    public static FieldIndexResult fixedResolveIndexes(
            List<String> idFieldNames, List<String> monotonicIdFields) {
        List<Integer> indexes = new ArrayList<>();
        long comparisons = idFieldNames.size(); // cost of building the map

        // Build reverse map once — O(N)
        Map<String, Integer> nameToIndex = new HashMap<>();
        for (int i = 0; i < idFieldNames.size(); i++) {
            nameToIndex.put(idFieldNames.get(i), i);
        }

        // O(1) per monotonic field
        for (String fieldName : monotonicIdFields) {
            comparisons++;
            Integer idx = nameToIndex.get(fieldName);
            if (idx != null) indexes.add(idx);
        }

        Collections.sort(indexes);
        return new FieldIndexResult(indexes, comparisons);
    }

    // ── 0002: Optimization loop stage tracking ────────────────────────────────
    //
    // Models pushDownComputedMetaProjection() loop.
    // Each "stage" is an integer ID (representing a DocumentSource* pointer).
    // The optimization processes S stages; each must be checked against the
    // already-processed set before being added.

    public static final class OptimizationResult {
        public final Set<Integer> processed;  // set of processed stage IDs
        public final long comparisons;

        public OptimizationResult(Set<Integer> processed, long comparisons) {
            this.processed = processed;
            this.comparisons = comparisons;
        }
    }

    /** Defective: std::vector<DocumentSource*> — O(n) scan per stage. */
    public static OptimizationResult defectiveTrackStages(List<Integer> stages) {
        List<Integer> tried = new ArrayList<>();
        long comparisons = 0;

        for (int stage : stages) {
            // std::find: scan tried list for pointer equality
            boolean found = false;
            for (int t : tried) {
                comparisons++;
                if (t == stage) { found = true; break; }
            }
            if (!found) tried.add(stage);
        }

        return new OptimizationResult(new HashSet<>(tried), comparisons);
    }

    /** Fixed: std::unordered_set<DocumentSource*> — O(1) per stage. */
    public static OptimizationResult fixedTrackStages(List<Integer> stages) {
        Set<Integer> tried = new HashSet<>();
        long comparisons = 0;

        for (int stage : stages) {
            comparisons++;   // O(1) hash lookup
            tried.add(stage);
        }

        return new OptimizationResult(tried, comparisons);
    }

    // ── Self-test ─────────────────────────────────────────────────────────────

    public static void main(String[] args) {
        System.out.println("MongodbPipelineAlgorithm — mongodb-0001 + mongodb-0002");
        System.out.println();

        // 0001: field index resolution
        System.out.println("── mongodb-0001: streaming group field index resolution ──");
        System.out.println("Defect: std::find on idFieldNames vector — O(N×M)");
        System.out.println("Fix:    HashMap<field→index> built once — O(N+M)");
        System.out.println();
        System.out.printf("%-8s  %-8s  %-14s  %-10s  %s%n",
            "N fields", "M mono", "Defective ops", "Fixed ops", "Speedup");
        System.out.println("─".repeat(56));

        for (int n : new int[]{5, 10, 20, 50, 100, 200}) {
            List<String> idFields = new ArrayList<>();
            for (int i = 0; i < n; i++) idFields.add("field_" + i);
            // monotonic = last half of id fields (worst case: found at end)
            List<String> mono = new ArrayList<>(idFields.subList(n / 2, n));

            FieldIndexResult def = defectiveResolveIndexes(idFields, mono);
            FieldIndexResult fix = fixedResolveIndexes(idFields, mono);
            assert def.indexes.equals(fix.indexes) : "index mismatch at N=" + n;

            System.out.printf("%-8d  %-8d  %-14d  %-10d  %.1fx%n",
                n, mono.size(), def.comparisons, fix.comparisons,
                (double) def.comparisons / Math.max(1, fix.comparisons));
        }

        System.out.println();

        // 0002: optimization loop tracking
        System.out.println("── mongodb-0002: optimization stage tracker ──");
        System.out.println("Defect: std::vector<DocumentSource*> + std::find — O(S²)");
        System.out.println("Fix:    std::unordered_set<DocumentSource*>     — O(S)");
        System.out.println();
        System.out.printf("%-10s  %-14s  %-10s  %s%n",
            "S stages", "Defective ops", "Fixed ops", "Speedup");
        System.out.println("─".repeat(46));

        Random rng = new Random(42);
        for (int s : new int[]{10, 20, 50, 100, 200, 500, 1000}) {
            // Mix of unique stages + some revisits (duplicates)
            List<Integer> stages = new ArrayList<>();
            for (int i = 0; i < s; i++) stages.add(i);
            for (int i = 0; i < s / 4; i++) stages.add(rng.nextInt(s)); // revisits
            Collections.shuffle(stages, rng);

            OptimizationResult def = defectiveTrackStages(stages);
            OptimizationResult fix = fixedTrackStages(stages);
            assert def.processed.equals(fix.processed) : "processed set mismatch at S=" + s;

            System.out.printf("%-10d  %-14d  %-10d  %.1fx%n",
                stages.size(), def.comparisons, fix.comparisons,
                (double) def.comparisons / Math.max(1, fix.comparisons));
        }
    }
}