skymap-gen/src/main.js

import './style.css';
import * as THREE from 'three';
import { OrbitControls } from 'three/examples/jsm/controls/OrbitControls.js';
import { convertFileSrc, invoke } from '@tauri-apps/api/tauri';
import { listen } from '@tauri-apps/api/event';

const app = document.getElementById('app');
app.innerHTML = `
  <div id="canvas-container"></div>
  <div id="thumb-dock">
    <div id="thumb-header">Maps</div>
    <div id="thumb-list"></div>
  </div>
  <div id="settings-panel" class="hidden">
    <div class="settings-header">
      <div>
        <div class="settings-title">Generation Settings</div>
        <div class="settings-sub">Defaults tuned for SDXL 360</div>
      </div>
      <button id="settings-close" aria-label="Close settings">×</button>
    </div>
    <div class="settings-grid">
      <label>Steps
        <input type="number" id="steps-input" min="1" max="80" />
      </label>
      <label>Guidance
        <input type="number" id="guidance-input" step="0.1" min="0" max="20" />
      </label>
      <label>Width
        <input type="number" id="width-input" min="256" max="4096" step="64" />
      </label>
      <label>Height
        <input type="number" id="height-input" min="256" max="4096" step="64" />
      </label>
      <label>Scheduler
        <select id="scheduler-input">
          <option value="dpmsolver-sde">DPM++ 2M SDE</option>
          <option value="dpmsolver">DPM++ 2M</option>
          <option value="euler">Euler</option>
          <option value="euler_a">Euler Ancestral</option>
          <option value="heun">Heun</option>
          <option value="ddim">DDIM</option>
        </select>
      </label>
      <label>Upscale
        <select id="upscale-input">
          <option value="none">None</option>
          <option value="realesrgan">Real-ESRGAN</option>
          <option value="topaz">Topaz</option>
        </select>
      </label>
      <label class="settings-check">Seam Inpainting
        <input type="checkbox" id="seam-inpaint-input" />
      </label>
      <label>Model Path
        <input type="text" id="model-path-input" />
      </label>
      <label>Base Model
        <input type="text" id="base-model-input" />
      </label>
      <label>VAE
        <input type="text" id="vae-model-input" />
      </label>
    </div>
    <div class="settings-footer">
      <button id="settings-reset">Reset defaults</button>
    </div>
  </div>
  <div id="progress-overlay">
    <div id="progress-ring">
      <div id="progress-fill"></div>
      <div id="progress-text">...</div>
    </div>
  </div>
  <div id="prompt-bar">
    <input id="prompt-input" type="text" placeholder="Describe the environment..." />
    <button id="settings-btn" title="Generation settings">Settings</button>
    <button id="generate-btn">Generate</button>
  </div>
  <div id="status"></div>
  <div id="delete-confirm" class="hidden" role="dialog" aria-modal="true" aria-labelledby="delete-confirm-title">
    <div id="delete-confirm-card">
      <div id="delete-confirm-title">Delete map?</div>
      <div id="delete-confirm-body">This will remove <span id="delete-confirm-filename"></span> from disk.</div>
      <div id="delete-confirm-actions">
        <button id="delete-confirm-cancel" type="button">Cancel</button>
        <button id="delete-confirm-delete" type="button">Delete</button>
      </div>
    </div>
  </div>
`;

const canvasContainer = document.getElementById('canvas-container');
const promptInput = document.getElementById('prompt-input');
const generateBtn = document.getElementById('generate-btn');
const settingsBtn = document.getElementById('settings-btn');
const statusEl = document.getElementById('status');
const thumbList = document.getElementById('thumb-list');
const thumbDock = document.getElementById('thumb-dock');
const deleteConfirm = document.getElementById('delete-confirm');
const deleteConfirmFilename = document.getElementById('delete-confirm-filename');
const deleteConfirmCancel = document.getElementById('delete-confirm-cancel');
const deleteConfirmDelete = document.getElementById('delete-confirm-delete');
const progressOverlay = document.getElementById('progress-overlay');
const progressFill = document.getElementById('progress-fill');
const progressText = document.getElementById('progress-text');
const settingsPanel = document.getElementById('settings-panel');
const settingsClose = document.getElementById('settings-close');
const settingsReset = document.getElementById('settings-reset');
const stepsInput = document.getElementById('steps-input');
const guidanceInput = document.getElementById('guidance-input');
const widthInput = document.getElementById('width-input');
const heightInput = document.getElementById('height-input');
const schedulerInput = document.getElementById('scheduler-input');
const upscaleInput = document.getElementById('upscale-input');
const seamInpaintInput = document.getElementById('seam-inpaint-input');
const modelPathInput = document.getElementById('model-path-input');
const baseModelInput = document.getElementById('base-model-input');
const vaeModelInput = document.getElementById('vae-model-input');

