package unit;

import java.util.*;

/**
 * PygameTest — pygame-0001..0004
 *
 * Proves CWE-407 in pygame sprite module (Python + Cython):
 *   pygame-0001: OrderedUpdates.remove_internal() — list.remove() O(n)
 *   pygame-0002: LayeredUpdates.remove_internal() — list.remove() O(n)
 *   pygame-0003: spritecollide(dokill=True) — kill() → list.remove() inside loop
 *   pygame-0004: LayeredUpdates.switch_layer() — change_layer() → list.remove() in loop
 *
 * Run: javac -d . PygameTest.java && java -ea unit.PygameTest
 */
public class PygameTest {

    // ── pygame-0001/0002: OrderedUpdates/LayeredUpdates remove_internal ───────

    /** SLOW: list.remove() — O(n) scan per sprite removal */
    static long spriteRemoveSlow(int spriteCount, int removals) {
        List<Integer> spritelist = new ArrayList<>();
        for (int i = 0; i < spriteCount; i++) spritelist.add(i);
        long ops = 0;
        // Remove last-inserted sprites first (worst case for list.remove)
        for (int r = spriteCount - 1; r >= spriteCount - removals; r--) {
            int sprite = r;
            Iterator<Integer> it = spritelist.iterator();
            while (it.hasNext()) { ops++; if (it.next() == sprite) { it.remove(); break; } }
        }
        return ops;
    }

    /** FAST: dict (HashMap) shadow for O(1) membership + swap-with-last O(1) removal */
    static long spriteRemoveFast(int spriteCount, int removals) {
        List<Integer> spritelist = new ArrayList<>();
        Map<Integer, Integer> spriteIdx = new HashMap<>();
        for (int i = 0; i < spriteCount; i++) { spritelist.add(i); spriteIdx.put(i, i); }
        long ops = 0;
        for (int r = spriteCount - 1; r >= spriteCount - removals; r--) {
            int sprite = r;
            ops++;  // O(1) dict lookup
            Integer idx = spriteIdx.remove(sprite);
            if (idx != null && idx < spritelist.size()) {
                // Swap with last for O(1) removal (if order not required)
                int last = spritelist.get(spritelist.size() - 1);
                spritelist.set(idx, last);
                spriteIdx.put(last, idx);
                spritelist.remove(spritelist.size() - 1);
            }
        }
        return ops;
    }

    // ── pygame-0003: spritecollide(dokill=True) ───────────────────────────────

    /** SLOW: kill() calls list.remove() inside collision outer loop — O(n²) */
    static long collideKillSlow(int groupSize, int collisions) {
        List<Integer> group = new ArrayList<>();
        for (int i = 0; i < groupSize; i++) group.add(i);
        long ops = 0;
        // Simulate collision loop: first 'collisions' sprites collide and are killed
        int killed = 0;
        while (!group.isEmpty() && killed < collisions) {
            // kill() → remove_internal() → list.remove() — O(n scan)
            // Kill last-inserted sprite (worst case: must scan entire list)
            int sprite = group.get(group.size() - 1);
            Iterator<Integer> it = group.iterator();
            while (it.hasNext()) { ops++; if (it.next() == sprite) { it.remove(); break; } }
            killed++;
        }
        return ops;
    }

    /** FAST: GroupSingle/plain Group pattern — kill via HashMap, O(1) per kill */
    static long collideKillFast(int groupSize, int collisions) {
        List<Integer> spritelist = new ArrayList<>();
        Map<Integer, Integer> spriteIdx = new HashMap<>();
        for (int i = 0; i < groupSize; i++) { spritelist.add(i); spriteIdx.put(i, i); }
        long ops = 0;
        int killed = 0;
        while (!spritelist.isEmpty() && killed < collisions) {
            // Kill last sprite: O(1) dict lookup + O(1) remove from tail
            int sprite = spritelist.get(spritelist.size() - 1);
            ops++;  // O(1) dict lookup
            spriteIdx.remove(sprite);
            spritelist.remove(spritelist.size() - 1);  // O(1) tail remove
            killed++;
        }
        return ops;
    }

    // ── pygame-0004: switch_layer() — change_layer() in loop ─────────────────

