import * as THREE from 'three';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
//import URDFLoader from './urdf/ExtendedURDFLoader.js';
import { createRotationSector, createAngleIndicator, createJointLabel } from './JointVisualiser.js';
import ExtendedURDFLoader from './urdf/ExtendedURDFLoader.js';
import { OBJLoader } from 'three/examples/jsm/loaders/OBJLoader.js';
import { MTLLoader } from 'three/examples/jsm/loaders/MTLLoader.js';
import { GLTFLoader } from 'three/examples/jsm/loaders/GLTFLoader.js';
import { ViewerOverlay } from './ViewerOverlay.js';
import { clamp } from 'three/src/math/MathUtils.js';
const gltfLoader = new GLTFLoader();

const SyncMode = {
  None: 0,
  SimToReal: 1,
  RealToSim: 2,
  Calibrate: 3
};

export class URDFEditor {
  constructor(canvas, sendMotorPosition, serial, curveEditor, tryConnect) {
    this.canvas = canvas;
    this.sendMotorPosition = sendMotorPosition;
    this.serial = serial;
    this.curveEditor = curveEditor;
    this.tryConnect = tryConnect;
    this.scene = new THREE.Scene();
    this.scene.background = new THREE.Color(0xaaaaaa);

    this.camera = new THREE.PerspectiveCamera(60, canvas.clientWidth / canvas.clientHeight, 0.1, 100);
    this.camera.position.set(-0.3, 0.5, -0.7);
    this.camera.lookAt(0, 0, 0);

    this.renderer = new THREE.WebGLRenderer({ canvas, antialias: true });
    this.renderer.setSize(canvas.clientWidth, canvas.clientHeight);



    this.controls = new OrbitControls(this.camera, this.renderer.domElement);
    this.controls.target.set(0, 0.1, 0);
    this.controls.update();

    this.raycaster = new THREE.Raycaster();
    this.mouse = new THREE.Vector2();

    this.currentSyncMode = SyncMode.None;

    this.hoveredJoint = null;
    this.draggedJoint = null;
    this.selectedJoint = null;
    this.worldAxis = null;
    this.lastX = null;
    this.isDragging = false;
    this.jointAngles = {};

    this.loader = new ExtendedURDFLoader();
    this.loader.loadMeshCb = (path, manager, onComplete) => {
      if (path.endsWith('.glb') || path.endsWith('.gltf')) {
        const gltfLoader = new GLTFLoader(manager);
        gltfLoader.load(path, (gltf) => {
          onComplete(gltf.scene);
        });
      } else {
        // fall back to default behavior (STL, DAE)
        URDFLoader.prototype.loadMeshCb.call(this.loader, path, manager, onComplete);
      }
    };

    this.loader.packages = {
      'Little_Sophia_Face': '/robots/LittleSophia'
    };
    this.loader.parseVisual = true;
    this.loader.parseCollision = false;


    this.setupScene();
    //this.loadURDF();
    //this.loadURDFFromIndexedDB();
    this.setupEvents();

    const editorCallbacks = {}


    this.overlay = new ViewerOverlay(
      this.renderer,
      this.robot,
      this.jointAngles,
      this.findObjectByName.bind(this),
      this
    );


    // Add some lights (optional if using MeshBasicMaterial)
    const ambient = new THREE.AmbientLight(0x404040);
    this.scene.add(ambient);
    const dirLight = new THREE.DirectionalLight(0xffffff, 1);
    dirLight.position.set(5, 5, 5);
    this.scene.add(dirLight);

    // this.mesh = null;
    // gltfLoader.load('meshes/Little_Sophia.glb', (gltf) => {
    //   const model = gltf.scene;
    //   model.scale.set(0.001, 0.001, 0.001); // adjust if needed
    //   model.rotation.y = Math.PI; // radians (180°)
    //   this.scene.add(model);

    //   // Optional: fit camera to model
    //   const box = new THREE.Box3().setFromObject(model);
    //   const center = box.getCenter(new THREE.Vector3());
    //   const size = box.getSize(new THREE.Vector3()).length();
    // });


    this.animate();
  }