// Three.js setup
let renderer, scene, camera, controls, skyMesh;
const defaultTextureUrl = '/default.png';
const skyFadeDurationMs = 1000;
let skyTransition = null;
let skyTextureRequestId = 0;
let autoSpin = true;
let lastInteraction = Date.now();
const minZoomDistance = 0.2;
const maxZoomDistance = 20;
let desiredDistance = 1;
const zoomLerpFactor = 0.1;
const zoomVec = new THREE.Vector3();
let progressState = {
  upscale: null,
  seamInpaint: null,
  phases: {
    gen: null,
    decode: null,
    inpaint: null,
    upscale: null,
  },
};
let progressUnlisten = null;
let currentMapPath = null;
let deleteConfirmResolve = null;
let generationRunning = false;
let cancelRequested = false;
const defaultSettings = {
  steps: 25,
  width: 1536,
  height: 768,
  guidance: 6.5,
  scheduler: 'dpmsolver-sde',
  upscale: 'none',
  seamInpaint: false,
  modelPath: 'proximasan/sdxl-360-diffusion',
  baseModel: 'stabilityai/stable-diffusion-xl-base-1.0',
  vaeModel: 'madebyollin/sdxl-vae-fp16-fix',
};
let currentSettings = { ...defaultSettings };

function setStatus(msg) {
  statusEl.textContent = msg || '';
  statusEl.style.opacity = msg ? '0.98' : '0';
  if (msg) {
    clearTimeout(setStatus._timer);
    setStatus._timer = setTimeout(() => {
      statusEl.style.opacity = '0';
    }, 5000);
  }
}

async function loadTexture(url) {
  const loader = new THREE.TextureLoader();
  return new Promise((resolve, reject) => {
    loader.load(
      url,
      tex => {
        tex.mapping = THREE.EquirectangularReflectionMapping;
        tex.colorSpace = THREE.SRGBColorSpace;
        resolve(tex);
      },
      undefined,
      err => reject(err)
    );
  });
}

async function setSkyFromUrl(url) {
  const requestId = ++skyTextureRequestId;
  try {
    const texture = await loadTexture(url);
    if (requestId !== skyTextureRequestId) {
      texture.dispose();
      return;
    }
    transitionSkyToTexture(texture);
  } catch (e) {
    console.error('Failed to load texture', e);
    if (requestId === skyTextureRequestId) {
      setStatus('Failed to load map texture');
    }
  }
}

function replaceSkyTexture(texture) {
  const previousTexture = skyMesh.material.map;
  skyMesh.material.map = texture;
  skyMesh.material.needsUpdate = true;

  if (previousTexture && previousTexture !== texture) {
    previousTexture.dispose();
  }
}

function removeSkyTransition(transition) {
  scene.remove(transition.mesh);
  transition.material.dispose();
}

function finishSkyTransition(transition) {
  if (!transition || skyTransition !== transition) return;
  removeSkyTransition(transition);
  skyTransition = null;
  replaceSkyTexture(transition.texture);
}

function transitionSkyToTexture(texture) {
  if (!skyMesh.material.map) {
    replaceSkyTexture(texture);
    return;
  }

  if (skyTransition) {
    finishSkyTransition(skyTransition);
  }

  const material = new THREE.MeshBasicMaterial({
    side: THREE.BackSide,
    color: 0xffffff,
    map: texture,
    transparent: true,
    opacity: 0,
    depthWrite: false,
  });
  const mesh = new THREE.Mesh(skyMesh.geometry, material);
  mesh.rotation.copy(skyMesh.rotation);
  mesh.renderOrder = skyMesh.renderOrder + 1;
  scene.add(mesh);

  skyTransition = {
    mesh,
    material,
    texture,
    start: performance.now(),
    duration: skyFadeDurationMs,
  };
}

