import eu.mhsl.minecraft.pixelpics.render.entity.*;
import eu.mhsl.minecraft.pixelpics.render.render.RenderJob;
import eu.mhsl.minecraft.pixelpics.render.snapshot.WorldSnapshot;
import org.bukkit.ChunkSnapshot;
import org.bukkit.Material;
import org.bukkit.block.data.BlockData;
import org.bukkit.util.Vector;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.lang.reflect.Proxy;
import java.util.*;
import java.util.List;

/**
 * Reproducible render benchmark (no server): farm scenes with hundreds of mobs / dropped items on a grass floor
 * plus a "normal" scene (few mobs, translucent + additive layers). Renders every scene {@code -Druns} times (after
 * {@code -Dwarmup} runs), prints median/min wall time and writes one PNG per scene to {@code -Dout} (default
 * ~/Dokumente/PixelPics-Entity-Renders/bench) for before/after pixel diffs. Args: scene names (default all):
 * farm_mobs, farm_empty (same view, no entities), farm_mobs_close, farm_items, village, normal.
 * {@code -Djava.util.concurrent.ForkJoinPool.common.parallelism=N} pins the thread count. Build: see {@link Harness}.
 */
public class EntityBenchmark {

    static final int W = 320, H = 240;

    record Scene(String name, Vector cam, Vector look, List<EntityState> mobs, List<ObjectState> objects) {}

    public static void main(String[] args) throws Exception {
        Harness h = new Harness("bench");
        WorldSnapshot world = farmWorld();
        int runs = Integer.getInteger("runs", 5), warmup = Integer.getInteger("warmup", 2);
        File outDir = new File(System.getProperty("out", "/home/elias/Dokumente/PixelPics-Entity-Renders/bench"));
        outDir.mkdirs();
        Set<String> only = new HashSet<>(Arrays.asList(args));

        for (Scene sc : scenes()) {
            if (!only.isEmpty() && !only.contains(sc.name)) continue;
            RenderJob job = Harness.job(world, sc.cam, sc.look, W, H, sc.mobs, List.of(), List.of(), sc.objects);
            BufferedImage img = null;
            for (int i = 0; i < warmup; i++) img = h.renderer.execute(job);
            long[] times = new long[runs];
            for (int i = 0; i < runs; i++) {
                long t0 = System.nanoTime();
                img = h.renderer.execute(job);
                times[i] = System.nanoTime() - t0;
            }
            Arrays.sort(times);
            ImageIO.write(img, "png", new File(outDir, sc.name + ".png"));
            System.out.printf("BENCH %-16s mobs=%4d objects=%4d  median %7.1f ms  min %7.1f ms%n", sc.name,
                sc.mobs.size(), sc.objects.size(), times[runs / 2] / 1e6, times[0] / 1e6);
        }
    }