  setupScene() {
    this.scene.add(new THREE.AmbientLight(0xffffff, 0.6));
    const light1 = new THREE.DirectionalLight(0xffffff, 0.8);
    light1.position.set(15, 10, 7);
    this.scene.add(light1);

    const light2 = new THREE.DirectionalLight(0xffffff, 0.8);
    light2.position.set(-15, 10, -7);
    this.scene.add(light2);

    this.scene.add(new THREE.AxesHelper(1));
    this.scene.add(new THREE.GridHelper(2, 10));


  }

  resetEditor() {
    if (!this.robot) return;
    const toRemove = this.scene.children.filter(c => c.type === 'URDFLink' || c.constructor.name === 'URDFRobot');
    toRemove.forEach(robot => {
      // detach from scene
      this.scene.remove(robot);

      // dispose GPU resources
      robot.traverse(obj => {
        if (obj.geometry) obj.geometry.dispose();
        if (obj.material) {
          if (Array.isArray(obj.material)) obj.material.forEach(m => m.dispose());
          else obj.material.dispose();
        }
      });
    });

    // Clear references
    this.robot.urdfRobotNode = null;
    this.robot.joints = {};
    this.robot.links = {};
    this.robot.sceneObject = null;
    this.robot = null;

    this.overlay.resetEditor();
  }

  setAnimationLength(totalFrames) {
    //this.overlay.
  }

  async loadURDF() {
    if (this.robot) {
      console.log("resetting editor");
      this.resetEditor();
    }

    this.urdfPath = '/robots/LittleSophia/urdf/LittleSophia.urdf';

    const urdfText = await fetch(this.urdfPath).then(res => res.text());
    const robot = await this.loader.loadFromString(urdfText);
    //addOriginMarkers(robot);
    robot.rotation.x = -Math.PI / 2;
    this.scene.add(robot);
    this.robot = robot;

    //this.overlay.robot = robot;
    this.overlay.init(robot);

  }

  async loadURDFFromIndexedDB() {
    if (this.robot) {
      console.log("resetting editor");
      this.resetEditor();
    }

    return new Promise((resolve, reject) => {
      const request = indexedDB.open("URDFEditorDB", 1);

      request.onerror = e => reject(e);

      request.onsuccess = e => {
        const db = e.target.result;
        const tx = db.transaction("urdfs", "readonly");
        const getReq = tx.objectStore("urdfs").get("current");

        getReq.onsuccess = async () => {
          const urdfText = getReq.result;
          if (!urdfText) {
            reject("No URDF stored in IndexedDB");
            return;
          }

          try {
            // ✅ feed the text into your loader
            const robot = await this.loader.loadFromString(urdfText);
            robot.rotation.x = -Math.PI / 2;
            this.scene.add(robot);
            this.robot = robot;
            this.overlay.init(robot);
            this.saveURDFToIndexedDB(robot);
            resolve(robot);
          } catch (err) {
            reject(err);
          }
        };
      };
    });
  }

  // WRITE TO URDF POSITION IN ROBOT
  // Find the transmission node for neck_roll
  // const neckTrans = this.robot.urdfRobotNode.querySelector('transmission[name="neck_roll_trans"]');
  // if (neckTrans) {
  //   const actuator = neckTrans.querySelector('actuator[name="neck_roll"]');
  //   if (actuator) {
  //     // Example: change encoderValidMin/Max
  //     actuator.querySelector('encoderValidMin').textContent = "100";
  //     actuator.querySelector('encoderValidMax').textContent = "3900";
  //   }
  // }

  downloadURDF() {

    const serializer = new XMLSerializer();
    const urdfString = serializer.serializeToString(this.robot.urdfRobotNode);

    const a = document.createElement("a");
    a.href = "data:text/xml;charset=utf-8," + encodeURIComponent(urdfString);
    a.download = "robot.urdf";
    a.click();
  }

