package unit;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;

/**
 * CWE-407 unit test: bitcoin-0001
 *
 * Models MiniMiner::DeleteAncestorPackage() + BuildMockTemplate().
 *
 * The real code maintains m_entries as a vector sorted by feerate. At each
 * iteration it calls std::sort (re-sorting the whole vector), picks best,
 * then calls DeleteAncestorPackage which does std::find(m_entries, anc) for
 * each ancestor. After sort the best is at index 0, but for ancestor packages
 * of size A > 1, the remaining ancestors are scattered throughout the sorted
 * vector, requiring linear scans averaging O(E/2) each.
 *
 * SLOW: std::find linear scan  — O(A × E) per DeleteAncestorPackage call
 * FAST: HashMap index lookup   — O(A) per DeleteAncestorPackage call
 *
 * To measure the defect accurately, we simulate multi-transaction ancestor
 * packages by building a graph where each transaction has a parent in the
 * mempool, so A grows with the chain depth.
 */
public class MiniMinerAncestorDeleteAlgorithm {

    static long slowOps = 0;
    static long fastOps = 0;

    static class TxEntry {
        final int txid;
        final int parentTxid; // -1 = no parent (coinbase-like root)
        int fee;

        TxEntry(int txid, int parentTxid, int fee) {
            this.txid = txid;
            this.parentTxid = parentTxid;
            this.fee = fee;
        }

        @Override public boolean equals(Object o) {
            return o instanceof TxEntry && ((TxEntry) o).txid == txid;
        }
        @Override public int hashCode() { return txid; }
        @Override public String toString() { return "Tx#" + txid; }
    }

    // -----------------------------------------------------------------------
    // Build ancestor set for a given txid (all ancestors inclusive)
    // -----------------------------------------------------------------------
    static Set<TxEntry> computeAncestors(TxEntry root, HashMap<Integer, TxEntry> txMap) {
        Set<TxEntry> ancs = new HashSet<>();
        TxEntry cur = root;
        while (cur != null) {
            ancs.add(cur);
            cur = cur.parentTxid >= 0 ? txMap.get(cur.parentTxid) : null;
        }
        return ancs;
    }

    // -----------------------------------------------------------------------
    // SLOW: linear std::find to locate each ancestor in the entries vector
    // -----------------------------------------------------------------------

    static void deleteAncestorPackageSlow(List<TxEntry> entries, Set<TxEntry> ancestors) {
        for (TxEntry anc : new ArrayList<>(ancestors)) {
            // std::find(m_entries.begin(), m_entries.end(), anc) — O(E)
            int idx = -1;
            for (int i = 0; i < entries.size(); i++) {
                slowOps++;
                if (entries.get(i).txid == anc.txid) {
                    idx = i;
                    break;
                }
            }
            if (idx >= 0) {
                entries.remove(idx);
            }
        }
    }

    static int buildMockTemplateSlow(int numTxns, int chainLen) {
        // Build chain: chain of length chainLen, repeated to fill numTxns
        List<TxEntry> entries = new ArrayList<>();
        HashMap<Integer, TxEntry> txMap = new HashMap<>();
        Random rng = new Random(42);

        for (int i = 0; i < numTxns; i++) {
            int parent = (i % chainLen == 0) ? -1 : (i - 1);
            TxEntry tx = new TxEntry(i, parent, rng.nextInt(100) + 1);
            entries.add(tx);
            txMap.put(i, tx);
        }

        // Shuffle entries to scatter them (simulating mempool arrival order ≠ feerate order)
        Collections.shuffle(entries, new Random(99));
        // Rebuild txMap based on txid (unchanged)

        int iterations = 0;
        while (!entries.isEmpty()) {
            // Pick best: highest fee entry that has no unprocessed parents
            TxEntry best = null;
            for (TxEntry e : entries) {
                if (e.parentTxid < 0 || !txMap.containsKey(e.parentTxid)) {
                    if (best == null || e.fee > best.fee) {
                        best = e;
                    }
                }
            }
            if (best == null) best = entries.get(0); // fallback

            // Compute ancestors (all in-mempool parents in the chain)
            Set<TxEntry> ancestors = computeAncestors(best, txMap);
            // Only include ancestors still in entries
            ancestors.retainAll(new HashSet<>(entries));

            deleteAncestorPackageSlow(entries, ancestors);

            // Remove from txMap
            for (TxEntry anc : ancestors) {
                txMap.remove(anc.txid);
            }
            iterations++;
        }
        return iterations;
    }