    static List<Scene> scenes() throws ReflectiveOperationException {
        List<Scene> out = new ArrayList<>();
        // Mob farm: 400 animals in a 20x20 pen (1.4 block spacing), seen from an elevated corner.
        List<EntityState> farm = new ArrayList<>();
        Random r = new Random(42);
        String[] types = {"cow", "sheep", "chicken", "pig"};
        for (int i = 0; i < 20; i++) {
            for (int j = 0; j < 20; j++) {
                String type = types[r.nextInt(types.length)];
                farm.add(mob(type, i * 1.4 + r.nextDouble() * 0.4, 1, j * 1.4 + r.nextDouble() * 0.4,
                    r.nextFloat() * 360, r.nextInt(5) == 0));
            }
        }
        out.add(new Scene("farm_mobs", new Vector(-6, 14, -6), new Vector(14, 1, 14), farm, List.of()));
        // Same view without entities: farm_mobs minus this = the entity share of the render time.
        out.add(new Scene("farm_empty", new Vector(-6, 14, -6), new Vector(14, 1, 14), List.of(), List.of()));
        out.add(new Scene("farm_mobs_close", new Vector(3, 4.5, -3), new Vector(10, 1, 8), farm, List.of()));

        // Item farm: 500 dropped items (flat sprites + block items) piled on a 10x10 collection area.
        List<ObjectState> items = new ArrayList<>();
        String[] flat = {"wheat", "carrot", "egg", "bone", "rotten_flesh", "string", "gunpowder", "iron_ingot"};
        Material[] blocks = {Material.COBBLESTONE, Material.OAK_LOG, Material.PUMPKIN};
        String[] blockState = {null, "axis=y", null};
        for (int i = 0; i < 500; i++) {
            double x = 2 + r.nextDouble() * 10, z = 2 + r.nextDouble() * 10, y = 1.1 + r.nextDouble() * 0.3;
            Affine a = Affine.translation(x, y, z).mul(Affine.rotY(r.nextDouble() * Math.PI * 2)).mul(Affine.scale(0.5));
            ItemSpec spec;
            if (r.nextInt(4) == 0) {
                int b = r.nextInt(blocks.length);
                spec = new ItemSpec(blocks[b].getKey().getKey(), Harness.block(blocks[b], blockState[b]), null, false);
            } else {
                spec = Harness.item(flat[r.nextInt(flat.length)]);
            }
            items.add(ObjectState.item(a, spec, "ground", true));
        }
        out.add(new Scene("farm_items", new Vector(0, 7, -1), new Vector(7, 1, 7), List.of(), items));

        // Village: 300 mobs with per-mob texture compositing (villager type/profession/level overlays, dyed
        // sheep wool, dyed leather armor, tinted wolf collars) — exercises the tinted-texture cache.
        List<EntityState> village = new ArrayList<>();
        String[] vTypes = {"plains", "desert", "jungle", "savanna", "snow", "swamp", "taiga"};
        String[] profs = {"armorer", "butcher", "cartographer", "cleric", "farmer", "fisherman", "fletcher",
            "leatherworker", "librarian", "mason", "shepherd", "toolsmith", "weaponsmith", "nitwit", "none"};
        Random rv = new Random(7);
        for (int i = 0; i < 15; i++) {
            for (int j = 0; j < 20; j++) {
                double x = i * 1.4 + rv.nextDouble() * 0.3, z = j * 1.4 + rv.nextDouble() * 0.3;
                float yaw = rv.nextFloat() * 360;
                int dye = DYES[rv.nextInt(DYES.length)];
                EntityState m = switch (rv.nextInt(4)) {
                    case 0 -> villager(x, z, yaw, vTypes[rv.nextInt(vTypes.length)], profs[rv.nextInt(profs.length)],
                        1 + rv.nextInt(5), rv.nextInt(6) == 0);
                    case 1 -> with(mob("sheep", x, 1, z, yaw, rv.nextInt(6) == 0), dye, null, 0);
                    case 2 -> {
                        EntityState.EquipPiece le = new EntityState.EquipPiece("leather", dye, null, null, false);
                        yield with(mob("zombie", x, 1, z, yaw, false), 0, new EntityState.Equipment(le, le, le, le), 0);
                    }
                    default -> with(mob("wolf", x, 1, z, yaw, rv.nextInt(6) == 0), 0, null, dye);
                };
                village.add(m);
            }
        }
        out.add(new Scene("village", new Vector(-5, 10, -5), new Vector(10, 1, 13), village, List.of()));

        // Normal scene: a handful of mobs incl. translucent (slime) and additive (charged creeper) layers,
        // a text display with the default translucent background and two items.
        List<EntityState> few = new ArrayList<>();
        few.add(mob("cow", 3, 1, 3, 30, false));
        few.add(mob("sheep", 5, 1, 4, 200, false));
        few.add(mob("villager", 4, 1, 6, 160, false));
        few.add(mob("zombie", 6.5, 1, 6, 220, false));
        few.add(mob("slime", 2, 1, 5.5, 120, false));
        few.add(Harness.copy(mob("creeper", 7, 1, 3.5, 250, false), Map.of("powered", true)));
        few.add(mob("chicken", 3.5, 1, 1.5, 90, false));
        List<ObjectState> objs = new ArrayList<>();
        objs.add(ObjectState.item(Affine.translation(4.5, 1.2, 2).mul(Affine.scale(0.5)), Harness.item("apple"), "ground", true));
        objs.add(ObjectState.text(Affine.translation(4.5, 3.2, 4).mul(Affine.rotY(Math.PI)).mul(Affine.scale(0.025)),
            List.of("PixelPics bench"), 200, 0xFFFFFFFF, 0x40000000, true));
        out.add(new Scene("normal", new Vector(0, 3.5, -2.5), new Vector(4.5, 1.5, 4), few, objs));
        return out;
    }