  async uploadURDF() {
    // 1. Let the user pick a file
    const input = document.createElement('input');
    input.type = 'file';
    input.accept = '.urdf,.xml'; // restrict to URDF/XML files

    if (this.robot) {
      console.log("resetting editor");
      this.resetEditor();
    }

    input.onchange = async (e) => {
      const file = e.target.files[0];
      if (!file) return;

      // 2. Reset editor if a robot is already loaded
      if (this.robot) {
        console.log("resetting editor");
        this.resetEditor();
      }

      // 3. Read the file contents
      const urdfText = await file.text();

      // 4. Parse and load robot
      const robot = await this.loader.loadFromString(urdfText);
      robot.rotation.x = -Math.PI / 2;
      this.scene.add(robot);
      this.robot = robot;

      // 5. Initialize overlay
      this.overlay.init(robot);
    };

    // trigger the file picker
    input.click();
  }


  saveURDFToIndexedDB(robot) {
    const serializer = new XMLSerializer();
    const urdfString = serializer.serializeToString(robot.urdfRobotNode);

    const request = indexedDB.open("URDFEditorDB", 1);
    request.onupgradeneeded = e => {
      const db = e.target.result;
      db.createObjectStore("urdfs");
    };
    request.onsuccess = e => {
      const db = e.target.result;
      const tx = db.transaction("urdfs", "readwrite");
      tx.objectStore("urdfs").put(urdfString, "current");
      console.log("saved URDF to IndexedDB");
    };

  }

  listUrdfsFromIndexedDB() {
    const request = indexedDB.open("URDFEditorDB", 1);

    request.onsuccess = e => {
      const db = e.target.result;
      const tx = db.transaction("urdfs", "readonly");
      const store = tx.objectStore("urdfs");
      const getAllReq = store.getAll();

      getAllReq.onsuccess = () => {
        console.log("Stored URDFs:", getAllReq.result);
      };
    };
  }


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


  setupEvents() {
    this.canvas.addEventListener('pointermove', this.onPointerMove.bind(this));
    this.canvas.addEventListener('pointerdown', this.onPointerDown.bind(this));
    this.canvas.addEventListener('pointerup', this.onPointerUp.bind(this));
  }

  onPointerMove(event) {
    const rect = this.canvas.getBoundingClientRect();

    this.mouse.x = ((event.clientX - rect.left) / rect.width) * 2 - 1;
    this.mouse.y = -((event.clientY - rect.top) / rect.height) * 2 + 1;

    this.raycaster.setFromCamera(this.mouse, this.camera);

    if (!this.robot) return;
    const intersects = this.raycaster.intersectObjects(this.robot.children, true);

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

      let jointName = this.draggedJoint.name;
      const { lower, upper } = getJointLimits(jointName, this.robot);

      const currentAngle = this.jointAngles[jointName] ?? 0;
      const proposedAngle = currentAngle + angleDelta;
      const clampedAngle = Math.max(lower, Math.min(upper, proposedAngle));
      const actualDelta = clampedAngle - currentAngle;

      // ✅ Rotate dragged joint
      this.draggedJoint.rotateOnAxis(this.worldAxis, actualDelta);
      this.jointAngles[jointName] = clampedAngle;
      const jointTransmission = getJointTransmission(jointName, this.robot);
      this.onRotateMotor(jointTransmission.motorID, radiansToTicks(clampedAngle, jointTransmission).toFixed(0));

      // ✅ Apply mimic relationships automatically
      for (const [name, jointObj] of Object.entries(this.robot.joints)) {
        if (jointObj.type === 'URDFMimicJoint') {
          const { mimicJoint, multiplier, offset } = jointObj;
          if (mimicJoint === jointName) {
            const mimicAngle = clampedAngle * multiplier + offset;
            const { lower, upper } = getJointLimits(name, this.robot);
            const mimicClamped = Math.max(lower, Math.min(upper, mimicAngle));
            const mimicDelta = mimicClamped - (this.jointAngles[name] ?? 0);

            const mimicNode = this.robot.getObjectByName(name);
            mimicNode.rotateOnAxis(this.worldAxis, mimicDelta);
            this.jointAngles[name] = mimicClamped;
          }
        }
      }

      // ✅ Update gizmo indicator
      const gizmo = this.draggedJoint.userData.gizmo;
      if (gizmo?.indicator) this.draggedJoint.parent.remove(gizmo.indicator);
      const newIndicator = createAngleIndicator(this.worldAxis, clampedAngle);
      newIndicator.position.copy(this.draggedJoint.position);
      this.draggedJoint.parent.add(newIndicator);
      this.draggedJoint.userData.gizmo.indicator = newIndicator;
    }



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

