package com.razz.dfashion.bbmodel;

import net.minecraft.client.model.geom.ModelPart;
import net.minecraft.core.Direction;
import org.joml.Vector3f;

import java.util.*;

public class BbmodelBaker {

    /** Backward-compat: cosmetic pipeline call site. {@code naturalUv=false} keeps vanilla's
     *  cube vertex order which is calibrated for the player avatar's pre-flip transforms. */
    public static Map<String, ModelPart> bake(Bbmodel model, int texWidth, int texHeight, boolean mirrorX) {
        return bake(model, texWidth, texHeight, mirrorX, false);
    }

    /**
     * @param naturalUv when true, polygons are constructed with a per-face vertex order that
     *                  produces a 1:1 UV→cube mapping (Blockbench's authoring convention),
     *                  matching decocraft's renderer. Use this for BlockEntity rendering, which
     *                  doesn't go through the avatar pre-flip pipeline. {@code mirrorX} should
     *                  be false in this mode.
     */
    public static Map<String, ModelPart> bake(Bbmodel model, int texWidth, int texHeight,
                                              boolean mirrorX, boolean naturalUv) {
        Map<String, BbCube> cubeIndex = new HashMap<>();
        for (BbCube c : model.elements()) cubeIndex.put(c.uuid(), c);

        Map<String, BbGroup> groupIndex = new HashMap<>();
        for (BbGroup g : model.groups()) groupIndex.put(g.uuid(), g);

        Map<String, ModelPart> result = new HashMap<>();
        for (BbOutlinerNode node : model.outliner()) {
            if (node instanceof BbOutlinerNode.GroupRef gref) {
                BbGroup g = groupIndex.get(gref.uuid());
                // Top-level: parentOrigin = the group's own origin so the ModelPart's
                // position resolves to (0, 0, 0). bone.translateAndRotate in the render
                // layer already places us at the bone, so the ModelPart itself doesn't
                // need to translate again.
                result.put(g.name(),
                        buildPart(gref, g, g.origin(), cubeIndex, groupIndex, texWidth, texHeight, mirrorX, naturalUv));
            }
            // ElementRef at root = loose cube, skip (no bone to attach to)
        }
        return result;
    }

    private static ModelPart buildPart(
            BbOutlinerNode.GroupRef gref,
            BbGroup group,
            Vector3f parentOrigin,
            Map<String, BbCube> cubeIndex,
            Map<String, BbGroup> groupIndex,
            int texWidth, int texHeight,
            boolean mirrorX,
            boolean naturalUv
    ) {
        List<ModelPart.Cube> cubes = new ArrayList<>();
        Map<String, ModelPart> children = new LinkedHashMap<>();

        for (BbOutlinerNode child : gref.children()) {
            switch (child) {
                case BbOutlinerNode.ElementRef er -> {
                    BbCube bb = cubeIndex.get(er.uuid());
                    if (bb == null) continue;                         // locator ref, skip

                    if (isRotated(bb)) {
                        // ModelPart.Cube is axis-aligned, so per-cube rotation is done
                        // by wrapping the cube in its own one-cube ModelPart whose pose
                        // carries the rotation.
                        ModelPart.Cube inner = bbCubeToCube(bb, bb.origin(), texWidth, texHeight, mirrorX, naturalUv);
                        ModelPart wrapper = new ModelPart(List.of(inner), Map.of());
                        wrapper.setPos(
                                bb.origin().x - group.origin().x,
                                bb.origin().y - group.origin().y,
                                bb.origin().z - group.origin().z
                        );
                        wrapper.setRotation(
                                (float) Math.toRadians(bb.rotation().x),
                                (float) Math.toRadians(bb.rotation().y),
                                (float) Math.toRadians(bb.rotation().z)
                        );
                        children.put("cube_" + bb.uuid(), wrapper);
                    } else {
                        cubes.add(bbCubeToCube(bb, group.origin(), texWidth, texHeight, mirrorX, naturalUv));
                    }
                }
                case BbOutlinerNode.GroupRef nested -> {
                    BbGroup ng = groupIndex.get(nested.uuid());
                    children.put(ng.name(),
                            buildPart(nested, ng, group.origin(), cubeIndex, groupIndex, texWidth, texHeight, mirrorX, naturalUv));
                }
            }
        }

        ModelPart part = new ModelPart(cubes, children);
        part.setPos(
                group.origin().x - parentOrigin.x,
                group.origin().y - parentOrigin.y,
                group.origin().z - parentOrigin.z
        );
        part.setRotation(
                (float) Math.toRadians(group.rotation().x),
                (float) Math.toRadians(group.rotation().y),
                (float) Math.toRadians(group.rotation().z)
        );
        return part;
    }

    private static boolean isRotated(BbCube bb) {
        Vector3f r = bb.rotation();
        return r != null && (r.x != 0f || r.y != 0f || r.z != 0f);
    }

