import * as THREE from "three";
import { OrbitControls } from "three/addons/controls/OrbitControls.js";
import Stats from "three/addons/libs/stats.module.js";
import { GUI } from "lil-gui";
import { 
  constructGrid, 
  SparkControls, 
  SparkRenderer, 
  SplatMesh, 
  textSplats, 
  dyno, 
  transcodeSpz, 
  isMobile, 
  isPcSogs, 
  LN_SCALE_MIN, 
  LN_SCALE_MAX 
} from "@sparkjsdev/spark";

const scene = new THREE.Scene();
const camera = new THREE.PerspectiveCamera(75, window.innerWidth / window.innerHeight, 0.01, 1000);
camera.position.set(0, 0, 1);

const canvas = document.getElementById("canvas");
const renderer = new THREE.WebGLRenderer({ canvas, antialias: false });
const spark = new SparkRenderer({ renderer });
scene.add(spark);

function handleResize() {
  const width = canvas.clientWidth;
  const height = canvas.clientHeight;
  renderer.setSize(width, height, false);
  camera.aspect = width / height;
  camera.updateProjectionMatrix();
}

handleResize();
window.addEventListener("resize", handleResize);

const frame = new THREE.Group();
frame.quaternion.set(1, 0, 0, 0);
scene.add(frame);

const inputs = [];

const grid = new SplatMesh({
  constructSplats: (splats) => constructGrid({
    splats,
    extents: new THREE.Box3(new THREE.Vector3(-10, -10, -10), new THREE.Vector3(10, 10, 10)),
  }),
});
grid.opacity = 0;
scene.add(grid);

// Initializing Stats correctly
const stats = new Stats();
document.body.appendChild(stats.dom);
stats.dom.style.display = "none";

const gui = new GUI({
  title: "Settings",
  container: document.getElementById("main-gui")
});
const secondGui = new GUI({
  title: "Splats",
  container: document.getElementById("second-gui")
}).close();

const controls = new SparkControls({ canvas });

const orbitControls = new OrbitControls(camera, renderer.domElement);
orbitControls.enabled = false;
orbitControls.target.set(0, 0, 0);
orbitControls.minDistance = 0.1;
orbitControls.maxDistance = 10;

const fileInput = document.querySelector("#file-input");
fileInput.onchange = (event) => {
  const allFiles = [...event.target.files];
  const validFiles = allFiles.filter(f => /\.(ply|spz|splat|ksplat|zip|sog)$/i.test(f.name));
  
  if (validFiles.length !== allFiles.length) {
    const invalidCount = allFiles.length - validFiles.length;
    alert(`${invalidCount} file(s) skipped. Only .ply, .spz, .splat, .ksplat, .zip, and .sog files are supported.`);
  }
  
  if (validFiles.length > 0) {
    loadFiles(validFiles);
  }
};

const guiOptions = {
  highDevicePixel: !isMobile(),
  stats: false,
  resetOnLoad: true,
  loadOffset: 0,
  openCv: true,
  autoRotate: false,
  orbit: false,
  cameraLocked: false,
  reversePointerDir: false,
  reversePointerSlide: false,
  backgroundColor: "#000000",
  viewBoundingBox: false,
  openFiles: () => { fileInput.click(); },
  loadFromText: "",
  loadFromTextAction: () => {
    if (guiOptions.loadFromText.trim()) {
      const urls = parseURLsFromText(guiOptions.loadFromText);
      if (urls.length > 0) {
        loadFiles(urls);
        guiOptions.loadFromText = ""; 
      } else {
        alert("No valid URLs found.");
      }
    }
  },
  resetPose: () => {
    camera.position.set(0, 0, 1);
    camera.quaternion.set(0, 0, 0, 1);
    camera.fov = 75;
    resetFrameQuaternion();
    camera.updateProjectionMatrix();
  },
};

function resetFrameQuaternion() {
  if (guiOptions.openCv) {
    frame.quaternion.set(1, 0, 0, 0);
  } else {
    frame.quaternion.set(0, 0, 0, 1);
  }
}

function parseURLsFromQuery() {
  const urlParams = new URLSearchParams(window.location.search);
  const urls = [];
  for (const [key, value] of urlParams.entries()) {
    if (key === "url") urls.push(value);
  }
  return urls;
}

guiOptions.resetPose();

