import * as THREE from 'three';
import * as GaussianSplats3D from '@mkkellogg/gaussian-splats-3d';

const ASSET_ROOT =
  'https://huggingface.co/datasets/mikafi/gaussiansplats3D/resolve/main/demo/assets';

const scenes = [
  {
    id: 'garden',
    title: 'Garden Walkthrough',
    shortName: 'Garden',
    mood: 'Outdoor capture',
    description:
      'Start in a dense outdoor capture, then walk through the scene instead of orbiting around it.',
    files: {
      standard: `${ASSET_ROOT}/data/garden/garden.ksplat`,
      ultra: `${ASSET_ROOT}/data/garden/garden_high.ksplat`
    },
    sizes: { standard: '75 MB', ultra: '181 MB' },
    thumb: `${ASSET_ROOT}/images/garden.png`,
    cameraUp: [0, -1, -0.54],
    start: [-3.15634, -0.16946, -0.51552],
    lookAt: [1.52976, 2.27776, 1.65898],
    walkSpeed: 1.45,
    sprintMultiplier: 2.4,
    liftSpeed: 0.85,
    tourRadius: 4.4,
    tourHeight: 0.35
  },
  {
    id: 'bonsai',
    title: 'Bonsai Atrium',
    shortName: 'Bonsai',
    mood: 'Indoor detail',
    description:
      'A compact capture with fine structure and view-dependent highlights for close inspection.',
    files: {
      standard: `${ASSET_ROOT}/data/bonsai/bonsai.ksplat`,
      ultra: `${ASSET_ROOT}/data/bonsai/bonsai_high.ksplat`
    },
    sizes: { standard: '21 MB', ultra: '46 MB' },
    thumb: `${ASSET_ROOT}/images/bonsai.png`,
    cameraUp: [0.01933, -0.7583, -0.65161],
    start: [1.54163, 2.68515, -6.37228],
    lookAt: [0.45622, 1.95338, 1.51278],
    walkSpeed: 1.2,
    sprintMultiplier: 2.2,
    liftSpeed: 0.75,
    tourRadius: 3.2,
    tourHeight: 0.25
  },
  {
    id: 'truck',
    title: 'Truck Yard',
    shortName: 'Truck',
    mood: 'Large object scan',
    description:
      'Circle a high-fidelity vehicle capture with a grounded first-person camera and cinematic tour mode.',
    files: {
      standard: `${ASSET_ROOT}/data/truck/truck.ksplat`,
      ultra: `${ASSET_ROOT}/data/truck/truck_high.ksplat`
    },
    sizes: { standard: '29 MB', ultra: '75 MB' },
    thumb: `${ASSET_ROOT}/images/truck.png`,
    cameraUp: [0, -1, -0.17],
    start: [-5, -1, -1],
    lookAt: [-1.72477, 0.05395, -0.00147],
    walkSpeed: 1.35,
    sprintMultiplier: 2.45,
    liftSpeed: 0.8,
    tourRadius: 4.1,
    tourHeight: 0.2
  },
  {
    id: 'stump',
    title: 'Forest Stump',
    shortName: 'Stump',
    mood: 'Organic texture',
    description:
      'Explore bark, foliage, and soft photogrammetry detail from a close-range natural capture.',
    files: {
      standard: `${ASSET_ROOT}/data/stump/stump.ksplat`,
      ultra: `${ASSET_ROOT}/data/stump/stump_high.ksplat`
    },
    sizes: { standard: '74 MB', ultra: '181 MB' },
    thumb: `${ASSET_ROOT}/images/stump.png`,
    cameraUp: [0, -1, -1],
    start: [-3.3816, 1.96931, -1.7189],
    lookAt: [-0.04979, 1.37519, 1.13443],
    walkSpeed: 1.1,
    sprintMultiplier: 2.35,
    liftSpeed: 0.75,
    tourRadius: 3.7,
    tourHeight: 0.2
  }
];

const dom = {
  app: document.querySelector('#app'),
  viewport: document.querySelector('#viewport'),
  sceneTitle: document.querySelector('#sceneTitle'),
  sceneDescription: document.querySelector('#sceneDescription'),
  sceneGrid: document.querySelector('#sceneGrid'),
  qualityLabel: document.querySelector('#qualityLabel'),
  loadStatus: document.querySelector('#loadStatus'),
  walkButton: document.querySelector('#walkButton'),
  tourButton: document.querySelector('#tourButton'),
  resetButton: document.querySelector('#resetButton'),
  loadingOverlay: document.querySelector('#loadingOverlay'),
  loaderText: document.querySelector('#loaderText'),
  toast: document.querySelector('#toast'),
  fps: document.querySelector('#fps'),
  speedLabel: document.querySelector('#speedLabel')
};