    private static ModelPart.Cube bbCubeToCube(
            BbCube bb, Vector3f referenceOrigin, int texWidth, int texHeight,
            boolean mirrorX, boolean naturalUv
    ) {
        float x = bb.from().x - referenceOrigin.x;
        float y = bb.from().y - referenceOrigin.y;
        float z = bb.from().z - referenceOrigin.z;

        float w = bb.to().x - bb.from().x;
        float h = bb.to().y - bb.from().y;
        float d = bb.to().z - bb.from().z;

        // Which faces are visible — texture == -1 means hidden in bbmodel.
        Set<Direction> visibleFaces = EnumSet.noneOf(Direction.class);
        if (bb.faces() != null) {
            bb.faces().forEach((key, face) -> {
                if (face.texture() != null && face.texture() >= 0) {
                    Direction dir = directionFromName(key);
                    if (dir != null) visibleFaces.add(dir);
                }
            });
        }

        // Vanilla ctor gives us vertex geometry + placeholder box-UV polygons.
        // mirrorX=true flips the cube on X so faces orient correctly after the
        // player render-layer's scale(-1, -1, 1) un-flip. Block-entity renderers
        // don't apply that pre-flip, so they pass mirrorX=false.
        ModelPart.Cube cube = new ModelPart.Cube(
                0, 0, x, y, z, w, h, d, 0, 0, 0, mirrorX, texWidth, texHeight, visibleFaces
        );

        // Vanilla writes polygons[] in this order, skipping any direction not in visibleFaces.
        Direction[] order = {
                Direction.DOWN, Direction.UP, Direction.WEST,
                Direction.NORTH, Direction.EAST, Direction.SOUTH
        };

        // The 8 cube vertices, named by their world-space corner. Used by the natural-UV
        // path to reorder polygon vertices so each face's UV maps Blockbench's [u0,v0,u1,v1]
        // rectangle 1:1 onto the cube — top-left of UV → top-left of face, etc.
        ModelPart.Vertex nwBot = naturalUv ? new ModelPart.Vertex(x,     y,     z,     0, 0) : null;
        ModelPart.Vertex neBot = naturalUv ? new ModelPart.Vertex(x + w, y,     z,     0, 0) : null;
        ModelPart.Vertex neTop = naturalUv ? new ModelPart.Vertex(x + w, y + h, z,     0, 0) : null;
        ModelPart.Vertex nwTop = naturalUv ? new ModelPart.Vertex(x,     y + h, z,     0, 0) : null;
        ModelPart.Vertex swBot = naturalUv ? new ModelPart.Vertex(x,     y,     z + d, 0, 0) : null;
        ModelPart.Vertex seBot = naturalUv ? new ModelPart.Vertex(x + w, y,     z + d, 0, 0) : null;
        ModelPart.Vertex seTop = naturalUv ? new ModelPart.Vertex(x + w, y + h, z + d, 0, 0) : null;
        ModelPart.Vertex swTop = naturalUv ? new ModelPart.Vertex(x,     y + h, z + d, 0, 0) : null;

        int idx = 0;
        for (Direction dir : order) {
            if (!visibleFaces.contains(dir)) continue;
            BbFace bbFace = bb.faces().get(nameOf(dir));
            if (bbFace != null && bbFace.uv() != null && bbFace.uv().length >= 4) {
                ModelPart.Vertex[] verts = naturalUv
                        ? naturalVertsFor(dir, nwBot, neBot, neTop, nwTop, swBot, seBot, seTop, swTop)
                        : cube.polygons[idx].vertices();
                cube.polygons[idx] = new ModelPart.Polygon(
                        verts,
                        bbFace.uv()[0], bbFace.uv()[1],
                        bbFace.uv()[2], bbFace.uv()[3],
                        texWidth, texHeight,
                        mirrorX,
                        dir
                );
            }
            idx++;
        }

        return cube;
    }

    /**
     * Returns the four cube vertices for {@code dir} in (TR, TL, BL, BR) visual order, where
     * "visual" is the orientation the bbmodel author sees in Blockbench's per-face UV editor.
     * Vanilla's {@link ModelPart.Polygon} ctor remaps {@code vertices[0..3]} to UV
     * {@code (u1,v0)→(u0,v0)→(u0,v1)→(u1,v1)} — pairing this with (TR, TL, BL, BR) gives a
     * 1:1 mapping. UP/DOWN use a north-up convention (UV top = +Z = south? — see decocraft).
     *
     * <p>UP face: vertex[0]=NE-top → texture top-right. So texture's "north" (top) = cube's
     * north edge, texture's "east" (right) = cube's east edge — Blockbench's standard preview
     * orientation when looking down at the top face.
     *
     * <p>DOWN face: vertex[0]=SE-bot → texture top-right. Texture's "south" (top) = cube's
     * south edge — viewing from below as if you tilted your head back, with south "above" you.
     */
    private static ModelPart.Vertex[] naturalVertsFor(
            Direction dir,
            ModelPart.Vertex nwBot, ModelPart.Vertex neBot, ModelPart.Vertex neTop, ModelPart.Vertex nwTop,
            ModelPart.Vertex swBot, ModelPart.Vertex seBot, ModelPart.Vertex seTop, ModelPart.Vertex swTop
    ) {
        return switch (dir) {
            case NORTH -> new ModelPart.Vertex[]{nwTop, neTop, neBot, nwBot};
            case SOUTH -> new ModelPart.Vertex[]{seTop, swTop, swBot, seBot};
            case EAST  -> new ModelPart.Vertex[]{neTop, seTop, seBot, neBot};
            case WEST  -> new ModelPart.Vertex[]{swTop, nwTop, nwBot, swBot};
            case UP    -> new ModelPart.Vertex[]{neTop, nwTop, swTop, seTop};
            case DOWN  -> new ModelPart.Vertex[]{seBot, swBot, nwBot, neBot};
        };
    }

    private static String nameOf(Direction d) {
        return d.getName();
    }

    private static Direction directionFromName(String name) {
        return switch (name) {
            case "down"  -> Direction.DOWN;
            case "up"    -> Direction.UP;
            case "north" -> Direction.NORTH;
            case "south" -> Direction.SOUTH;
            case "west"  -> Direction.WEST;
            case "east"  -> Direction.EAST;
            default      -> null;
        };
    }
}