const cameraFolder = gui.addFolder("Camera");
const cameraPose = cameraFolder.addFolder("Camera Pose");
cameraPose.add(camera.position, "x", -10, 10, 0.01).name("X").listen();
cameraPose.add(camera.position, "y", -10, 10, 0.01).name("Y").listen();
cameraPose.add(camera.position, "z", -10, 10, 0.01).name("Z").listen();
const rotX = cameraPose.add(camera.rotation, "x", -Math.PI, Math.PI, 0.01).name("RotateX").listen();
const rotY = cameraPose.add(camera.rotation, "y", -Math.PI, Math.PI, 0.01).name("RotateY").listen();
const rotZ = cameraPose.add(camera.rotation, "z", -Math.PI, Math.PI, 0.01).name("RotateZ").listen();
cameraPose.add(camera, "fov", 1, 179, 1).name("Fov Y").listen().onChange(() => camera.updateProjectionMatrix());
cameraPose.close();

const progressBar = document.getElementById('progress-bar');
const progressFill = document.getElementById('progress-fill');

function updateProgress(progress) {
  progressFill.style.width = `${Math.min(100, Math.max(0, progress * 100))}%`;
}

async function fetchWithProgress(url) {
  const response = await fetch(url, { mode: "cors" });
  if (!response.ok) throw new Error(`HTTP error! status: ${response.status}`);
  const total = response.headers.get('content-length') ? parseInt(response.headers.get('content-length'), 10) : null;
  const reader = response.body.getReader();
  const chunks = [];
  let loadedBytes = 0;
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value);
    loadedBytes += value.length;
    updateProgress(total ? (loadedBytes / total) : (loadedBytes / (loadedBytes + 10 * 1024 * 1024)));
  }
  const allChunks = new Uint8Array(loadedBytes);
  let offset = 0;
  for (const chunk of chunks) {
    allChunks.set(chunk, offset);
    offset += chunk.length;
  }
  return allChunks.buffer;
}

async function loadFiles(splatFiles) {
  if (guiOptions.resetOnLoad) {
    const toRemove = frame.children.filter((child) => child instanceof SplatMesh || child instanceof THREE.Box3Helper);
    for (const child of toRemove) frame.remove(child);
    inputs.length = 0;
    splatsFolder.foldersRecursive().forEach((child) => child.destroy());
  }
  
  const hasUrls = splatFiles.some(file => typeof file === "string");
  if (hasUrls) progressBar.style.display = "block";

  let index = 0;
  for (const splatFile of splatFiles) {
    try {
      let fileBytes, fileName, url = null;
      if (typeof splatFile === "string") {
        fileBytes = new Uint8Array(await fetchWithProgress(splatFile));
        fileName = splatFile.split("/").pop().split("?")[0] || "downloaded";
        if (isPcSogs(fileBytes)) url = splatFile;
      } else {
        fileBytes = new Uint8Array(await splatFile.arrayBuffer());
        fileName = splatFile.name;
      }

      const init = url ? { url } : { fileBytes: fileBytes.slice(), fileName };
      const splatMesh = new SplatMesh(init);
      const translate = guiOptions.loadOffset * index;
      splatMesh.position.set(translate, 0.5 * translate, 0.1 * translate);
      splatMesh.enableWorldToView = true;
      splatMesh.worldModifier = makeWorldModifier(splatMesh);
      await splatMesh.initialized;

      if (!url) inputs.push({ fileBytes, pathOrUrl: fileName, object: splatMesh });
      frame.add(splatMesh);
      addBoundingBoxHelper(splatMesh);

      const splatFolder = splatsFolder.addFolder(fileName).close();
      splatFolder.add(splatMesh, "opacity", 0, 1, 0.01).name("Opacity").listen();
      splatFolder.add(splatMesh.position, "x", -10, 10, 0.01).name("X").listen();
      splatFolder.add(splatMesh.scale, "x", 0.01, 4, 0.01).name("Scale").listen().onChange((v) => splatMesh.scale.setScalar(v));
    } catch (error) {
      console.error("Error loading splat:", error);
    }
    index++;
  }
  progressBar.style.display = "none";
  canvas.focus();
}

async function addBoundingBoxHelper(splatMesh) {
  await splatMesh.initialized;
  const box = splatMesh.getBoundingBox();
  const boxHelper = new THREE.Box3Helper(box, 0x00ff00);
  boxHelper.visible = guiOptions.viewBoundingBox;
  frame.add(boxHelper);
}

secondGui.add(guiOptions, "resetOnLoad").name("Reset on load");
secondGui.add(guiOptions, "loadOffset", -2, 2, 0.01).name("Offset");
secondGui.add(guiOptions, "openFiles").name("Select Files");
secondGui.add(guiOptions, "loadFromText").name("Paste URL");
secondGui.add(guiOptions, "loadFromTextAction").name("Load URL");

