package unit;

import java.util.Arrays;

/**
 * Ffmpeg0003MpegtsProgramDiscardTest — CWE-407 ffmpeg-0003
 *
 * Models discard_pid() in libavformat/mpegts.c:
 *   slow() = O(P × J × K) inner linear program-id scan (current defect)
 *   fast() = O(P × J × log K) using sorted arrays + bsearch (patch)
 *
 * Parameters: K programs, P internal program entries, J pids per entry.
 * Assert: slowOps > fastOps * Nx at K=50 programs, P=50, J=10.
 */
public class Ffmpeg0003MpegtsProgramDiscardTest {

    static long slowOps;
    static long fastOps;

    /**
     * Slow: O(P × J × K) — models original discard_pid inner loop.
     * For each (program, pid) match, linearly scans all K AVPrograms by id.
     */
    static int discardPidSlow(int[] prg_ids, int[] prg_pids, boolean[] prg_disc_flag,
                               int[] avprg_ids, boolean[] avprg_discard,
                               int target_pid, int nb_prg, int pids_per_prg, int nb_avprg)
    {
        // Early exit: any discarded?
        boolean any_disc = false;
        for (int k = 0; k < nb_avprg; k++) {
            slowOps++;
            if (avprg_discard[k]) { any_disc = true; break; }
        }
        if (!any_disc) return 0;

        int used = 0, discarded = 0;
        for (int i = 0; i < nb_prg; i++) {
            for (int j = 0; j < pids_per_prg; j++) {
                int idx = i * pids_per_prg + j;
                if (prg_pids[idx] != target_pid) continue;
                // Inner O(K) scan — the defect
                for (int k = 0; k < nb_avprg; k++) {
                    slowOps++;
                    if (avprg_ids[k] == prg_ids[i]) {
                        if (avprg_discard[k]) discarded++;
                        else                  used++;
                    }
                }
            }
        }
        return (used == 0 && discarded > 0) ? 1 : 0;
    }

    /**
     * Fast: O(K log K + P × J × log K) — models bsearch fix.
     * Pre-sorts discarded and used id arrays, then uses binary search.
     */
    static int discardPidFast(int[] prg_ids, int[] prg_pids, boolean[] prg_disc_flag,
                               int[] avprg_ids, boolean[] avprg_discard,
                               int target_pid, int nb_prg, int pids_per_prg, int nb_avprg)
    {
        // Build sorted disc_ids and used_ids in O(K)
        int[] disc_ids = new int[nb_avprg];
        int[] used_ids = new int[nb_avprg];
        int nb_disc = 0, nb_used = 0;
        for (int k = 0; k < nb_avprg; k++) {
            fastOps++;
            if (avprg_discard[k]) disc_ids[nb_disc++] = avprg_ids[k];
            else                  used_ids[nb_used++] = avprg_ids[k];
        }
        if (nb_disc == 0) return 0;  // no programs discarded

        // Sort O(K log K)
        int[] disc_sorted = Arrays.copyOf(disc_ids, nb_disc);
        int[] used_sorted = Arrays.copyOf(used_ids, nb_used);
        Arrays.sort(disc_sorted);
        Arrays.sort(used_sorted);

        int used = 0, discarded = 0;
        for (int i = 0; i < nb_prg; i++) {
            for (int j = 0; j < pids_per_prg; j++) {
                int idx = i * pids_per_prg + j;
                if (prg_pids[idx] != target_pid) continue;
                // O(log K) binary search
                fastOps++;
                if (nb_disc > 0 && Arrays.binarySearch(disc_sorted, 0, nb_disc, prg_ids[i]) >= 0)
                    discarded++;
                if (nb_used > 0 && Arrays.binarySearch(used_sorted, 0, nb_used, prg_ids[i]) >= 0)
                    used++;
            }
        }
        return (used == 0 && discarded > 0) ? 1 : 0;
    }

    public static void main(String[] args) {
        final int NX = 5;

        // Simulate: K=50 AVPrograms, some discarded
        final int K = 50;      // nb_programs (AVProgram count)
        final int P = 50;      // nb_prg (internal Program count)
        final int J = 10;      // pids_per_program
        final int TARGET_PID = 42;

        int[] avprg_ids     = new int[K];
        boolean[] avprg_disc = new boolean[K];
        int[] prg_ids       = new int[P];
        int[] prg_pids      = new int[P * J];

        // Set up: AVPrograms 0..K-1, mark last 10 as discarded
        for (int k = 0; k < K; k++) {
            avprg_ids[k]  = k + 1000;  // IDs 1000..1049
            avprg_disc[k] = (k >= K - 10);  // last 10 discarded
        }

        // Internal programs map to AVProgram IDs
        for (int i = 0; i < P; i++) {
            prg_ids[i] = i + 1000;  // matches avprg_ids
            for (int j = 0; j < J; j++) {
                prg_pids[i * J + j] = (i == 5 && j == 2) ? TARGET_PID : (i * J + j + 1);
            }
        }
        boolean[] dummy_flag = new boolean[P];

        // Warm up
        slowOps = 0; fastOps = 0;
        discardPidSlow(prg_ids, prg_pids, dummy_flag, avprg_ids, avprg_disc,
                       TARGET_PID, P, J, K);
        discardPidFast(prg_ids, prg_pids, dummy_flag, avprg_ids, avprg_disc,
                       TARGET_PID, P, J, K);

        // Measure
        slowOps = 0; fastOps = 0;
        final int CALLS = 200;
        int slowResult = 0, fastResult = 0;
        for (int c = 0; c < CALLS; c++) {
            slowResult = discardPidSlow(prg_ids, prg_pids, dummy_flag, avprg_ids, avprg_disc,
                                        TARGET_PID, P, J, K);
            fastResult = discardPidFast(prg_ids, prg_pids, dummy_flag, avprg_ids, avprg_disc,
                                        TARGET_PID, P, J, K);
        }

        boolean correctnessOk = (slowResult == fastResult);
        boolean speedupOk = slowOps > fastOps * NX;

        System.out.printf("K=%d programs, P=%d internal prg, J=%d pids, CALLS=%d%n", K, P, J, CALLS);
        System.out.printf("slow (linear scan) ops: %d%n", slowOps);
        System.out.printf("fast (bsearch)     ops: %d%n", fastOps);
        System.out.printf("speedup ratio: %.1fx (required >%dx)%n",
            (double) slowOps / fastOps, NX);
        System.out.printf("result: slow=%d fast=%d%n", slowResult, fastResult);

        int passed = 0, total = 2;
        if (correctnessOk) {
            System.out.println("1/2 PASS correctness: both return same discard decision");
            passed++;
        } else {
            System.out.printf("1/2 FAIL correctness: slow=%d fast=%d%n", slowResult, fastResult);
        }
        if (speedupOk) {
            System.out.printf("2/2 PASS speedup: %d > %d * %d%n", slowOps, fastOps, NX);
            passed++;
        } else {
            System.out.printf("2/2 FAIL speedup: %d not > %d * %d%n", slowOps, fastOps, NX);
        }

        System.out.printf("%d/%d PASS%n", passed, total);
        if (passed < total) System.exit(1);
    }
}