import java.util.*;

/**
 * Unit test for pinot-0002: PartialUpsertHandler + ColumnarMerger List.contains per column
 *
 * Simulates the O(C×P) defective pattern vs O(C) fixed pattern.
 */
public class TestPartialUpsertPrimaryKeyContains {

    // --- Defective: List.contains per column ---
    static int defectiveMerge(List<String> allColumns,
                               List<String> primaryKeyColumns,
                               List<String> comparisonColumns) {
        int cost = 0;
        for (String column : allColumns) {
            cost += primaryKeyColumns.size();   // cost of List.contains
            boolean isPK = primaryKeyColumns.contains(column);
            cost += comparisonColumns.size();   // cost of List.contains
            boolean isCmp = comparisonColumns.contains(column);
            if (isPK || isCmp) continue;
            // process column (no-op in test)
        }
        return cost;
    }

    // --- Fixed: Set.contains per column ---
    static int fixedMerge(List<String> allColumns,
                           Set<String> primaryKeySet,
                           Set<String> comparisonSet) {
        int cost = 0;
        for (String column : allColumns) {
            cost += 1;   // O(1) HashSet.contains
            boolean isPK = primaryKeySet.contains(column);
            cost += 1;   // O(1) HashSet.contains
            boolean isCmp = comparisonSet.contains(column);
            if (isPK || isCmp) continue;
            // process column
        }
        return cost;
    }

    static List<String> buildColumns(int count, String prefix) {
        List<String> cols = new ArrayList<>();
        for (int i = 0; i < count; i++) cols.add(prefix + i);
        return cols;
    }

    public static void main(String[] args) {
        int pass = 0, fail = 0;

        // Test 1: basic correctness - same columns skipped
        {
            List<String> pk = Arrays.asList("id", "ts");
            List<String> cmp = Arrays.asList("version");
            List<String> all = new ArrayList<>(pk);
            all.addAll(cmp);
            all.addAll(Arrays.asList("col1", "col2", "col3", "col4", "col5"));

            Set<String> pkSet = new HashSet<>(pk);
            Set<String> cmpSet = new HashSet<>(cmp);

            // Count how many non-key columns processed in each version
            int defSkipped = 0, fixSkipped = 0;
            for (String col : all) {
                if (pk.contains(col) || cmp.contains(col)) defSkipped++;
                if (pkSet.contains(col) || cmpSet.contains(col)) fixSkipped++;
            }
            boolean ok = (defSkipped == fixSkipped);
            System.out.printf("[%s] Correctness: defective skipped=%d, fixed skipped=%d%n",
                ok ? "PASS" : "FAIL", defSkipped, fixSkipped);
            if (ok) pass++; else fail++;
        }

        // Test 2: cost ratio at C=50, P=3, K=2
        {
            int C = 50, P = 3, K = 2;
            List<String> pk = buildColumns(P, "pk");
            List<String> cmp = buildColumns(K, "cmp");
            List<String> all = buildColumns(C, "col");
            all.addAll(pk);
            all.addAll(cmp);

            Set<String> pkSet = new HashSet<>(pk);
            Set<String> cmpSet = new HashSet<>(cmp);

            int defCost = defectiveMerge(all, pk, cmp);
            int fixCost = fixedMerge(all, pkSet, cmpSet);
            double ratio = (double) defCost / fixCost;
            boolean ok = ratio >= (P + K) * 0.8;  // ratio should approach (P+K)
            System.out.printf("[%s] Cost ratio C=%d P=%d K=%d: defective=%d, fixed=%d, ratio=%.1fx%n",
                ok ? "PASS" : "FAIL", C, P, K, defCost, fixCost, ratio);
            if (ok) pass++; else fail++;
        }

        // Test 3: cost ratio at C=100, P=5, K=3
        {
            int C = 100, P = 5, K = 3;
            List<String> pk = buildColumns(P, "pk");
            List<String> cmp = buildColumns(K, "cmp");
            List<String> all = buildColumns(C, "col");
            all.addAll(pk);
            all.addAll(cmp);

            Set<String> pkSet = new HashSet<>(pk);
            Set<String> cmpSet = new HashSet<>(cmp);

            int defCost = defectiveMerge(all, pk, cmp);
            int fixCost = fixedMerge(all, pkSet, cmpSet);
            double ratio = (double) defCost / fixCost;
            boolean ok = ratio >= (P + K) * 0.8;
            System.out.printf("[%s] Cost ratio C=%d P=%d K=%d: defective=%d, fixed=%d, ratio=%.1fx%n",
                ok ? "PASS" : "FAIL", C, P, K, defCost, fixCost, ratio);
            if (ok) pass++; else fail++;
        }

        // Test 4: throughput model - at 100k rows/sec, cost reduction
        {
            int C = 100, P = 3, K = 2;
            long rowsPerSec = 100_000L;
            List<String> pk = buildColumns(P, "pk");
            List<String> cmp = buildColumns(K, "cmp");
            List<String> all = buildColumns(C, "col");

            Set<String> pkSet = new HashSet<>(pk);
            Set<String> cmpSet = new HashSet<>(cmp);

            int defCostPerRow = defectiveMerge(all, pk, cmp);
            int fixCostPerRow = fixedMerge(all, pkSet, cmpSet);
            long defTotal = defCostPerRow * rowsPerSec;
            long fixTotal = fixCostPerRow * rowsPerSec;
            double ratio = (double) defTotal / fixTotal;
            boolean ok = ratio >= 3.0;
            System.out.printf("[%s] Throughput C=%d P=%d K=%d: defective=%d ops/row, fixed=%d ops/row, ratio=%.1fx%n",
                ok ? "PASS" : "FAIL", C, P, K, defCostPerRow, fixCostPerRow, ratio);
            if (ok) pass++; else fail++;
        }

        // Test 5: edge case - no primary keys (degenerate)
        {
            List<String> pk = new ArrayList<>();
            List<String> cmp = new ArrayList<>();
            List<String> all = buildColumns(10, "col");
            Set<String> pkSet = new HashSet<>(pk);
            Set<String> cmpSet = new HashSet<>(cmp);

            int defCost = defectiveMerge(all, pk, cmp);
            int fixCost = fixedMerge(all, pkSet, cmpSet);
            // With empty lists, defective cost = 0 (P+K=0), fixed cost = 2*C (two O(1) calls each)
            // But logically same — just verify no crash
            boolean ok = true;
            System.out.printf("[%s] Edge empty PK/CMP: defective=%d, fixed=%d%n",
                ok ? "PASS" : "FAIL", defCost, fixCost);
            if (ok) pass++; else fail++;
        }

        System.out.printf("%nResult: %d PASS, %d FAIL%n", pass, fail);
        System.exit(fail > 0 ? 1 : 0);
    }
}