function updateSkyTransition(now) {
  if (!skyTransition) return;

  const progress = Math.min(1, (now - skyTransition.start) / skyTransition.duration);
  skyTransition.material.opacity = THREE.MathUtils.smoothstep(progress, 0, 1);
  skyTransition.mesh.rotation.copy(skyMesh.rotation);

  if (progress >= 1) {
    finishSkyTransition(skyTransition);
  }
}

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

  renderer = new THREE.WebGLRenderer({ antialias: true });
  renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
  renderer.setSize(window.innerWidth, window.innerHeight);
  renderer.outputColorSpace = THREE.SRGBColorSpace;
  renderer.toneMapping = THREE.ACESFilmicToneMapping;
  canvasContainer.appendChild(renderer.domElement);

  desiredDistance = camera.position.length();

  const geometry = new THREE.SphereGeometry(50, 64, 32);
  const material = new THREE.MeshBasicMaterial({ side: THREE.BackSide, color: 0xffffff });
  skyMesh = new THREE.Mesh(geometry, material);
  scene.add(skyMesh);

  controls = new OrbitControls(camera, renderer.domElement);
  controls.enablePan = false;
  controls.enableZoom = false; // we implement custom smooth zooming
  controls.enableDamping = true;
  controls.dampingFactor = 0.08;
  controls.rotateSpeed = 0.3;
  controls.autoRotate = false;
  controls.autoRotateSpeed = 0.15;
  controls.minDistance = 0.05;
  controls.maxDistance = 20;

  controls.addEventListener('start', () => {
    autoSpin = false;
    lastInteraction = Date.now();
  });
  controls.addEventListener('end', () => {
    lastInteraction = Date.now();
  });

  window.addEventListener('resize', onResize);
  renderer.domElement.addEventListener(
    'wheel',
    (e) => {
      e.preventDefault();
      const delta = e.deltaY;
      // Scale desired distance exponentially for smooth feel
      const factor = Math.exp(delta * 0.0015);
      desiredDistance = THREE.MathUtils.clamp(desiredDistance * factor, minZoomDistance, maxZoomDistance);
      lastInteraction = Date.now();
      autoSpin = false;
    },
    { passive: false }
  );
  animate();
}

function onResize() {
  if (!renderer || !camera) return;
  const w = window.innerWidth;
  const h = window.innerHeight;
  renderer.setSize(w, h);
  camera.aspect = w / h;
  camera.updateProjectionMatrix();
}

function animate() {
  requestAnimationFrame(animate);
  const now = performance.now();
  const idle = Date.now() - lastInteraction > 2000;
  if (idle) autoSpin = true;
  if (autoSpin) {
    skyMesh.rotation.y += 0.0008;
  }
  updateSkyTransition(now);
  // Smooth zoom toward desired distance
  const target = controls.target;
  const currentDistance = camera.position.distanceTo(target);
  const nextDistance = THREE.MathUtils.lerp(currentDistance, desiredDistance, zoomLerpFactor);
  zoomVec.copy(camera.position).sub(target).normalize().multiplyScalar(nextDistance).add(target);
  camera.position.copy(zoomVec);

  controls.update();
  renderer.render(scene, camera);
}

async function pathToObjectUrl(path) {
  return convertFileSrc(path);
}

function resetProgressState() {
  progressState = {
    upscale: null,
    seamInpaint: null,
    phases: {
      gen: null,
      decode: null,
      inpaint: null,
      upscale: null,
    },
  };
}

function computeProgress() {
  const upscaleOn = progressState.upscale === true;
  const seamOn = progressState.seamInpaint === true;
  const weights = {
    gen: seamOn || upscaleOn ? 0.45 : 0.9,
    decode: seamOn || upscaleOn ? 0.05 : 0.1,
    inpaint: seamOn ? 0.5 : 0,
    upscale: upscaleOn ? 0.5 : 0,
  };
  if (seamOn && upscaleOn) {
    weights.gen = 0.25;
    weights.decode = 0.05;
    weights.inpaint = 0.25;
    weights.upscale = 0.45;
  }
  const frac = phase => {
    if (!phase || !phase.total) return 0;
    const pct = phase.current / phase.total;
    return Math.max(0, Math.min(1, pct));
  };
  return (
    frac(progressState.phases.gen) * weights.gen +
    frac(progressState.phases.decode) * weights.decode +
    frac(progressState.phases.inpaint) * weights.inpaint +
    frac(progressState.phases.upscale) * weights.upscale
  );
}

function updateProgressDisplay(pct) {
  if (!progressOverlay) return;
  const clamped = Math.max(0, Math.min(1, pct || 0));
  const deg = clamped * 100;
  progressFill.style.background = `conic-gradient(#ffffff ${deg}%, rgba(255,255,255,0.12) ${deg}%)`;
  progressText.textContent = `${Math.round(clamped * 100)}%`;
}

