import * as THREE from 'three';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
import URDFLoader from './urdf/URDFLoader.js';

// Setup scene, camera, renderer
const canvas = document.getElementById('urdf-canvas');
const renderer = new THREE.WebGLRenderer({ canvas, antialias: true });
renderer.setSize(window.innerWidth, window.innerHeight);

const scene = new THREE.Scene();
scene.background = new THREE.Color(0x111111);

const camera = new THREE.PerspectiveCamera(45, window.innerWidth / window.innerHeight, 0.1, 100);
camera.position.set(2, 2, 2);
camera.lookAt(0, 0, 0);

// Lights
scene.add(new THREE.AmbientLight(0xffffff, 0.6));
const directional = new THREE.DirectionalLight(0xffffff, 0.8);
directional.position.set(15, 10, 7);
scene.add(directional);

// Controls
const controls = new OrbitControls(camera, renderer.domElement);
controls.target.set(0, 0.5, 0);
controls.update();

// Helpers
scene.add(new THREE.AxesHelper(1));
const gridHelper = new THREE.GridHelper(2, 10);
//gridHelper.rotation.x = Math.PI / 2;
gridHelper.position.z = 0;
scene.add(gridHelper);

// Raycasting
const raycaster = new THREE.Raycaster();
const mouse = new THREE.Vector2();
let hoveredJoint = null;
let isDragging = false;
let lastX = null;
let worldAxis = null;

let draggedJoint = null;
let arc = null;
const jointAngles = {}; // key: jointName, value: current angle in radians




function findJointAncestor(object) {
    while (object && object.parent) {
        if (object.type === 'URDFJoint') return object;
        object = object.parent;
    }
    return null;
}

function createRotationArc(axis, lower, upper, radius = 0.1, segments = 32) {
    const points = [];
    console.log(axis);
    console.log(axis.clone());
    // ✅ Ensure axis is normalized
    const normalizedAxis = axis.clone().normalize();

    // ✅ Find a perpendicular vector
    let perp = new THREE.Vector3(0, 1, 0);
    if (Math.abs(normalizedAxis.dot(perp)) > 0.99) {
        perp = new THREE.Vector3(1, 0, 0); // fallback if axis is nearly parallel to Y
    }
    const tangent = perp.cross(normalizedAxis).normalize();

    for (let i = 0; i <= segments; i++) {
        const angle = lower + (upper - lower) * (i / segments);
        const point = tangent.clone().applyAxisAngle(normalizedAxis, angle).multiplyScalar(radius);
        points.push(point);
    }

    const geometry = new THREE.BufferGeometry().setFromPoints(points);
    const material = new THREE.LineBasicMaterial({ color: 0xffff00 });
    return new THREE.Line(geometry, material);
}

function createRotationSector(axis, lower, upper, radius = 0.1, segments = 64) {
    const shape = new THREE.Shape();
    shape.moveTo(0, 0);

    for (let i = 0; i <= segments; i++) {
        const angle = lower + (upper - lower) * (i / segments);
        shape.lineTo(Math.cos(angle) * radius, Math.sin(angle) * radius);
    }

    shape.lineTo(0, 0); // close the shape

    const geometry = new THREE.ShapeGeometry(shape);
    const material = new THREE.MeshBasicMaterial({
        color: 0x888888,         // light gray
        transparent: true,
        opacity: 0.3,            // translucent
        side: THREE.DoubleSide,  // show both sides of the fill
        depthTest: false,        // ✅ always on top
        depthWrite: false
    });


    const mesh = new THREE.Mesh(geometry, material);

    // Orient the sector to match the axis
    const basis = new THREE.Matrix4();
    const up = new THREE.Vector3(0, 0, 1);
    const quaternion = new THREE.Quaternion().setFromUnitVectors(up, axis.clone().normalize());
    basis.makeRotationFromQuaternion(quaternion);
    mesh.applyMatrix4(basis);

    return mesh;
}

function createAngleIndicator(axis, angle, radius = 0.1) {
    const dir = new THREE.Vector3(Math.cos(angle), Math.sin(angle), 0).multiplyScalar(radius);

    const geometry = new THREE.BufferGeometry().setFromPoints([
        new THREE.Vector3(0, 0, 0),
        dir,
    ]);

    const material = new THREE.LineBasicMaterial({
  color: 0xffffff,         // white
  linewidth: 2,            // note: linewidth only works in some renderers
  depthTest: false,        // ✅ always on top
  depthWrite: false
});

    const line = new THREE.Line(geometry, material);

    // Rotate into axis plane
    const up = new THREE.Vector3(0, 0, 1);
    const quaternion = new THREE.Quaternion().setFromUnitVectors(up, axis.clone().normalize());
    line.applyQuaternion(quaternion);

    return line;
}



// Load URDF
const loader = new URDFLoader();
loader.packages = { '': './urdf/' };

function findObjectByName(root, name) {
    let result = null;
    root.traverse(child => {
        if (child.name === name) result = child;
    });
    return result;
}