cameraFolder.add(guiOptions, "resetPose").name("Reset pose");
cameraFolder.add(guiOptions, "autoRotate").name("Auto rotate").listen();
cameraFolder.add(guiOptions, "cameraLocked").name("Lock Camera").onChange((v) => {
  if (guiOptions.orbit) orbitControls.enabled = !v;
});
// cameraFolder.add(guiOptions, "orbit").name("Orbit controls").listen().onChange((v) => {
//   orbitControls.enabled = v;
//   canvas.focus();
//   rotX.enable(!v); rotY.enable(!v); rotZ.enable(!v);
// });

gui.add(guiOptions, "stats").name("Stats").onChange((v) => stats.dom.style.display = v ? "block" : "none");
gui.addColor(guiOptions, "backgroundColor").name("Background").onChange((v) => scene.background.set(v));

const splatsFolder = secondGui.addFolder("Files");
const clipFolder = gui.addFolder("Clip Splats").close();

function updateFrameSplats() {
  frame.children.forEach((child) => { if (child instanceof SplatMesh) child.updateVersion(); });
}

const clipEnable = dyno.dynoBool(false);
const clipMinX = dyno.dynoFloat(-5), clipMaxX = dyno.dynoFloat(5);
clipFolder.add(clipEnable, "value").name("Enable clip").onChange(updateFrameSplats);
clipFolder.add(clipMinX, "value", -50, 50, 0.01).name("Min X").onChange(updateFrameSplats);
clipFolder.add(clipMaxX, "value", -50, 50, 0.01).name("Max X").onChange(updateFrameSplats);

// Deformation uniforms
const dragPoint = dyno.dynoVec3(new THREE.Vector3(0, 0, 0));
const dragDisplacement = dyno.dynoVec3(new THREE.Vector3(0, 0, 0));
const dragRadius = dyno.dynoFloat(isMobile() ? 0.05 : 0.25);
const dragActive = dyno.dynoFloat(0.0);
const bounceTime = dyno.dynoFloat(0.0);
const bounceBaseDisplacement = dyno.dynoVec3(new THREE.Vector3(0, 0, 0));
const dragIntensity = dyno.dynoFloat(5.0);
const bounceAmount = dyno.dynoFloat(0.5);
const bounceSpeed = dyno.dynoFloat(isMobile() ? 0.15 : 0.35);
const mode1Decay = dyno.dynoFloat(2.0);
const mode2Decay = dyno.dynoFloat(4.0);
const mode2FreqMult = dyno.dynoFloat(2.5);
const modeMix = dyno.dynoFloat(0.8);

let isBouncing = false;
let isDragging = false;
let dragStartPoint = null;
let currentDragPoint = null;
let dragStartNDC = null;
let dragScale = 1.0;

const raycaster = new THREE.Raycaster();
raycaster.params.Points = { threshold: 0.5 };

const deformFolder = gui.addFolder("Deformation").close();
deformFolder.add(dragIntensity, "value", 0, 10.0, 0.1).name("Strength").onChange(updateFrameSplats);
deformFolder.add(dragRadius, "value", 0.01, 1.0, 0.01).name("Radius").onChange(updateFrameSplats);
deformFolder.add(bounceAmount, "value", 0, 1.0, 0.1).name("Bounce Strength").onChange(updateFrameSplats);
deformFolder.add(bounceSpeed, "value", 0, 1.0, 0.01).name("Bounce Speed").onChange(updateFrameSplats);
deformFolder.add(mode1Decay, "value", 0.1, 10.0, 0.1).name("Primary Decay").onChange(updateFrameSplats);
deformFolder.add(mode2Decay, "value", 0.1, 10.0, 0.1).name("Secondary Decay").onChange(updateFrameSplats);
deformFolder.add(mode2FreqMult, "value", 1.0, 10.0, 0.1).name("Secondary Freq").onChange(updateFrameSplats);
deformFolder.add(modeMix, "value", 0.0, 1.0, 0.01).name("Mode Mix").onChange(updateFrameSplats);