      // reset previous hover if it's not selected
      if (this.hoveredJoint && this.hoveredJoint !== this.selectedJoint) {
        if (this.hoveredJoint?.material?.emissive) {
          this.hoveredJoint.material.emissive.setHex(0x000000);
        }
      }

      this.hoveredJoint = target;

      if (this.hoveredJoint.material?.emissive) {
        // use a different color if it's also selected
        const color = (this.hoveredJoint === this.selectedJoint) ? 0x555555 : 0x333333;
        this.hoveredJoint.material.emissive.setHex(color);
      }
    } else {
      if (this.hoveredJoint && this.hoveredJoint !== this.selectedJoint) {
        if (this.hoveredJoint?.material?.emissive) {
          this.hoveredJoint.material.emissive.setHex(0x000000);
        }
      }
      this.hoveredJoint = null;
    }

  }

  setSyncMode(newMode) {
    console.log("changing from " + this.currentSyncMode + " to " + newMode);
    this.currentSyncMode = newMode;
    if (this.currentSyncMode == SyncMode.RealToSim) {
      this.serial.requestPositionStreaming(true);
    } else {
      this.serial.requestPositionStreaming(false);

    }
  }

  findMotorByID(motorID) {
    for (const jointName in this.robot.joints) {
      const joint = this.robot.joints[jointName];
      const transmission = joint.transmission;
      if (!transmission) continue;

      if (transmission.motorID === motorID) {
        return transmission;
      }
    }
  }

  findAllMotorIDs() {
    let list = [];
    for (const jointName in this.robot.joints) {
      const joint = this.robot.joints[jointName];
      const transmission = joint.transmission;
      if (!transmission) continue;
      list.push(transmission.motorID);
    }
    return list;
  }

  getMotorTicks(motorID) {
    if (!this.robot?.joints) return null;
    for (const [jointName, jointData] of Object.entries(this.robot.joints)) {
      const t = jointData.transmission;
      if (!t) continue;
      //console.log(t);
      if (t.motorID === motorID) {

        const angle = this.jointAngles?.[jointName] ?? 0;
        return parseInt(radiansToTicks(angle, t));
      }
    }

    return null; // not found
  }

  syncWithCurveEditor() {
    let allIDs = this.findAllMotorIDs();

    let ids = [];
    let positions = [];
    allIDs.forEach(id => {
      let pos = this.curveEditor.getMotorPositionAtTime(id, this.curveEditor.currentTime);
      if (this.currentSyncMode !== SyncMode.RealToSim) {
        this.setMotorPosition(id, pos);
      }

      ids.push(id);
      positions.push(pos);
    });

    if (this.currentSyncMode === SyncMode.SimToReal) {
      this.sendMotorPosition(ids, positions);
    }
  }

  setMotorPosition(motorID, positionTicks) {
    for (const jointName in this.robot.joints) {
      const joint = this.robot.joints[jointName];
      const transmission = joint.transmission;
      if (!transmission) continue;

      if (transmission.motorID === motorID) {
        // 1) Convert ticks → target joint angle (radians)
        const targetAngle = ticksToRadians(positionTicks, transmission);

        // 2) Compute delta vs tracked state
        const currentAngle = this.jointAngles[jointName] ?? 0;
        const delta = targetAngle - currentAngle;

        // 3) Rotate driven joint
        joint.rotateOnAxis(joint.axis, delta);
        this.jointAngles[jointName] = targetAngle;

        // 4) ✅ Apply mimic relationships
        for (const [name, jointObj] of Object.entries(this.robot.joints)) {
          if (jointObj.type === 'URDFMimicJoint') {
            const { mimicJoint, multiplier = 1.0, offset = 0.0 } = jointObj;
            if (mimicJoint === jointName) {
              const mimicAngle = targetAngle * multiplier + offset;
              const mimicCurrent = this.jointAngles[name] ?? 0;
              const mimicDelta = mimicAngle - mimicCurrent;

              const mimicNode = this.robot.getObjectByName(name);
              mimicNode.rotateOnAxis(mimicNode.axis, mimicDelta);
              this.jointAngles[name] = mimicAngle;
            }
          }
        }

        break; // found and updated the matching joint
      }
    }
  }

  onRotateMotor(motorID, positionTicks) {
    if (this.currentSyncMode === SyncMode.SimToReal) {
      this.sendMotorPosition(motorID, positionTicks);
    }
  }

  onPointerDown(event) {
    if (!this.hoveredJoint) return;

    if (this.hoveredJoint) {
      // reset previous selection
      if (this.selectedJoint) {
        this.selectedJoint.material.emissive.setHex(0x000000);
      }
      this.selectedJoint = this.hoveredJoint;
      this.selectedJoint.material.emissive.setHex(0x555555); // selection color
    }

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

    // Decide how far up the chain to go
    let levelsUp = 0; // default: climb one level from link → joint
    if (event.ctrlKey) levelsUp = 1;
    if (event.altKey) levelsUp = 2;

    this.draggedJoint = this.findJointAncestor(this.hoveredJoint, levelsUp);
    if (!this.draggedJoint) return;

    // ✅ Skip mimic joints
    if (this.draggedJoint.type === 'URDFMimicJoint') {
      this.draggedJoint = findJointAncestor(this.draggedJoint, 1);
    }


    const jointName = this.draggedJoint.name;
    const jointData = this.robot.joints?.[jointName];
    if (!(jointData?.axis instanceof THREE.Vector3)) return;

    const axis = jointData.axis.clone();
    if (axis.lengthSq() === 0) return;
    this.worldAxis = axis.normalize();

    const { lower, upper } = getJointLimits(jointName, this.robot);
    const transmission = jointData.transmission;
    const encoderRange = transmission
      ? transmission.encoderRange / transmission.mechanicalReduction
      : Math.PI;

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

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

    this.draggedJoint.parent.add(sector);
    this.draggedJoint.parent.add(indicator);
    this.draggedJoint.userData.gizmo = { sector, indicator };
  }



  onPointerUp() {
    this.isDragging = false;
    this.controls.enabled = true;

    const gizmo = this.draggedJoint?.userData.gizmo;
    if (gizmo) {
      this.draggedJoint.parent.remove(gizmo.sector);
      this.draggedJoint.parent.remove(gizmo.indicator);
    }

    this.draggedJoint = null;
    this.worldAxis = null;
  }

  findJointAncestor(joint, levelsUp = 0) {
    let current = joint;
    let steps = 0;

    while (current && steps <= levelsUp) {
      current = current.parent;
      // climb until we find a URDFJoint (not a mimic)
      while (current && !(current.type === 'URDFJoint')) {
        current = current.parent;
      }
      steps++;
    }

    return current;
  }

  applyCalibrationOffsets() {
    console.log(this.jointAngles);
    if (!this.robot?.urdfRobotNode || !this.jointAngles) return;

    for (const jointName in this.jointAngles) {
      const offset = this.jointAngles[jointName]; // radians
      const jointNode = this.robot.urdfRobotNode.querySelector(`joint[name="${jointName}"]`);
      if (!jointNode) continue;

      // get or create <origin>
      let originNode = jointNode.querySelector("origin");
      let [rx, ry, rz] = originNode?.getAttribute("rpy")?.split(" ").map(parseFloat) || [0, 0, 0];

      // determine axis of rotation
      const axisNode = jointNode.querySelector("axis");
      const [ax, ay, az] = axisNode ? axisNode.getAttribute("xyz").split(" ").map(parseFloat) : [0, 0, 1];

      // apply offset along the correct axis
      if (Math.abs(ax) === 1) rx += offset * ax;
      else if (Math.abs(ay) === 1) ry += offset * ay;
      else if (Math.abs(az) === 1) rz += offset * az;

      // update the URDF DOM
      console.log(jointName, [rx, ry, rz]);
      this.updateJointOriginRpy(jointName, [rx, ry, rz]);
    }
  }


  updateJointOriginRpy(jointName, rpyArray) {
    //console.log("Updating URDF joint rpy:", jointName, rpyArray);
    const urdfNode = this.robot?.urdfRobotNode;
    if (!urdfNode) return;

    // find joint anywhere in the URDF
    const jointNode = urdfNode.querySelector(`joint[name="${jointName}"]`);
    if (!jointNode) return;

    // find or create the <origin> child
    let originNode = jointNode.querySelector("origin");
    if (!originNode) {
      originNode = urdfNode.ownerDocument.createElement("origin");
      jointNode.appendChild(originNode);
    }

    // ensure we have 3 numbers [rx, ry, rz]
    const [rx, ry, rz] = rpyArray;
    originNode.setAttribute("rpy", `${rx} ${ry} ${rz}`);
  }



  animate() {
    requestAnimationFrame(() => this.animate());
    //console.log(this.robot.joints)
    this.renderer.render(this.scene, this.camera);
    this.overlay.draw();
    //console.log(this.mesh);
  }
}