loader.load('./urdf/sample.urdf', robot => {
    //scene.rotation.x = -Math.PI / 2;
    robot.rotation.x = -Math.PI / 2;
    scene.add(robot);
    document.getElementById('status').textContent = 'URDF Loaded';

    loader.parseVisual = true;
    loader.parseCollision = false;

    // Add AxesHelpers to all meshes
    robot.traverse(obj => {
        if (obj.isMesh) obj.add(new THREE.AxesHelper(0.05));
    });

    // Setup pointer events
    canvas.addEventListener('pointermove', event => {
        mouse.x = (event.clientX / canvas.clientWidth) * 2 - 1;
        mouse.y = -(event.clientY / canvas.clientHeight) * 2 + 1;

        raycaster.setFromCamera(mouse, camera);
        const intersects = raycaster.intersectObjects(robot.children, true);

        if (isDragging && draggedJoint && worldAxis) {
            const deltaX = event.clientX - lastX;
            lastX = event.clientX;
            const angleDelta = deltaX * 0.005;

            const jointName = draggedJoint.name;
            const limits = robot.joints[jointName]?.limit;
            const lower = limits?.lower ?? -Math.PI;
            const upper = limits?.upper ?? Math.PI;

            const currentAngle = jointAngles[jointName] ?? 0;
            const proposedAngle = currentAngle + angleDelta;

            // ✅ Clamp to limits
            const clampedAngle = Math.max(lower, Math.min(upper, proposedAngle));
            const actualDelta = clampedAngle - currentAngle;

            // ✅ Apply rotation
            draggedJoint.rotateOnAxis(worldAxis, actualDelta);
            jointAngles[jointName] = clampedAngle;


            // Remove old indicator
            if (draggedJoint.userData.gizmo?.indicator) {
                draggedJoint.parent.remove(draggedJoint.userData.gizmo.indicator);
            }

            // Add updated indicator
            const newIndicator = createAngleIndicator(worldAxis, jointAngles[jointName]);
            newIndicator.position.copy(draggedJoint.position);
            draggedJoint.parent.add(newIndicator);
            draggedJoint.userData.gizmo.indicator = newIndicator;

        }



        if (intersects.length > 0) {
            const target = intersects[0].object;

            // Reset previous highlight
            if (
                hoveredJoint &&
                hoveredJoint.material &&
                'emissive' in hoveredJoint.material &&
                typeof hoveredJoint.material.emissive.setHex === 'function'
            ) {
                hoveredJoint.material.emissive.setHex(0x000000);
            }

            hoveredJoint = target;

            // Highlight new joint
            if (
                hoveredJoint.material &&
                'emissive' in hoveredJoint.material &&
                typeof hoveredJoint.material.emissive.setHex === 'function'
            ) {
                hoveredJoint.material.emissive.setHex(0x333333);
            }

        } else {
            if (
                hoveredJoint &&
                hoveredJoint.material &&
                'emissive' in hoveredJoint.material &&
                typeof hoveredJoint.material.emissive.setHex === 'function'
            ) {
                hoveredJoint.material.emissive.setHex(0x000000);
            }
            hoveredJoint = null;
        }

    });

    canvas.addEventListener('pointerdown', event => {
        if (!hoveredJoint) return;

        isDragging = true;
        lastX = event.clientX;
        controls.enabled = false;

        draggedJoint = findJointAncestor(hoveredJoint);


        if (!draggedJoint) return;

        const jointName = draggedJoint.name;
        const jointData = robot.joints?.[jointName];

        // ✅ Check if axis is a valid THREE.Vector3
        if (!(jointData?.axis instanceof THREE.Vector3)) {
            console.warn(`Invalid axis for joint "${jointName}":`, jointData?.axis);
            return;
        }

        const urdfAxis = jointData.axis.clone();
        if (urdfAxis.lengthSq() === 0) {
            console.warn(`Zero-length axis for joint "${jointName}"`);
            return;
        }

        worldAxis = urdfAxis.normalize();
        const limits = robot.joints[jointName]?.limit;
        const lower = limits?.lower ?? -Math.PI;
        const upper = limits?.upper ?? Math.PI;

        const sector = createRotationSector(worldAxis, lower, upper);
        const indicator = createAngleIndicator(worldAxis, jointAngles[jointName] ?? 0);

        sector.position.copy(draggedJoint.position);
        indicator.position.copy(draggedJoint.position);

        draggedJoint.parent.add(sector);
        draggedJoint.parent.add(indicator);

        // Store for cleanup
        draggedJoint.userData.gizmo = { sector, indicator };




        console.log(`Dragging joint "${jointName}" on axis`, worldAxis.toArray());
    });






    canvas.addEventListener('pointerup', () => {
        isDragging = false;
        controls.enabled = true;
        const gizmo = draggedJoint?.userData.gizmo;
        if (gizmo) {
            draggedJoint.parent.remove(gizmo.sector);
            draggedJoint.parent.remove(gizmo.indicator);
        }


        draggedJoint = null;
        worldAxis = null;
    });


    animate();
});

function animate() {
    requestAnimationFrame(animate);
    renderer.render(scene, camera);
}