const params = new URLSearchParams(window.location.search);
let currentScene =
  scenes.find((scene) => scene.id === params.get('scene')) ??
  scenes.find((scene) => scene.id === 'truck') ??
  scenes[0];
let quality = params.get('quality') === 'ultra' ? 'ultra' : 'standard';

let renderer;
let camera;
let viewer;
let rafId = 0;
let lastTime = performance.now();
let frameBucketStart = lastTime;
let frameCount = 0;
let isWalking = false;
let tourEnabled = true;
let tourStartTime = 0;
let yaw = 0;
let pitch = 0;
let basis;

const keys = new Set();
const cameraState = {
  position: new THREE.Vector3(),
  direction: new THREE.Vector3(),
  height: 0
};

function vectorFromArray(value) {
  return new THREE.Vector3(value[0], value[1], value[2]);
}

function makeBasis(scene) {
  const up = vectorFromArray(scene.cameraUp).normalize();
  const start = vectorFromArray(scene.start);
  const target = vectorFromArray(scene.lookAt);
  const initialForward = target.clone().sub(start).normalize();
  let flatForward = initialForward.clone().sub(up.clone().multiplyScalar(initialForward.dot(up)));

  if (flatForward.lengthSq() < 0.0001) {
    flatForward = new THREE.Vector3(0, 0, -1).sub(up.clone().multiplyScalar(up.z)).normalize();
  } else {
    flatForward.normalize();
  }

  const right = new THREE.Vector3().crossVectors(flatForward, up).normalize();
  const correctedForward = new THREE.Vector3().crossVectors(up, right).normalize();
  const initialPitch = Math.asin(THREE.MathUtils.clamp(initialForward.dot(up), -0.92, 0.92));

  return {
    up,
    right,
    forward: correctedForward,
    focus: target,
    start,
    initialPitch
  };
}

function setCameraFromYawPitch() {
  const yawQuat = new THREE.Quaternion().setFromAxisAngle(basis.up, yaw);
  const yawedForward = basis.forward.clone().applyQuaternion(yawQuat).normalize();
  const yawedRight = basis.right.clone().applyQuaternion(yawQuat).normalize();
  const pitchQuat = new THREE.Quaternion().setFromAxisAngle(yawedRight, pitch);
  const direction = yawedForward.applyQuaternion(pitchQuat).normalize();

  cameraState.direction.copy(direction);
  camera.up.copy(basis.up);
  camera.position.copy(cameraState.position);
  camera.lookAt(cameraState.position.clone().add(direction));
}

function resetCamera() {
  yaw = 0;
  pitch = basis.initialPitch;
  cameraState.position.copy(basis.start);
  cameraState.height = cameraState.position.dot(basis.up);
  setCameraFromYawPitch();
}

function stopTour() {
  tourEnabled = false;
  dom.tourButton.classList.remove('active');
  dom.tourButton.innerHTML = '<span class="button-icon">◎</span> Tour';
}

function startTour() {
  tourEnabled = true;
  tourStartTime = performance.now();
  dom.tourButton.classList.add('active');
  dom.tourButton.innerHTML = '<span class="button-icon">Ⅱ</span> Tour';
}

function updateTour(time) {
  const t = Math.max(0, time - tourStartTime) * 0.00009;
  const radius = currentScene.tourRadius;
  const side = basis.right.clone().multiplyScalar(Math.cos(t) * radius);
  const forward = basis.forward.clone().multiplyScalar(Math.sin(t) * radius);
  const lift = basis.up.clone().multiplyScalar(currentScene.tourHeight + Math.sin(t * 1.7) * 0.24);
  const position = basis.focus.clone().add(side).add(forward).add(lift);

  cameraState.position.copy(position);
  cameraState.height = cameraState.position.dot(basis.up);
  camera.up.copy(basis.up);
  camera.position.copy(position);
  camera.lookAt(basis.focus.clone().add(basis.up.clone().multiplyScalar(0.12)));
}