function ticksToRadians(ticks, actuator) {
  const ticksPerDeg = actuator.encoderTicks / actuator.encoderRange; // ≈ 22.75
  const centered = ticks - actuator.encoderTicks / 2; // shift so mid = 0
  const actuatorDeg = centered / ticksPerDeg;
  const jointDeg = actuatorDeg / actuator.mechanicalReduction;
  return jointDeg * (Math.PI / 180); // convert degrees → radians
}

function radiansToTicks(radians, actuator) {
  const ticksPerDeg = actuator.encoderTicks / actuator.encoderRange; // ≈ 22.75
  const jointDeg = radians * (180 / Math.PI); // radians → degrees
  const actuatorDeg = jointDeg * actuator.mechanicalReduction; // apply reduction
  const centered = actuatorDeg * ticksPerDeg; // convert to ticks offset from mid
  const ticks = centered + actuator.encoderTicks / 2; // shift back to full range
  return ticks;
}


function getJointLimits(jointName, robot) {
  const transmission = robot.joints[jointName].transmission;

  if (!transmission) {
    return { lower: -Math.PI, upper: Math.PI };
  }
  //console.log(transmission.encoderValidMin, transmission.encoderValidMax);
  const lowerRad = ticksToRadians(transmission.encoderValidMin, transmission);
  const upperRad = ticksToRadians(transmission.encoderValidMax, transmission);
  //console.log(lowerRad, upperRad);
  return { lower: lowerRad, upper: upperRad };
}

function getJointTransmission(jointName, robot) {
  return robot.joints[jointName].transmission;
}

function loadObjWithMtl(objPath, mtlPath, onLoad) {
  const mtlLoader = new MTLLoader();
  mtlLoader.load(mtlPath, (materials) => {
    materials.preload();
    const objLoader = new OBJLoader();
    objLoader.setMaterials(materials);
    objLoader.load(objPath, (object) => {
      onLoad(object);
    });
  });
}

function addOriginMarkers(robot) {
  console.log(robot.joints);
  for (const [name, joint] of Object.entries(robot.joints)) {
    const marker = new THREE.Mesh(
      new THREE.SphereGeometry(0.01),
      new THREE.MeshBasicMaterial({ color: 0x00ff00 })
    );
    joint.add(marker);

    const axes = new THREE.AxesHelper(0.05);
    joint.add(axes);

    console.log("Added marker for joint:", name);
  }

}