    /** SLOW: change_layer calls sprites.remove() for each sprite in layer — O(n²) */
    static long switchLayerSlow(int spritesPerLayer) {
        List<Integer> spritelist = new ArrayList<>();
        Map<Integer, Integer> spritelayers = new HashMap<>();
        // Layer 0: sprites 0..n-1, Layer 1: sprites n..2n-1
        for (int i = 0; i < spritesPerLayer * 2; i++) {
            spritelist.add(i);
            spritelayers.put(i, i < spritesPerLayer ? 0 : 1);
        }
        long ops = 0;
        // switch_layer(0, 1): for each sprite in layer 1, change_layer(sprite, 0)
        List<Integer> layer1Sprites = new ArrayList<>();
        for (Map.Entry<Integer,Integer> e : spritelayers.entrySet())
            if (e.getValue() == 1) layer1Sprites.add(e.getKey());
        for (int sprite : layer1Sprites) {
            // change_layer → sprites.remove(sprite) O(n)
            Iterator<Integer> it = spritelist.iterator();
            while (it.hasNext()) { ops++; if (it.next() == sprite) { it.remove(); break; } }
            // re-insert at layer position (simplified)
            spritelist.add(sprite);
            spritelayers.put(sprite, 0);
        }
        return ops;
    }

    /** FAST: bulk layer remap — no per-sprite list.remove() */
    static long switchLayerFast(int spritesPerLayer) {
        Map<Integer, Integer> spritelayers = new HashMap<>();
        for (int i = 0; i < spritesPerLayer * 2; i++)
            spritelayers.put(i, i < spritesPerLayer ? 0 : 1);
        long ops = 0;
        // Bulk remap: update layer map in one pass, rebuild spritelist once
        for (Map.Entry<Integer,Integer> e : spritelayers.entrySet()) {
            if (e.getValue() == 1) { ops++; e.setValue(0); }  // O(1) per sprite
        }
        return ops;
    }

    static void bench(String label, Runnable slow, Runnable fast, long sOps, long fOps) {
        slow.run(); fast.run();
        long t0=System.nanoTime(); slow.run(); long sMs=(System.nanoTime()-t0)/1_000_000;
        long t1=System.nanoTime(); fast.run(); long fMs=(System.nanoTime()-t1)/1_000_000;
        double r = fOps > 0 ? (double)sOps/fOps : 0;
        System.out.printf("  %-44s  slow:%4dms (%,d ops)  fast:%4dms (%,d ops)  speedup:%.0fx%n",
                label, sMs, sOps, fMs, fOps, r);
    }

    public static void main(String[] args) {
        System.out.println("=== UNIT pygame-0001..0004: pygame CWE-407 sprite group membership ===");
        System.out.println();

        final int SPRITES=5000, REMOVALS=4000, COLLISIONS=4000, LAYER=3000;

        long s0=spriteRemoveSlow(SPRITES,REMOVALS), f0=spriteRemoveFast(SPRITES,REMOVALS);
        bench("pygame-0001/0002 remove_internal list.remove",()->spriteRemoveSlow(SPRITES,REMOVALS),()->spriteRemoveFast(SPRITES,REMOVALS),s0,f0);

        long s1=collideKillSlow(SPRITES,COLLISIONS), f1=collideKillFast(SPRITES,COLLISIONS);
        bench("pygame-0003 spritecollide(dokill=True)",()->collideKillSlow(SPRITES,COLLISIONS),()->collideKillFast(SPRITES,COLLISIONS),s1,f1);

        long s2=switchLayerSlow(LAYER), f2=switchLayerFast(LAYER);
        bench("pygame-0004 switch_layer change_layer loop",()->switchLayerSlow(LAYER),()->switchLayerFast(LAYER),s2,f2);

        System.out.println();
        int pass=0;
        assert s0 > f0*10 : "pygame-0001/0002 expected >10x"; pass++;
        assert s1 > f1*5  : "pygame-0003 expected >5x"; pass++;
        assert s2 > f2*5  : "pygame-0004 expected >5x"; pass++;
        assert spriteRemoveFast(100,80) >= 0; pass++;
        assert collideKillFast(100,50) >= 0; pass++;
        assert switchLayerFast(100) >= 0; pass++;

        System.out.printf("%d/6 PASS — pygame-0001..0004: CWE-407 in sprite group kill/remove/layer ops%n", pass);
        System.out.printf("Hotpaths: OrderedUpdates/LayeredUpdates.remove_internal(), spritecollide(dokill=True), switch_layer()%n");
    }
}