    static EntityState mob(String type, double x, double y, double z, float yaw, boolean baby) {
        int tint = type.equals("sheep") ? 0xFFF9FFFE : 0;
        return new EntityState(type, x, y, z, yaw, baby, 0.8, 1.0,
            false, null, false, variant(type), tint, 1.0,
            type.equals("villager") ? "farmer" : null, type.equals("villager") ? 2 : 0,
            null, false, false, false, null, null,
            false, 0, 0, 0, false, null, null, null, 0, false, null, null, null, null, null);
    }

    /** Vanilla DyeColor texture colours (white … black). */
    static final int[] DYES = {0xFFF9FFFE, 0xFFF9801D, 0xFFC74EBD, 0xFF3AB3DA, 0xFFFED83D, 0xFF80C71F, 0xFFF38BAA,
        0xFF474F52, 0xFF9D9D97, 0xFF169C9C, 0xFF8932B8, 0xFF3C44AA, 0xFF835432, 0xFF5E7C16, 0xFFB02E26, 0xFF1D1D21};

    static EntityState villager(double x, double z, float yaw, String type, String prof, int level, boolean baby)
            throws ReflectiveOperationException {
        return Harness.copy(mob("villager", x, 1, z, yaw, baby), Map.of("width", 0.6, "height", 1.95, "variant", type,
            "profession", prof, "villagerLevel", level));
    }

    /** {@code s} with wool tint, humanoid armor and collar dye replaced. */
    static EntityState with(EntityState s, int tint, EntityState.Equipment eq, int collar)
            throws ReflectiveOperationException {
        Map<String, Object> o = new HashMap<>();
        o.put("tint", tint);
        o.put("equipment", eq);
        o.put("collar", collar);
        return Harness.copy(s, o);
    }

    /** The variant the game always supplies (cow/pig/chicken have no texture without one). */
    static String variant(String type) {
        return switch (type) {
            case "villager" -> "plains";
            case "cow", "pig", "chicken" -> "temperate";
            case "wolf" -> "pale";
            default -> null;
        };
    }

    /**
     * Grass floor at y=0 over [-48, 48)². Nothing stacked on it: ambient occlusion asks Material.isOccluding()
     * for non-air neighbours, which needs a server registry (not available standalone).
     */
    static WorldSnapshot farmWorld() {
        BlockData air = Harness.block(Material.AIR, null);
        BlockData grass = Harness.block(Material.GRASS_BLOCK, "snowy=false");
        Map<Long, ChunkSnapshot> chunks = new HashMap<>();
        for (int cx = -3; cx < 3; cx++) {
            for (int cz = -3; cz < 3; cz++) {
                ChunkSnapshot cs = (ChunkSnapshot) Proxy.newProxyInstance(EntityBenchmark.class.getClassLoader(),
                    new Class<?>[]{ChunkSnapshot.class}, (p, m, a) -> switch (m.getName()) {
                        case "getBlockData" -> (int) a[1] == 0 ? grass : air;
                        case "equals" -> p == a[0];
                        case "hashCode" -> System.identityHashCode(p);
                        case "toString" -> "chunk";
                        default -> Harness.defaultValue(m.getReturnType());
                    });
                chunks.put(WorldSnapshot.chunkKey(cx, cz), cs);
            }
        }
        return new WorldSnapshot(chunks, 0, 8, air);
    }
}