function startProgress() {
  resetProgressState();
  if (progressOverlay) {
    progressOverlay.style.display = 'flex';
  }
  updateProgressDisplay(0);
}

function stopProgress() {
  updateProgressDisplay(1);
  if (progressOverlay) {
    progressOverlay.style.display = 'none';
  }
}

async function refreshThumbnails(selectedPath, shouldLoadSelection = true) {
  let maps;
  try {
    maps = await invoke('list_maps');
  } catch (e) {
    console.error(e);
    setStatus('Failed to read map list');
    thumbDock.style.display = 'none';
    return;
  }
  thumbList.innerHTML = '';
  if (!maps || maps.length === 0) {
    thumbDock.style.display = 'none';
    return;
  }
  thumbDock.style.display = 'block';

  for (let idx = 0; idx < maps.length; idx++) {
    const item = maps[idx];
    const el = document.createElement('div');
    el.className = 'thumb-item';
    const img = document.createElement('img');
    const deleteBtn = document.createElement('button');
    deleteBtn.className = 'thumb-delete';
    deleteBtn.type = 'button';
    deleteBtn.textContent = '×';
    deleteBtn.title = 'Delete map';
    deleteBtn.setAttribute('aria-label', `Delete ${item.filename}`);
    deleteBtn.addEventListener('click', async (event) => {
      event.preventDefault();
      event.stopPropagation();
      await deleteMap(item, el);
    });
    let fileUrl;
    try {
      fileUrl = await pathToObjectUrl(item.path);
      img.src = fileUrl;
    } catch (err) {
      console.error('Failed to load thumbnail', err);
      img.alt = 'Failed to load';
    }
    el.appendChild(img);
    el.appendChild(deleteBtn);
    el.title = item.filename;
    el.addEventListener('click', () => {
      if (fileUrl) {
        currentMapPath = item.path;
        setSkyFromUrl(fileUrl);
        setStatus(`Loaded ${item.filename}`);
      }
    });
    thumbList.appendChild(el);

    if (fileUrl) {
      if (idx === 0 && !selectedPath) {
        currentMapPath = item.path;
        setSkyFromUrl(fileUrl);
        setStatus(`Showing ${item.filename}`);
      }
      if (selectedPath && selectedPath === item.path) {
        currentMapPath = item.path;
        setSkyFromUrl(fileUrl);
        setStatus(`Showing ${item.filename}`);
      }
    }
  }
}

async function deleteMap(item, tileEl) {
  if (!item?.path) return;
  const confirmed = await requestDeleteConfirmation(item.filename);
  if (!confirmed) return;

  try {
    await invoke('delete_map', { path: item.path });
    const wasCurrent = currentMapPath === item.path;
    tileEl?.remove();
    setStatus(`Deleted ${item.filename}`);

    if (thumbList.children.length === 0) {
      thumbDock.style.display = 'none';
      currentMapPath = null;
      await setSkyFromUrl(defaultTextureUrl);
      return;
    }

    if (wasCurrent) {
      thumbList.querySelector('.thumb-item')?.click();
    }
  } catch (err) {
    console.error(err);
    setStatus('Failed to delete map');
  }
}

function requestDeleteConfirmation(filename) {
  if (deleteConfirmResolve) {
    closeDeleteConfirmation(false);
  }
  deleteConfirmFilename.textContent = `"${filename}"`;
  deleteConfirm.classList.remove('hidden');
  deleteConfirmDelete.focus();

  return new Promise(resolve => {
    deleteConfirmResolve = resolve;
  });
}

function closeDeleteConfirmation(result) {
  deleteConfirm.classList.add('hidden');
  if (deleteConfirmResolve) {
    const resolve = deleteConfirmResolve;
    deleteConfirmResolve = null;
    resolve(result);
  }
}

