let scene, camera, renderer, skeletonHelper;
let joints = [];
let lines = [];
let currentMotion = null;
let frameIndex = 0;

function init() {
    scene = new THREE.Scene();
    camera = new THREE.PerspectiveCamera(75, window.innerWidth / window.innerHeight, 0.1, 1000);
    camera.position.set(2, 2, 5);
    camera.lookAt(0, 1, 0);

    renderer = new THREE.WebGLRenderer({ antialias: true });
    renderer.setSize(window.innerWidth, window.innerHeight);
    document.body.appendChild(renderer.domElement);

    const grid = new THREE.GridHelper(10, 10, 0x555555, 0x333333);
    scene.add(grid);

    const light = new THREE.PointLight(0xffffff, 1);
    light.position.set(10, 10, 10);
    scene.add(light);
    scene.add(new THREE.AmbientLight(0x444444));

    animate();
}

const PRESETS = {
    humanoid: [
        { name: "root", pos: [0, 1, 0], parent: -1 },
        { name: "spine", pos: [0, 1.2, 0], parent: 0 },
        { name: "head", pos: [0, 1.5, 0], parent: 1 },
        { name: "l_shoulder", pos: [-0.2, 1.3, 0], parent: 1 },
        { name: "r_shoulder", pos: [0.2, 1.3, 0], parent: 1 },
        { name: "l_elbow", pos: [-0.4, 1.3, 0], parent: 3 },
        { name: "r_elbow", pos: [0.4, 1.3, 0], parent: 4 },
        { name: "l_leg", pos: [-0.1, 0.9, 0], parent: 0 },
        { name: "r_leg", pos: [0.1, 0.9, 0], parent: 0 },
        { name: "l_knee", pos: [-0.1, 0.5, 0], parent: 7 },
        { name: "r_knee", pos: [0.1, 0.5, 0], parent: 8 },
        { name: "l_foot", pos: [-0.1, 0, 0], parent: 9 },
        { name: "r_foot", pos: [0.1, 0, 0], parent: 10 },
    ],
    centipede: [
        { name: "seg1", pos: [0, 0.2, 0], parent: -1 },
        { name: "seg2", pos: [0.5, 0.2, 0], parent: 0 },
        { name: "seg3", pos: [1.0, 0.2, 0], parent: 1 },
        { name: "seg4", pos: [1.5, 0.2, 0], parent: 2 },
        { name: "l1", pos: [0, 0, 0.3], parent: 0 },
        { name: "r1", pos: [0, 0, -0.3], parent: 0 },
        { name: "l2", pos: [0.5, 0, 0.3], parent: 1 },
        { name: "r2", pos: [0.5, 0, -0.3], parent: 1 },
        { name: "l3", pos: [1.0, 0, 0.3], parent: 2 },
        { name: "r3", pos: [1.0, 0, -0.3], parent: 2 },
    ]
};

function buildSkeleton(type) {
    joints.forEach(j => scene.remove(j));
    lines.forEach(l => scene.remove(l));
    joints = [];
    lines = [];

    const data = PRESETS[type] || PRESETS.humanoid;
    data.forEach(d => {
        const geo = new THREE.SphereGeometry(0.05);
        const mat = new THREE.MeshBasicMaterial({ color: 0x00ffcc });
        const mesh = new THREE.Mesh(geo, mat);
        mesh.position.set(...d.pos);
        scene.add(mesh);
        joints.push(mesh);
    });

    data.forEach((d, i) => {
        if (d.parent !== -1) {
            const geo = new THREE.BufferGeometry().setFromPoints([
                new THREE.Vector3(...d.pos),
                new THREE.Vector3(...data[d.parent].pos)
            ]);
            const mat = new THREE.LineBasicMaterial({ color: 0x888888 });
            const line = new THREE.Line(geo, mat);
            scene.add(line);
            lines.push({ line, start: i, end: d.parent });
        }
    });
}

async function runInference() {
    const type = document.getElementById("preset").value;
    const steps = document.getElementById("steps").value;
    const status = document.getElementById("status");
    
    status.innerText = "Computing Enrichment & Transformer Layers...";
    status.style.color = "#ff0";

    const response = await fetch("/infer", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({
            skeleton: { joints: PRESETS[type] || PRESETS.humanoid },
            steps: parseInt(steps)
        })
    });

    const data = await response.json();
    currentMotion = data.motion;
    frameIndex = 0;
    status.innerText = "Inference Complete. Playing Motion.";
    status.style.color = "#0f0";
}

function animate() {
    requestAnimationFrame(animate);
    
    if (currentMotion) {
        const frame = currentMotion[frameIndex];
        frame.forEach((pos, i) => {
            if (joints[i]) joints[i].position.set(pos[0], pos[1], pos[2]);
        });

        lines.forEach(l => {
            const p1 = joints[l.start].position;
            const p2 = joints[l.end].position;
            l.line.geometry.setFromPoints([p1, p2]);
        });

        frameIndex = (frameIndex + 1) % currentMotion.length;
    }

    renderer.render(scene, camera);
}

window.addEventListener('resize', () => {
    camera.aspect = window.innerWidth / window.innerHeight;
    camera.updateProjectionMatrix();
    renderer.setSize(window.innerWidth, window.innerHeight);
});

document.getElementById("run").onclick = runInference;
document.getElementById("preset").onchange = (e) => buildSkeleton(e.target.value);

init();
buildSkeleton("humanoid");