function makeWorldModifier(mesh) {
  const context = mesh.context;
  return dyno.dynoBlock({ gsplat: dyno.Gsplat }, { gsplat: dyno.Gsplat }, ({ gsplat }) => {
    // 1. Define Deformation Shader
    const deformShader = new dyno.Dyno({
      inTypes: {
        gsplat: dyno.Gsplat,
        dragPoint: "vec3",
        dragDisplacement: "vec3",
        dragRadius: "float",
        dragActive: "float",
        bounceTime: "float",
        bounceBaseDisplacement: "vec3",
        dragIntensity: "float",
        bounceAmount: "float",
        bounceSpeed: "float",
        mode1Decay: "float",
        mode2Decay: "float",
        mode2FreqMult: "float",
        modeMix: "float",
      },
      outTypes: { gsplat: dyno.Gsplat },
      statements: ({ inputs, outputs }) => dyno.unindentLines(`
        ${outputs.gsplat} = ${inputs.gsplat};
        vec3 originalPos = ${inputs.gsplat}.center;
        
        // Calculate influence based on distance from drag point
        float distToDrag = distance(originalPos, ${inputs.dragPoint});
        
        // IMPROVEMENT: Gaussian Falloff for organic soft-body feel
        // Replaces linear smoothstep with a bell curve
        float sigma = max(0.001, ${inputs.dragRadius});
        float dragInfluence = exp(-(distToDrag * distToDrag) / (2.0 * sigma * sigma));
        
        // Optimization: Cutoff influence at 3 sigma to prevent infinite tails
        if (distToDrag > sigma * 3.0) dragInfluence = 0.0;

        float time = ${inputs.bounceTime};

        // Apply drag deformation
        if (${inputs.dragActive} > 0.5 && ${inputs.dragRadius} > 0.0) {
          vec3 dragOffset = ${inputs.dragDisplacement} * dragInfluence * ${inputs.dragIntensity} * 50.0;
          originalPos += dragOffset;
        }
        
        // Apply elastic bounce effect
        float bounceFrequency = 1.0 + ${inputs.bounceSpeed} * 8.0;
        vec3 bounceOffset = ${inputs.bounceBaseDisplacement} * dragInfluence * ${inputs.dragIntensity} * 50.0;
        
        // IMPROVEMENT: Multi-frequency Oscillation (Harmonics)
        // Primary low-frequency mode (bulk motion)
        float mode1 = cos(time * bounceFrequency) * exp(-time * ${inputs.mode1Decay} * (1.0 - ${inputs.bounceAmount} * 0.9));
        
        // Secondary high-frequency mode (surface wobble/jiggle)
        // Higher frequency (2.5x) and faster decay
        float mode2 = cos(time * bounceFrequency * ${inputs.mode2FreqMult}) * exp(-time * ${inputs.mode2Decay} * (1.0 - ${inputs.bounceAmount} * 0.9));
        
        // Blend the modes (80% bulk, 20% wobble)
        float oscillation = ${inputs.modeMix} * mode1 + (1.0 - ${inputs.modeMix}) * mode2;

        originalPos += bounceOffset * oscillation;

        ${outputs.gsplat}.center = originalPos;
      `),
    });

    // 2. Apply Deformation
    const deformedGsplat = deformShader.apply({
      gsplat,
      dragPoint,
      dragDisplacement,
      dragRadius,
      dragActive,
      bounceTime,
      bounceBaseDisplacement,
      dragIntensity,
      bounceAmount,
      bounceSpeed,
      mode1Decay,
      mode2Decay,
      mode2FreqMult,
      modeMix,
    }).gsplat;

    // 3. Apply Clipping (on the deformed splat)
    let { rgb, center, opacity } = dyno.splitGsplat(deformedGsplat).outputs;
    const { x } = dyno.split(center).outputs;
    const within = dyno.and(dyno.greaterThanEqual(x, clipMinX), dyno.lessThanEqual(x, clipMaxX));
    opacity = dyno.select(dyno.or(dyno.not(clipEnable), within), opacity, dyno.dynoConst("float", 0));
    
    return { gsplat: dyno.combineGsplat({ gsplat: deformedGsplat, rgb, opacity }) };
  });
}

function makeInstructions() {
  const instr = textSplats({ text: "Drag and Drop\na ply or splat file\nhere to view", textAlign: "center", fontSize: 64, objectScale: 0.1 / 64 });
  instr.quaternion.set(1, 0, 0, 0);
  instr.enableWorldToView = true;
  instr.worldModifier = makeWorldModifier(instr);
  instr.updateGenerator();
  return instr;
}

const instructions = makeInstructions();
frame.add(instructions);

const urlsToLoad = parseURLsFromQuery();
if (urlsToLoad.length > 0) loadFiles(urlsToLoad);

// Convert mouse coordinates to normalized device coordinates
function getMouseNDC(event) {
  const rect = renderer.domElement.getBoundingClientRect();
  return new THREE.Vector2(
    ((event.clientX - rect.left) / rect.width) * 2 - 1,
    -((event.clientY - rect.top) / rect.height) * 2 + 1,
  );
}