function updateMovement(delta) {
  if (!isWalking) return;

  const sprint = keys.has('ShiftLeft') || keys.has('ShiftRight');
  const speed = currentScene.walkSpeed * (sprint ? currentScene.sprintMultiplier : 1);
  const movement = new THREE.Vector3();
  const flatForward = cameraState.direction
    .clone()
    .sub(basis.up.clone().multiplyScalar(cameraState.direction.dot(basis.up)));

  if (flatForward.lengthSq() < 0.0001) {
    flatForward.copy(basis.forward);
  } else {
    flatForward.normalize();
  }

  const right = new THREE.Vector3().crossVectors(flatForward, basis.up).normalize();

  if (keys.has('KeyW') || keys.has('ArrowUp')) movement.add(flatForward);
  if (keys.has('KeyS') || keys.has('ArrowDown')) movement.sub(flatForward);
  if (keys.has('KeyD') || keys.has('ArrowRight')) movement.add(right);
  if (keys.has('KeyA') || keys.has('ArrowLeft')) movement.sub(right);
  if (keys.has('Space')) movement.add(basis.up);
  if (keys.has('KeyC') || keys.has('ControlLeft') || keys.has('ControlRight')) movement.sub(basis.up);

  if (movement.lengthSq() > 0) {
    movement.normalize().multiplyScalar(speed * delta);
    cameraState.position.add(movement);
  }

  const vertical = cameraState.position.dot(basis.up);
  const minHeight = cameraState.height - 2.2;
  const maxHeight = cameraState.height + 2.2;
  if (vertical < minHeight || vertical > maxHeight) {
    cameraState.position.add(
      basis.up.clone().multiplyScalar(THREE.MathUtils.clamp(vertical, minHeight, maxHeight) - vertical)
    );
  }

  dom.speedLabel.textContent = sprint ? 'Sprint' : 'Walk';
  setCameraFromYawPitch();
}

function animate(time) {
  rafId = requestAnimationFrame(animate);
  const delta = Math.min(0.05, (time - lastTime) / 1000);
  lastTime = time;

  if (tourEnabled && !isWalking) updateTour(time);
  updateMovement(delta);

  viewer?.update();
  viewer?.render();

  frameCount += 1;
  if (time - frameBucketStart >= 650) {
    const fps = Math.round((frameCount * 1000) / (time - frameBucketStart));
    dom.fps.textContent = fps >= 5 ? `${fps} fps` : 'Live';
    frameBucketStart = time;
    frameCount = 0;
  }
}

function showToast(message) {
  dom.toast.textContent = message;
  dom.toast.classList.add('visible');
  window.clearTimeout(showToast.timeoutId);
  showToast.timeoutId = window.setTimeout(() => {
    dom.toast.classList.remove('visible');
  }, 2600);
}

function setLoading(isLoading, message) {
  dom.loaderText.textContent = message;
  dom.loadStatus.textContent = isLoading ? 'Loading scene' : 'Ready';
  dom.loadingOverlay.classList.toggle('hidden', !isLoading);
}

function updateUrl() {
  const next = new URL(window.location.href);
  next.searchParams.set('scene', currentScene.id);
  next.searchParams.set('quality', quality);
  window.history.replaceState({}, '', next);
}

function renderSceneButtons() {
  dom.sceneGrid.innerHTML = '';
  for (const scene of scenes) {
    const button = document.createElement('button');
    button.type = 'button';
    button.className = `scene-option${scene.id === currentScene.id ? ' active' : ''}`;
    button.style.setProperty('--thumb', `url("${scene.thumb}")`);
    button.innerHTML = `
      <span class="scene-name">
        <span>${scene.shortName}</span>
        <small>${scene.mood}</small>
      </span>
    `;
    button.addEventListener('click', () => {
      if (scene.id === currentScene.id) return;
      currentScene = scene;
      updateUrl();
      loadScene();
    });
    dom.sceneGrid.appendChild(button);
  }
}

function updateText() {
  dom.sceneTitle.textContent = currentScene.title;
  dom.sceneDescription.textContent = currentScene.description;
  dom.qualityLabel.textContent = quality === 'ultra' ? `Ultra ${currentScene.sizes.ultra}` : `Standard ${currentScene.sizes.standard}`;

  for (const segment of document.querySelectorAll('.segment')) {
    segment.classList.toggle('active', segment.dataset.quality === quality);
  }
}

function disposeCurrentViewer() {
  if (rafId) cancelAnimationFrame(rafId);
  rafId = 0;
  try {
    viewer?.dispose?.();
  } catch (error) {
    console.warn('Viewer disposal warning:', error);
  }
  viewer = undefined;
  renderer?.dispose?.();
  dom.viewport.replaceChildren();
}

