import com.jme3.animation.AnimControl;
import com.jme3.animation.Animation;
import com.jme3.animation.Bone;
import com.jme3.animation.BoneTrack;
import com.jme3.animation.Skeleton;
import com.jme3.animation.Track;
import com.jme3.export.binary.BinaryImporter;
import com.jme3.math.Quaternion;
import com.jme3.math.Vector3f;
import com.jme3.scene.Geometry;
import com.jme3.scene.Mesh;
import com.jme3.scene.Node;
import com.jme3.scene.Spatial;
import com.jme3.scene.VertexBuffer;
import com.jme3.util.IntMap;
import com.jme3.util.SafeArrayList;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.FloatBuffer;
import java.nio.ShortBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Logger;

public class J3OConverter {

    private static final Logger logger = Logger.getLogger(J3OConverter.class.getName());

    public static void main(String[] args) {
        if (args.length < 2) {
            System.err.println("Usage: J3OConverter <input.j3o> <output.json>");
            System.exit(1);
        }

        File inputFile = new File(args[0]);
        File outputFile = new File(args[1]);

        try {
            BinaryImporter importer = new BinaryImporter();
            Object root = importer.load(inputFile);

            Map<String, Object> result = new HashMap<>();
            List<Map<String, Object>> geometries = new ArrayList<>();
            List<Map<String, Object>> skeletons = new ArrayList<>();
            Map<String, Object> animations = new HashMap<>();

            if (root instanceof Node) {
                extractFromNode((Node) root, geometries, skeletons, animations, new HashMap<>());
            }

            result.put("geometries", geometries);
            result.put("skeletons", skeletons);
            result.put("animations", animations);

            writeJson(outputFile, result);
            System.out.println("Wrote " + outputFile);
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
    }

    private static void extractFromNode(Node node,
                                         List<Map<String, Object>> geometries,
                                         List<Map<String, Object>> skeletons,
                                         Map<String, Object> animations,
                                         Map<Integer, Integer> boneIdMap) {
        for (Spatial child : node.getChildren()) {
            if (child instanceof Node) {
                extractFromNode((Node) child, geometries, skeletons, animations, boneIdMap);
            } else if (child instanceof Geometry) {
                Geometry geo = (Geometry) child;
                Map<String, Object> meshData = extractMesh(geo.getMesh());
                if (meshData != null) {
                    meshData.put("name", geo.getName());
                    geometries.add(meshData);
                }
            }
        }

        // Look for AnimControl
        AnimControl animControl = node.getControl(AnimControl.class);
        if (animControl != null) {
            Skeleton skeleton = animControl.getSkeleton();
            if (skeleton != null) {
                Map<String, Object> skeletonData = extractSkeleton(skeleton);
                if (skeletonData != null) {
                    skeletons.add(skeletonData);
                }
            }

            for (String animName : animControl.getAnimationNames()) {
                Animation anim = animControl.getAnim(animName);
                Map<String, Object> animData = extractAnimation(anim);
                if (animData != null) {
                    animations.put(animName, animData);
                }
            }
        }
    }

    private static Map<String, Object> extractMesh(Mesh mesh) {
        Map<String, Object> result = new HashMap<>();
        result.put("vertexCount", mesh.getVertexCount());
        result.put("triangleCount", mesh.getTriangleCount());
        result.put("elementCount", mesh.getTriangleCount() * 3);

        // Positions
        VertexBuffer posBuf = mesh.getBuffer(VertexBuffer.Type.Position);
        if (posBuf != null) {
            FloatBuffer fb = (FloatBuffer) posBuf.getData();
            float[] positions = new float[fb.limit()];
            fb.rewind();
            fb.get(positions);
            result.put("positions", positions);
        }

        // Normals
        VertexBuffer normBuf = mesh.getBuffer(VertexBuffer.Type.Normal);
        if (normBuf != null) {
            FloatBuffer fb = (FloatBuffer) normBuf.getData();
            float[] normals = new float[fb.limit()];
            fb.rewind();
            fb.get(normals);
            result.put("normals", normals);
        }

        // TexCoords
        VertexBuffer texBuf = mesh.getBuffer(VertexBuffer.Type.TexCoord);
        if (texBuf != null) {
            FloatBuffer fb = (FloatBuffer) texBuf.getData();
            float[] texcoords = new float[fb.limit()];
            fb.rewind();
            fb.get(texcoords);
            result.put("texcoords", texcoords);
        }

        // BoneWeights
        VertexBuffer weightBuf = mesh.getBuffer(VertexBuffer.Type.BoneWeight);
        if (weightBuf != null) {
            FloatBuffer fb = (FloatBuffer) weightBuf.getData();
            float[] weights = new float[fb.limit()];
            fb.rewind();
            fb.get(weights);
            result.put("boneWeights", weights);
        }

        // BoneIndices
        VertexBuffer indexBuf = mesh.getBuffer(VertexBuffer.Type.BoneIndex);
        if (indexBuf != null) {
            java.nio.Buffer buf = indexBuf.getData();
            short[] indices;
            if (buf instanceof ShortBuffer) {
                ShortBuffer sb = (ShortBuffer) buf;
                indices = new short[sb.limit()];
                sb.rewind();
                sb.get(indices);
                int[] intIndices = new int[indices.length];
                for (int i = 0; i < indices.length; i++) {
                    intIndices[i] = indices[i] & 0xFFFF;
                }
                result.put("boneIndices", intIndices);
            } else if (buf instanceof java.nio.ByteBuffer) {
                java.nio.ByteBuffer bb = (java.nio.ByteBuffer) buf;
                byte[] bytes = new byte[bb.limit()];
                bb.rewind();
                bb.get(bytes);
                int[] intIndices = new int[bytes.length];
                for (int i = 0; i < bytes.length; i++) {
                    intIndices[i] = bytes[i] & 0xFF;
                }
                result.put("boneIndices", intIndices);
            }
        }

        // Indices - JME3 doesn't always have an index buffer (uses non-indexed rendering)
        // Try to get the indices from the mesh's LOD levels or mode
        VertexBuffer indexBuffer = mesh.getBuffer(VertexBuffer.Type.Index);
        if (indexBuffer != null) {
            java.nio.Buffer buf = indexBuffer.getData();
            if (buf instanceof ShortBuffer) {
                ShortBuffer sb = (ShortBuffer) buf;
                short[] indices = new short[sb.limit()];
                sb.rewind();
                sb.get(indices);
                int[] intIndices = new int[indices.length];
                for (int i = 0; i < indices.length; i++) {
                    intIndices[i] = indices[i] & 0xFFFF;
                }
                result.put("indices", intIndices);
            } else if (buf instanceof java.nio.IntBuffer) {
                java.nio.IntBuffer ib = (java.nio.IntBuffer) buf;
                int[] indices = new int[ib.limit()];
                ib.rewind();
                ib.get(indices);
                result.put("indices", indices);
            }
        }

        if (!result.containsKey("indices")) {
            // Non-indexed rendering: generate sequential indices
            int count = mesh.getVertexCount();
            int[] indices = new int[count];
            for (int i = 0; i < count; i++) {
                indices[i] = i;
            }
            result.put("indices", indices);
            result.put("nonIndexed", true);
        }

        return result;
    }

    private static Map<String, Object> extractSkeleton(Skeleton skeleton) {
        Map<String, Object> result = new HashMap<>();

        Bone[] bones = skeleton.getBones();
        List<Map<String, Object>> boneList = new ArrayList<>();
        Map<String, Integer> nameToIndex = new HashMap<>();

        for (int i = 0; i < bones.length; i++) {
            Bone bone = bones[i];
            nameToIndex.put(bone.getName(), i);

            Map<String, Object> bdata = new HashMap<>();
            bdata.put("name", bone.getName());
            bdata.put("id", i);

            Vector3f pos = bone.getBindPosition();
            bdata.put("position", new float[]{pos.x, pos.y, pos.z});

            Quaternion rot = bone.getBindRotation();
            bdata.put("rotation", new float[]{rot.getX(), rot.getY(), rot.getZ(), rot.getW()});

            Vector3f scale = bone.getBindScale();
            bdata.put("scale", new float[]{scale.x, scale.y, scale.z});

            // Get children bone names
            List<String> children = new ArrayList<>();
            for (Bone child : bone.getChildren()) {
                children.add(child.getName());
            }
            bdata.put("children", children);

            boneList.add(bdata);
        }

        result.put("bones", boneList);

        // Roots
        List<String> roots = new ArrayList<>();
        Bone[] rootBones = skeleton.getRoots();
        for (Bone root : rootBones) {
            roots.add(root.getName());
        }
        result.put("roots", roots);

        return result;
    }

    private static Map<String, Object> extractAnimation(Animation anim) {
        Map<String, Object> result = new HashMap<>();
        result.put("name", anim.getName());
        result.put("length", anim.getLength());

        Track[] tracks = anim.getTracks();
        List<Map<String, Object>> trackList = new ArrayList<>();

        for (Track track : tracks) {
            if (track instanceof BoneTrack) {
                BoneTrack boneTrack = (BoneTrack) track;
                Map<String, Object> trackData = new HashMap<>();
                trackData.put("boneIndex", boneTrack.getTargetBoneIndex());
                trackData.put("times", boneTrack.getTimes());

                // Translations
                Vector3f[] translations = boneTrack.getTranslations();
                float[] transData = new float[translations.length * 3];
                for (int i = 0; i < translations.length; i++) {
                    transData[i * 3] = translations[i].x;
                    transData[i * 3 + 1] = translations[i].y;
                    transData[i * 3 + 2] = translations[i].z;
                }
                trackData.put("translations", transData);

                // Rotations
                Quaternion[] rotations = boneTrack.getRotations();
                float[] rotData = new float[rotations.length * 4];
                for (int i = 0; i < rotations.length; i++) {
                    rotData[i * 4] = rotations[i].getX();
                    rotData[i * 4 + 1] = rotations[i].getY();
                    rotData[i * 4 + 2] = rotations[i].getZ();
                    rotData[i * 4 + 3] = rotations[i].getW();
                }
                trackData.put("rotations", rotData);

                // Scales
                Vector3f[] scales = boneTrack.getScales();
                float[] scaleData = new float[scales.length * 3];
                for (int i = 0; i < scales.length; i++) {
                    scaleData[i * 3] = scales[i].x;
                    scaleData[i * 3 + 1] = scales[i].y;
                    scaleData[i * 3 + 2] = scales[i].z;
                }
                trackData.put("scales", scaleData);

                trackList.add(trackData);
            }
        }

        result.put("tracks", trackList);
        return result;
    }

    private static void writeJson(File file, Object data) throws IOException {
        StringBuilder sb = new StringBuilder();
        writeValue(sb, data);
        FileWriter writer = new FileWriter(file);
        writer.write(sb.toString());
        writer.close();
    }

    private static void writeValue(StringBuilder sb, Object value) {
        if (value == null) {
            sb.append("null");
        } else if (value instanceof String) {
            sb.append('"').append(escapeString((String) value)).append('"');
        } else if (value instanceof Integer || value instanceof Long ||
                   value instanceof Float || value instanceof Double ||
                   value instanceof Boolean) {
            sb.append(value);
        } else if (value instanceof float[]) {
            sb.append('[');
            float[] arr = (float[]) value;
            for (int i = 0; i < arr.length; i++) {
                if (i > 0) sb.append(',');
                sb.append(arr[i]);
            }
            sb.append(']');
        } else if (value instanceof int[]) {
            sb.append('[');
            int[] arr = (int[]) value;
            for (int i = 0; i < arr.length; i++) {
                if (i > 0) sb.append(',');
                sb.append(arr[i]);
            }
            sb.append(']');
        } else if (value instanceof List) {
            sb.append('[');
            List<?> list = (List<?>) value;
            for (int i = 0; i < list.size(); i++) {
                if (i > 0) sb.append(',');
                writeValue(sb, list.get(i));
            }
            sb.append(']');
        } else if (value instanceof Map) {
            sb.append('{');
            Map<?, ?> map = (Map<?, ?>) value;
            boolean first = true;
            for (Map.Entry<?, ?> entry : map.entrySet()) {
                if (!first) sb.append(',');
                sb.append('"').append(escapeString(entry.getKey().toString())).append('"').append(':');
                writeValue(sb, entry.getValue());
                first = false;
            }
            sb.append('}');
        }
    }

    private static String escapeString(String s) {
        StringBuilder sb = new StringBuilder();
        for (char c : s.toCharArray()) {
            switch (c) {
                case '"': sb.append("\\\""); break;
                case '\\': sb.append("\\\\"); break;
                case '\n': sb.append("\\n"); break;
                case '\r': sb.append("\\r"); break;
                case '\t': sb.append("\\t"); break;
                default:
                    if (c < 32) {
                        sb.append(String.format("\\u%04x", (int) c));
                    } else {
                        sb.append(c);
                    }
            }
        }
        return sb.toString();
    }
}