    // -----------------------------------------------------------------------
    // FAST: HashMap index for O(1) lookup
    // -----------------------------------------------------------------------

    static void deleteAncestorPackageFast(List<TxEntry> entries,
                                           HashMap<Integer, Integer> index,
                                           Set<TxEntry> ancestors) {
        List<TxEntry> toDelete = new ArrayList<>(ancestors);
        for (TxEntry anc : toDelete) {
            fastOps++;                                      // O(1) index lookup
            Integer idx = index.remove(anc.txid);
            if (idx == null) continue;

            // Swap-and-pop: O(1) removal
            int last = entries.size() - 1;
            if (idx != last) {
                TxEntry moved = entries.get(last);
                entries.set(idx, moved);
                index.put(moved.txid, idx);
            }
            entries.remove(last);
        }
    }

    static int buildMockTemplateFast(int numTxns, int chainLen) {
        List<TxEntry> entries = new ArrayList<>();
        HashMap<Integer, TxEntry> txMap = new HashMap<>();
        HashMap<Integer, Integer> index = new HashMap<>();
        Random rng = new Random(42);

        for (int i = 0; i < numTxns; i++) {
            int parent = (i % chainLen == 0) ? -1 : (i - 1);
            TxEntry tx = new TxEntry(i, parent, rng.nextInt(100) + 1);
            entries.add(tx);
            txMap.put(i, tx);
            index.put(i, i);
        }

        Collections.shuffle(entries, new Random(99));
        // Rebuild index after shuffle
        index.clear();
        for (int i = 0; i < entries.size(); i++) {
            index.put(entries.get(i).txid, i);
        }

        int iterations = 0;
        while (!entries.isEmpty()) {
            TxEntry best = null;
            for (TxEntry e : entries) {
                if (e.parentTxid < 0 || !txMap.containsKey(e.parentTxid)) {
                    if (best == null || e.fee > best.fee) {
                        best = e;
                    }
                }
            }
            if (best == null) best = entries.get(0);

            Set<TxEntry> ancestors = computeAncestors(best, txMap);
            ancestors.retainAll(new HashSet<>(entries));

            deleteAncestorPackageFast(entries, index, ancestors);

            for (TxEntry anc : ancestors) {
                txMap.remove(anc.txid);
            }
            iterations++;
        }
        return iterations;
    }

    // -----------------------------------------------------------------------
    // Tests
    // -----------------------------------------------------------------------

    static boolean runTest(int numTxns, int chainLen) {
        slowOps = 0;
        fastOps = 0;

        int slowIter = buildMockTemplateSlow(numTxns, chainLen);
        int fastIter = buildMockTemplateFast(numTxns, chainLen);

        // Both should produce same number of iterations (ancestor packages processed)
        if (slowIter != fastIter) {
            System.out.printf("FAIL N=%d chain=%d: slowIter=%d != fastIter=%d%n",
                numTxns, chainLen, slowIter, fastIter);
            return false;
        }

        // Avoid divide-by-zero
        double ratio = fastOps > 0 ? (double) slowOps / fastOps : (double) slowOps;
        System.out.printf("  N=%d chain=%d iters=%d | slowOps=%d fastOps=%d ratio=%.1fx%n",
            numTxns, chainLen, slowIter, slowOps, fastOps, ratio);

        if (ratio < 5.0) {
            System.out.printf("FAIL N=%d chain=%d: ratio %.1fx < 5x threshold%n",
                numTxns, chainLen, ratio);
            return false;
        }
        return true;
    }

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

        System.out.println("=== bitcoin-0001: MiniMiner DeleteAncestorPackage O(A×E) ===");
        System.out.println();

        // (numTxns, chainLen): chainLen>1 means multi-tx ancestor packages
        int[][] tests = {
            {50,  5},
            {100, 5},
            {200, 5},
            {500, 5},
            {200, 10},
        };

        for (int[] t : tests) {
            boolean ok = runTest(t[0], t[1]);
            if (ok) { pass++; } else { fail++; }
        }

        System.out.println();
        System.out.printf("%d/%d PASS%n", pass, pass + fail);
        if (fail > 0) {
            System.exit(1);
        }
    }
}