async function loadScene() {
  setLoading(true, `Loading ${currentScene.shortName} in ${quality} quality (${currentScene.sizes[quality]}).`);
  disposeCurrentViewer();
  renderSceneButtons();
  updateText();
  basis = makeBasis(currentScene);

  renderer = new THREE.WebGLRenderer({
    antialias: false,
    alpha: false,
    powerPreference: 'high-performance'
  });
  renderer.setClearColor(0x020303, 1);
  renderer.setPixelRatio(Math.min(window.devicePixelRatio, 1.75));
  renderer.setSize(window.innerWidth, window.innerHeight);
  renderer.outputColorSpace = THREE.SRGBColorSpace;
  dom.viewport.appendChild(renderer.domElement);

  camera = new THREE.PerspectiveCamera(64, window.innerWidth / window.innerHeight, 0.05, 500);
  camera.up.copy(basis.up);

  resetCamera();

  const canUseSharedMemory = typeof SharedArrayBuffer !== 'undefined' && window.crossOriginIsolated;
  viewer = new GaussianSplats3D.Viewer({
    selfDrivenMode: false,
    renderer,
    camera,
    useBuiltInControls: false,
    ignoreDevicePixelRatio: false,
    gpuAcceleratedSort: canUseSharedMemory,
    sharedMemoryForWorkers: canUseSharedMemory,
    integerBasedSort: true,
    halfPrecisionCovariancesOnGPU: true,
    dynamicScene: false,
    renderMode: GaussianSplats3D.RenderMode.Always,
    sceneRevealMode: GaussianSplats3D.SceneRevealMode.Gradual,
    sceneFadeInRateMultiplier: 1.7,
    antialiased: false,
    sphericalHarmonicsDegree: 2,
    inMemoryCompressionLevel: 2,
    freeIntermediateSplatData: true,
    logLevel: GaussianSplats3D.LogLevel.None
  });

  try {
    await viewer.addSplatScene(currentScene.files[quality], {
      progressiveLoad: false,
      showLoadingUI: false,
      splatAlphaRemovalThreshold: 5
    });
    setLoading(false, 'Ready');
    startTour();
    lastTime = performance.now();
    frameBucketStart = lastTime;
    frameCount = 0;
    rafId = requestAnimationFrame(animate);
    showToast('Click Walk, then use WASD and mouse to explore.');
  } catch (error) {
    console.error(error);
    setLoading(false, 'Unable to load the splat.');
    dom.loadStatus.textContent = 'Load failed';
    showToast('The scene failed to load. Try Standard quality or refresh.');
  }
}

function resize() {
  if (!renderer || !camera) return;
  camera.aspect = window.innerWidth / window.innerHeight;
  camera.updateProjectionMatrix();
  renderer.setPixelRatio(Math.min(window.devicePixelRatio, 1.75));
  renderer.setSize(window.innerWidth, window.innerHeight);
}

function enterWalkMode() {
  if (!renderer) return;
  stopTour();
  renderer.domElement.requestPointerLock?.();
}

document.addEventListener('pointerlockchange', () => {
  isWalking = document.pointerLockElement === renderer?.domElement;
  dom.app.classList.toggle('is-walking', isWalking);
  dom.walkButton.innerHTML = isWalking
    ? '<span class="button-icon">×</span> Walking'
    : '<span class="button-icon">↗</span> Walk';
  if (!isWalking) keys.clear();
});

document.addEventListener('mousemove', (event) => {
  if (!isWalking) return;
  yaw -= event.movementX * 0.0021;
  pitch = THREE.MathUtils.clamp(pitch - event.movementY * 0.0021, -1.28, 1.28);
});

document.addEventListener('keydown', (event) => {
  if (
    [
      'KeyW',
      'KeyA',
      'KeyS',
      'KeyD',
      'ArrowUp',
      'ArrowDown',
      'ArrowLeft',
      'ArrowRight',
      'Space',
      'ShiftLeft',
      'ShiftRight',
      'KeyC',
      'ControlLeft',
      'ControlRight'
    ].includes(event.code)
  ) {
    keys.add(event.code);
    if (isWalking) event.preventDefault();
  }

  if (event.code === 'Escape') keys.clear();
});

document.addEventListener('keyup', (event) => {
  keys.delete(event.code);
});

dom.walkButton.addEventListener('click', () => {
  if (isWalking) {
    document.exitPointerLock?.();
  } else {
    enterWalkMode();
  }
});

dom.tourButton.addEventListener('click', () => {
  if (isWalking) document.exitPointerLock?.();
  if (tourEnabled) {
    stopTour();
  } else {
    startTour();
  }
});

dom.resetButton.addEventListener('click', () => {
  if (isWalking) document.exitPointerLock?.();
  resetCamera();
  stopTour();
  showToast('View reset.');
});

for (const segment of document.querySelectorAll('.segment')) {
  segment.addEventListener('click', () => {
    const nextQuality = segment.dataset.quality;
    if (nextQuality === quality) return;
    quality = nextQuality;
    updateUrl();
    loadScene();
  });
}

window.addEventListener('resize', resize);
window.addEventListener('blur', () => keys.clear());

renderSceneButtons();
updateText();
loadScene();