async function generateMap() {
  if (generationRunning) {
    await cancelGeneration();
    return;
  }
  const prompt = promptInput.value.trim();
  if (!prompt) {
    setStatus('Please enter a prompt');
    return;
  }
  const settings = {
    steps: Number(stepsInput.value) || defaultSettings.steps,
    guidance: Number(guidanceInput.value) || defaultSettings.guidance,
    width: Number(widthInput.value) || defaultSettings.width,
    height: Number(heightInput.value) || defaultSettings.height,
    scheduler: schedulerInput.value || defaultSettings.scheduler,
    upscale: upscaleInput.value || defaultSettings.upscale,
    seamInpaint: seamInpaintInput.checked,
    modelPath: modelPathInput.value.trim() || defaultSettings.modelPath,
    baseModel: baseModelInput.value.trim() || defaultSettings.baseModel,
    vaeModel: vaeModelInput.value.trim() || defaultSettings.vaeModel,
  };
  currentSettings = settings;
  generationRunning = true;
  cancelRequested = false;
  generateBtn.disabled = false;
  generateBtn.textContent = 'Cancel';
  setStatus('Generating...');
  startProgress();
  try {
    const result = await invoke('generate_map', { prompt, settings });
    const outputPath = result.outputPath || result.output_path || result;
    if (outputPath) {
      const fileUrl = await pathToObjectUrl(outputPath);
      await refreshThumbnails(outputPath, false);
      currentMapPath = outputPath;
      await setSkyFromUrl(fileUrl);
      setStatus('New environment loaded');
    } else {
      setStatus('Generation finished, but no output path reported');
    }
  } catch (e) {
    console.error(e);
    setStatus(e === 'Generation cancelled' ? 'Generation cancelled' : (typeof e === 'string' ? e : 'Generation failed'));
  } finally {
    stopProgress();
    generationRunning = false;
    cancelRequested = false;
    generateBtn.textContent = 'Generate';
    generateBtn.disabled = false;
  }
}

async function cancelGeneration() {
  if (!generationRunning || cancelRequested) return;
  cancelRequested = true;
  generateBtn.disabled = true;
  generateBtn.textContent = 'Cancelling...';
  setStatus('Cancelling generation...');

  try {
    await invoke('cancel_generation');
  } catch (e) {
    console.error(e);
    cancelRequested = false;
    generateBtn.disabled = false;
    generateBtn.textContent = 'Cancel';
    setStatus(typeof e === 'string' ? e : 'Failed to cancel generation');
  }
}

function setupUI() {
  generateBtn.addEventListener('click', () => {
    if (generationRunning) {
      cancelGeneration();
    } else {
      generateMap();
    }
  });
  promptInput.addEventListener('keydown', (e) => {
    if (e.key === 'Enter' && !generationRunning) generateMap();
  });
  settingsBtn.addEventListener('click', () => {
    settingsPanel.classList.remove('hidden');
  });
  settingsClose.addEventListener('click', () => {
    settingsPanel.classList.add('hidden');
  });
  settingsReset.addEventListener('click', () => {
    applySettings(defaultSettings);
  });
  deleteConfirmCancel.addEventListener('click', () => {
    closeDeleteConfirmation(false);
  });
  deleteConfirmDelete.addEventListener('click', () => {
    closeDeleteConfirmation(true);
  });
  deleteConfirm.addEventListener('click', (event) => {
    if (event.target === deleteConfirm) {
      closeDeleteConfirmation(false);
    }
  });
  document.addEventListener('keydown', (event) => {
    if (event.key === 'Escape' && !deleteConfirm.classList.contains('hidden')) {
      closeDeleteConfirmation(false);
    }
  });
}

async function setupProgressEvents() {
  try {
    progressUnlisten = await listen('gen-progress', (event) => {
      const data = event?.payload || {};
      if (typeof data.upscale === 'boolean') {
        progressState.upscale = data.upscale;
      }
      if (typeof data.seamInpaint === 'boolean') {
        progressState.seamInpaint = data.seamInpaint;
      }
      if (data.phase && typeof data.current === 'number' && typeof data.total === 'number') {
        progressState.phases[data.phase] = {
          current: data.current,
          total: data.total,
        };
      }
      updateProgressDisplay(computeProgress());
    });
  } catch (err) {
    console.error('Failed to bind progress listener', err);
  }
}

async function bootstrap() {
  initScene();
  setupUI();
  applySettings(defaultSettings);
  await setupProgressEvents();
  await setSkyFromUrl(defaultTextureUrl);
  await refreshThumbnails();
  setStatus('Ready');
}

function applySettings(cfg) {
  currentSettings = { ...cfg };
  stepsInput.value = cfg.steps;
  guidanceInput.value = cfg.guidance;
  widthInput.value = cfg.width;
  heightInput.value = cfg.height;
  schedulerInput.value = cfg.scheduler;
  upscaleInput.value = cfg.upscale;
  seamInpaintInput.checked = Boolean(cfg.seamInpaint);
  modelPathInput.value = cfg.modelPath;
  baseModelInput.value = cfg.baseModel;
  vaeModelInput.value = cfg.vaeModel;
}

bootstrap();