// Raycast to find intersection point on splat
function getHitPoint(ndc) {
  raycaster.setFromCamera(ndc, camera);
  // Check all splat meshes in the frame
  const splatMeshes = frame.children.filter(child => child instanceof SplatMesh);
  if (splatMeshes.length === 0) return null;
  
  const hits = raycaster.intersectObjects(splatMeshes, false);
  if (hits && hits.length > 0) {
    return hits[0].point.clone();
  }
  return null;
}

renderer.domElement.addEventListener("pointerdown", (event) => {
  // Only allow deformation when camera is locked
  if (!guiOptions.cameraLocked) return;
  // Only start drag if we hit a splat
  const ndc = getMouseNDC(event);
  const hitPoint = getHitPoint(ndc);

  if (hitPoint) {
    // Disable orbit controls while dragging
    if (guiOptions.orbit) orbitControls.enabled = false;
    
    isDragging = true;
    dragStartNDC = ndc.clone();
    dragStartPoint = hitPoint.clone();
    currentDragPoint = hitPoint.clone();

    // Calculate scale factor for screen-to-world conversion
    const distanceToCamera = camera.position.distanceTo(hitPoint);
    const fov = camera.fov * (Math.PI / 180);
    const screenHeight = 2.0 * Math.tan(fov / 2.0) * distanceToCamera;
    dragScale = screenHeight / window.innerHeight;

    dragPoint.value.copy(hitPoint);
    dragActive.value = 1.0;
    dragDisplacement.value.set(0, 0, 0);

    bounceTime.value = -1.0;
    bounceBaseDisplacement.value.set(0, 0, 0);
    isBouncing = false;
    updateFrameSplats();
  }
});

renderer.domElement.addEventListener("pointermove", (event) => {
  if (!isDragging || !dragStartPoint || !dragStartNDC) return;

  const ndc = getMouseNDC(event);

  // Convert screen space movement to world space
  const mouseDelta = new THREE.Vector2(
    (ndc.x - dragStartNDC.x) * dragScale,
    (ndc.y - dragStartNDC.y) * dragScale,
  );

  const cameraRight = new THREE.Vector3();
  const cameraUp = new THREE.Vector3();
  camera.getWorldDirection(new THREE.Vector3());
  cameraRight.setFromMatrixColumn(camera.matrixWorld, 0).normalize();
  cameraUp.setFromMatrixColumn(camera.matrixWorld, 1).normalize();

  const worldDisplacement = new THREE.Vector3()
    .addScaledVector(cameraRight, mouseDelta.x)
    .addScaledVector(cameraUp, mouseDelta.y);

  currentDragPoint = dragStartPoint.clone().add(worldDisplacement);
  dragDisplacement.value.copy(worldDisplacement);
  updateFrameSplats();
});

renderer.domElement.addEventListener("pointerup", (event) => {
  if (!isDragging) return;

  isDragging = false;
  // Re-enable orbit controls if they were enabled
  if (guiOptions.orbit && !guiOptions.cameraLocked) orbitControls.enabled = true;

  // Start bounce animation with final displacement
  if (currentDragPoint && dragStartPoint) {
    bounceBaseDisplacement.value.copy(dragDisplacement.value);
    bounceTime.value = 0.0;
    isBouncing = true;
  }

  dragActive.value = 0.0;
  dragDisplacement.value.set(0, 0, 0);
  dragStartNDC = null;
  updateFrameSplats();
});

canvas.addEventListener("dragover", (e) => { e.preventDefault(); canvas.style.opacity = "0.5"; });
canvas.addEventListener("dragleave", (e) => { e.preventDefault(); canvas.style.opacity = "1"; });
canvas.addEventListener("drop", (e) => {
  e.preventDefault(); canvas.style.opacity = "1";
  const files = Array.from(e.dataTransfer.files).filter(f => /\.(ply|spz|splat|ksplat|zip|sog)$/i.test(f.name));
  if (files.length > 0) loadFiles(files);
});

function parseURLsFromText(text) {
  const supported = [".ply", ".spz", ".splat", ".ksplat", ".zip", ".json"];
  return text.trim().split(/[\r\n,;]+/).map(p => p.trim()).filter(p => (p.startsWith('http') && supported.some(ext => p.toLowerCase().includes(ext))));
}

let lastTime = null;
renderer.setAnimationLoop(function animate(time) {
  const deltaTime = time - (lastTime || time);
  lastTime = time;
  stats.begin();
  
  if (isBouncing) {
    bounceTime.value += 0.1;
    updateFrameSplats();
  }

  if (guiOptions.autoRotate) frame.rotation.y += deltaTime / 5000;
  if (!guiOptions.cameraLocked) {
    if (guiOptions.orbit) orbitControls.update(); else controls.update(camera);
  }
  renderer.render(scene, camera);
  stats.end();
});