import concaveman from 'concaveman';
import gsap from "gsap";
import JSZip from "jszip";
import * as THREE from 'three';
import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
import { TransformControls } from 'three/addons/controls/TransformControls.js';
import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';
import { RGBELoader } from 'three/addons/loaders/RGBELoader.js';

const ASSET_BASE_URL = 'https://huggingface.co/spaces/aaurelions/drones/resolve/main';
const defaultPanoramas = Array.from({ length: 10 }, (_, i) => ({ name: `bg${i + 1}.jpg`, url: `${ASSET_BASE_URL}/jpg/bg${i + 1}.jpg`, thumb: `${ASSET_BASE_URL}/jpg/thumbnails/bg${i + 1}.jpg` }));
const defaultModels = ['gerbera.glb', 'shahed1.glb', 'shahed2.glb', 'shahed3.glb', 'supercam.glb', 'zala.glb', 'beaver.glb'].map(n => ({ name: n, url: `${ASSET_BASE_URL}/glb/${n}`, thumb: `${ASSET_BASE_URL}/glb/thumbnails/${n.replace('.glb', '.png')}` }));

let panoramaAssets = [...defaultPanoramas];
let modelAssets = [...defaultModels];

let scene, camera, renderer, orbitControls, transformControls, pmremGenerator, panoramaSphere;
let selectedObject = null;
let activeControlMode = null;
let activePanorama = null;

let maskScene, maskMaterial;

const loaderOverlay = document.getElementById('loader-overlay');
const loaderText = document.getElementById('loader-text');
const canvas = document.getElementById('main-canvas');

async function init() {
    loaderText.textContent = 'Setting up 3D scene...';
    scene = new THREE.Scene();
    camera = new THREE.PerspectiveCamera(70, window.innerWidth / window.innerHeight, 0.1, 2000);
    camera.position.set(0, 1.5, 8);

    renderer = new THREE.WebGLRenderer({ canvas, antialias: true, preserveDrawingBuffer: true });
    renderer.setPixelRatio(window.devicePixelRatio);
    renderer.setSize(window.innerWidth, window.innerHeight);
    renderer.toneMapping = THREE.ACESFilmicToneMapping;
    renderer.toneMappingExposure = 1.0;
    renderer.outputColorSpace = THREE.SRGBColorSpace;

    pmremGenerator = new THREE.PMREMGenerator(renderer);
    pmremGenerator.compileEquirectangularShader();

    const panoGeometry = new THREE.SphereGeometry(1000, 60, 40);
    panoGeometry.scale(-1, 1, 1);
    const panoMaterial = new THREE.MeshBasicMaterial({ color: 0xffffff });
    panoramaSphere = new THREE.Mesh(panoGeometry, panoMaterial);
    scene.add(panoramaSphere);

    orbitControls = new OrbitControls(camera, renderer.domElement);
    orbitControls.enableDamping = true;
    orbitControls.target.set(0, 1, 0);

    transformControls = new TransformControls(camera, renderer.domElement);
    transformControls.enabled = false;
    scene.add(transformControls);

    const ambientLight = new THREE.AmbientLight(0xffffff, 0.5);
    const dirLight = new THREE.DirectionalLight(0xffffff, 1.0);
    dirLight.position.set(8, 15, 10);
    scene.add(ambientLight, dirLight);

    maskScene = new THREE.Scene();
    maskMaterial = new THREE.MeshBasicMaterial({ color: 0xffffff });

    setupUI();
    setupEventListeners();
    animate();

    try {
        await loadAsset(panoramaAssets[6], 'panorama');
        await loadAsset(modelAssets[1], 'model');
        setControlMode('translate');
    } catch (error) {
        showNotification({ text: `Initialization failed: ${error.message}`, type: 'error' });
        console.error("Initialization failed:", error);
    } finally {
        gsap.to(loaderOverlay, { opacity: 0, duration: 0.5, onComplete: () => loaderOverlay.classList.add('hidden') });
    }
}

async function loadAsset(assetData, assetType, button = null) {
    const isInitialLoad = !loaderOverlay.classList.contains('hidden');
    let spinner, textSpan;
    if (button) {
        spinner = button.querySelector('.btn-spinner');
        textSpan = button.querySelector('.btn-text');
        spinner.classList.remove('hidden');
        textSpan.classList.add('hidden');
        button.disabled = true;
    } else if (!isInitialLoad) {
        loaderText.textContent = `Loading ${assetType}...`;
        loaderOverlay.classList.remove('hidden');
        gsap.to(loaderOverlay, { opacity: 1, duration: 0.3 });
        canvas.classList.add('loading');
    } else {
        loaderText.textContent = `Loading ${assetType}: ${assetData.name}`;
    }

    try {
        if (assetType === 'panorama') {
            await new Promise((resolve, reject) => {
                const isHDR = assetData.url.endsWith('.hdr');
                const loader = isHDR ? new RGBELoader() : new THREE.TextureLoader();
                loader.load(assetData.url, (texture) => {
                    texture.mapping = THREE.EquirectangularReflectionMapping;
                    if (!isHDR) texture.colorSpace = THREE.SRGBColorSpace;

                    if (panoramaSphere) {
                        panoramaSphere.material.map = texture;
                        panoramaSphere.material.needsUpdate = true;
                    }

                    scene.environment = pmremGenerator.fromEquirectangular(texture).texture;

                    scene.background = texture;
                    scene.environment = pmremGenerator.fromEquirectangular(texture).texture;
                    pmremGenerator.dispose();
                    texture.dispose();
                    activePanorama = assetData;
                    setActiveCard('panorama-gallery', assetData.name);
                    resolve();
                }, undefined, reject);
            });
        } else {
            if (selectedObject) transformControls.detach();

            const existingAsset = modelAssets.find(m => m.name === assetData.name && m.mesh);
            if (existingAsset) {
                modelAssets.forEach(m => { if (m.mesh) m.mesh.visible = false; });
                existingAsset.mesh.visible = true;
                selectedObject = existingAsset.mesh;
            } else {
                await new Promise((resolve, reject) => {
                    const loader = new GLTFLoader();
                    loader.load(assetData.url, (gltf) => {
                        modelAssets.forEach(m => { if (m.mesh) m.mesh.visible = false; });
                        const model = gltf.scene;
                        model.name = assetData.name;
                        const box = new THREE.Box3().setFromObject(model);
                        const size = box.getSize(new THREE.Vector3());
                        const center = box.getCenter(new THREE.Vector3());
                        model.position.sub(center);
                        const maxDim = Math.max(size.x, size.y, size.z);
                        model.scale.setScalar(4.0 / maxDim);
                        scene.add(model);

                        const assetToUpdate = modelAssets.find(m => m.name === assetData.name);
                        if (assetToUpdate) {
                            assetToUpdate.mesh = model;

                            assetToUpdate.originalMaterials = new Map();
                            model.traverse(node => {
                                if (node.isMesh) assetToUpdate.originalMaterials.set(node, node.material);
                            });
                        }

                        selectedObject = model;
                        resolve();
                    }, undefined, reject);
                });
            }
            if (activeControlMode) transformControls.attach(selectedObject);
            setActiveCard('model-selector', assetData.name);
        }
    } catch (error) {
        showNotification({ text: `Failed to load ${assetType}: ${assetData.name}`, type: 'error', duration: 4000 });
        console.error(`Failed to load ${assetType}:`, error);
    } finally {
        if (button) {
            spinner.classList.add('hidden');
            textSpan.classList.remove('hidden');
            button.disabled = false;
        } else if (!isInitialLoad) {
            gsap.to(loaderOverlay, { opacity: 0, duration: 0.5, onComplete: () => loaderOverlay.classList.add('hidden') });
            canvas.classList.remove('loading');
        }
    }
}

function createAssetCard(type, assetData) {
    const container = document.createElement('div');
    container.className = 'asset-card flex flex-col rounded-lg cursor-pointer overflow-hidden bg-stone-900/50';
    container.dataset.name = assetData.name;
    const thumbWrapper = document.createElement('div');
    thumbWrapper.className = 'w-full h-28 flex items-center justify-center';
    const thumb = document.createElement('img');
    thumb.className = 'max-w-full max-h-full object-contain';

    container.onclick = () => loadAsset(assetData, type);

    if (type === 'model') {
        thumbWrapper.classList.add('p-2');
        thumb.src = assetData.thumb;
        const nameWrapper = document.createElement('div');
        nameWrapper.textContent = assetData.name.replace(/\.[^/.]+$/, "").substring(0, 15);
        nameWrapper.className = 'text-white text-xs font-semibold text-center block truncate p-2 bg-black/20 mt-auto';
        thumbWrapper.appendChild(thumb);
        container.append(thumbWrapper, nameWrapper);
    } else {
        thumb.className = 'w-full h-full object-cover';
        thumb.src = assetData.thumb;
        thumbWrapper.appendChild(thumb);
        container.appendChild(thumbWrapper);
    }

    thumb.onerror = () => {
        thumbWrapper.innerHTML = `<i data-feather="package" class="w-12 h-12 text-stone-500"></i>`;
        feather.replace();
        thumbWrapper.parentElement.classList.add('bg-stone-800');
    }
    document.getElementById(type === 'model' ? 'model-selector' : 'panorama-gallery').appendChild(container);
}

function setupUI() {
    feather.replace();
    modelAssets.forEach(m => createAssetCard('model', m));
    panoramaAssets.forEach(p => createAssetCard('panorama', p));
}

function setControlMode(mode) {
    if (mode === activeControlMode) {
        activeControlMode = null;
        transformControls.detach();
        transformControls.enabled = false;
    } else {
        activeControlMode = mode;
        transformControls.setMode(activeControlMode);
        transformControls.enabled = true;
        if (selectedObject) transformControls.attach(selectedObject);
    }
    document.querySelectorAll('.control-btn').forEach(b => b.classList.toggle('bg-blue-600', b.dataset.mode === activeControlMode));
}

function setupEventListeners() {
    ['model', 'panorama'].forEach(type => {
        document.getElementById(`${type}-panel-trigger`).addEventListener('click', e => { e.stopPropagation(); document.getElementById(`${type}-panel`).classList.add('is-open'); });
        document.getElementById(`close-${type}-panel`).addEventListener('click', () => document.getElementById(`${type}-panel`).classList.remove('is-open'));
    });

    document.getElementById('day-night-toggle').addEventListener('change', (e) => {
        const isNight = e.target.checked;
        const tintColor = isNight ? 0x202040 : 0xffffff;

        if (panoramaSphere) {
            gsap.to(panoramaSphere.material.color, {
                r: (tintColor >> 16 & 255) / 255,
                g: (tintColor >> 8 & 255) / 255,
                b: (tintColor & 255) / 255,
                duration: 0.5
            });
        }

        gsap.to(renderer, { toneMappingExposure: isNight ? 0.3 : 1.0, duration: 0.5 });
        gsap.to(scene.children.find(c => c.isDirectionalLight), { intensity: isNight ? 0.2 : 1.0, duration: 0.5 });
    });

    document.querySelectorAll('.control-btn').forEach(btn => btn.addEventListener('click', () => setControlMode(btn.dataset.mode)));
    transformControls.addEventListener('dragging-changed', e => { orbitControls.enabled = !e.value; });

    canvas.addEventListener('click', (e) => {
        if (transformControls.dragging) return;
        const rect = renderer.domElement.getBoundingClientRect();
        const mouse = new THREE.Vector2(((e.clientX - rect.left) / rect.width) * 2 - 1, -((e.clientY - rect.top) / rect.height) * 2 + 1);
        const raycaster = new THREE.Raycaster();
        raycaster.setFromCamera(mouse, camera);
        const visibleMeshObjects = modelAssets.map(m => m.mesh).filter(m => m && m.visible);
        if (visibleMeshObjects.length === 0) return;
        const intersects = raycaster.intersectObjects(visibleMeshObjects, true);

        if (intersects.length > 0) {
            const object = findTopLevelGroup(intersects[0].object);
            if (object && object.isGroup && object !== selectedObject) {
                const modelData = modelAssets.find(m => m.name === object.name);
                if (modelData) loadAsset(modelData, 'model');
            }
        }
    });

    window.addEventListener('keydown', (e) => {
        if (document.activeElement.tagName === 'INPUT' || document.activeElement.tagName === 'SELECT') return;
        const key = e.key.toLowerCase();
        if (key === 'w' || key === 'e' || key === 'r') { e.preventDefault(); setControlMode(key === 'w' ? 'translate' : key === 'e' ? 'rotate' : 'scale'); }
        if (key === 'g') { e.preventDefault(); document.getElementById('generate-dataset-btn').click(); }
        if (key === 'escape') { document.querySelectorAll('.side-panel.is-open, .modal-overlay.flex').forEach(p => p.classList.add('hidden')); }
    });

    setupUploadModal();
    setupDatasetModal();
}

function setupUploadModal() {
    const modal = document.getElementById('upload-modal');
    const title = document.getElementById('upload-title');
    const fileInput = document.getElementById('upload-file-input');
    const urlInput = document.getElementById('upload-url-input');
    const dropZone = document.getElementById('upload-drop-zone');
    const loadBtn = document.getElementById('upload-load-btn');
    let currentType, currentHandler;

    const openModal = (type, titleText, accept, placeholder, handler) => {
        currentType = type;
        title.textContent = titleText;
        fileInput.value = ''; urlInput.value = '';
        fileInput.accept = accept;
        urlInput.placeholder = placeholder;
        currentHandler = handler;
        modal.classList.remove('hidden'); modal.classList.add('flex');
    };

    document.getElementById('add-model-btn').addEventListener('click', () => openModal('model', 'Upload New Model', '.glb,.gltf', 'Enter .glb URL', (type, data) => handleNewAsset(type, data, loadBtn)));
    document.getElementById('add-panorama-btn').addEventListener('click', () => openModal('panorama', 'Upload New Panorama', 'image/*,.hdr', 'Enter image URL', (type, data) => handleNewAsset(type, data, loadBtn)));
    document.getElementById('upload-cancel-btn').addEventListener('click', () => modal.classList.add('hidden'));

    loadBtn.addEventListener('click', () => {
        const url = urlInput.value.trim();
        if (url) currentHandler(currentType, url, loadBtn);
    });
    dropZone.addEventListener('click', () => fileInput.click());
    fileInput.addEventListener('change', e => { if (e.target.files.length) currentHandler(currentType, e.target.files[0], loadBtn) });
    dropZone.addEventListener('dragover', e => { e.preventDefault(); dropZone.classList.add('drag-over'); });
    dropZone.addEventListener('dragleave', () => dropZone.classList.remove('drag-over'));
    dropZone.addEventListener('drop', e => { e.preventDefault(); dropZone.classList.remove('drag-over'); if (e.dataTransfer.files.length) currentHandler(currentType, e.dataTransfer.files[0], loadBtn); });
}

async function handleNewAsset(type, fileOrUrl, button) {
    document.getElementById('upload-modal').classList.add('hidden');
    const isUrl = typeof fileOrUrl === 'string';
    const url = isUrl ? fileOrUrl : URL.createObjectURL(fileOrUrl);
    const name = (isUrl ? new URL(url).pathname.split('/').pop() : fileOrUrl.name).substring(0, 25);

    const newAssetData = { name, url, mesh: null };
    if (type === 'model') {
        if (modelAssets.some(m => m.name === name)) {
            showNotification({ text: `Model "${name}" already exists.`, type: 'error' });
            return;
        }
        newAssetData.thumb = 'placeholder';
        modelAssets.push(newAssetData);
        createAssetCard('model', newAssetData);
        await loadAsset(newAssetData, 'model', button);
    } else {
        if (panoramaAssets.some(p => p.name === name)) {
            showNotification({ text: `Panorama "${name}" already exists.`, type: 'error' });
            return;
        }
        newAssetData.thumb = url;
        panoramaAssets.push(newAssetData);
        createAssetCard('panorama', newAssetData);
        await loadAsset(newAssetData, 'panorama', button);
    }
}

function setupDatasetModal() {
    const modal = document.getElementById('dataset-modal');

    document.getElementById('generate-dataset-btn').addEventListener('click', () => {
        modal.classList.remove('hidden');
        modal.classList.add('flex');
        feather.replace();
        updateDatasetModalUI();
    });

    document.getElementById('cancel-dataset').addEventListener('click', () => modal.classList.add('hidden'));
    document.getElementById('start-dataset').addEventListener('click', handleDatasetGeneration);

    document.querySelectorAll('#dataset-modal input[type="radio"], #dataset-modal input[type="checkbox"]').forEach(input => {
        input.addEventListener('change', updateDatasetModalUI);
    });

    const setupSlider = (sliderId, displayId, unit) => {
        const slider = document.getElementById(sliderId);
        const display = document.getElementById(displayId);
        if (slider && display) {
            const update = () => { display.textContent = slider.value + unit; };
            update();
            slider.addEventListener('input', update);
        }
    };
    setupSlider('position-variance', 'position-variance-value', '%');
    setupSlider('rotation-variance', 'rotation-variance-value', '%');
    setupSlider('scale-variance', 'scale-variance-value', '%');
    setupSlider('horizontal-variance', 'horizontal-variance-value', '°');
    setupSlider('vertical-variance', 'vertical-variance-value', '°');
    setupSlider('concave-hull-concavity', 'concave-hull-concavity-value', '');
    setupSlider('background-ratio-slider', 'background-ratio-value', '%');
    setupSlider('mask-simplification-slider', 'mask-simplification-value', '');

    const trainSlider = document.getElementById('split-train-ratio');
    const valSlider = document.getElementById('split-val-ratio');

    const bringToFront = (el, topZ, bottomZ) => {
        el.style.zIndex = topZ;
        (el === trainSlider ? valSlider : trainSlider).style.zIndex = bottomZ;
    }

    trainSlider.addEventListener('input', () => { if (parseInt(trainSlider.value) >= parseInt(valSlider.value)) valSlider.value = parseInt(trainSlider.value); updateSplitRatios(); });
    valSlider.addEventListener('input', () => { if (parseInt(valSlider.value) <= parseInt(trainSlider.value)) trainSlider.value = parseInt(valSlider.value); updateSplitRatios(); });

    trainSlider.addEventListener('mousedown', () => bringToFront(trainSlider, 25, 20));
    trainSlider.addEventListener('touchstart', () => bringToFront(trainSlider, 25, 20));
    valSlider.addEventListener('mousedown', () => bringToFront(valSlider, 25, 20));
    valSlider.addEventListener('touchstart', () => bringToFront(valSlider, 25, 20));
}

function updateDatasetModalUI() {

    const useCurrentModel = document.querySelector('input[name="model-source"]:checked')?.value === 'current';
    const modelInfo = document.getElementById('model-source-current-info');
    modelInfo.textContent = `Using: ${selectedObject ? selectedObject.name : 'None'}`;
    modelInfo.classList.toggle('hidden', !useCurrentModel);

    const useCurrentPano = document.querySelector('input[name="pano-source"]:checked')?.value === 'current';
    const panoInfo = document.getElementById('pano-source-current-info');
    panoInfo.textContent = `Using: ${activePanorama ? activePanorama.name : 'None'}`;
    panoInfo.classList.toggle('hidden', !useCurrentPano);

    const randomizeModel = document.getElementById('randomize-model-toggle')?.checked;
    gsap.to("#model-sliders", { maxHeight: randomizeModel ? 200 : 0, opacity: randomizeModel ? 1 : 0, paddingTop: randomizeModel ? '0.75rem' : 0, duration: 0.4 });

    const cameraSection = document.getElementById('camera-randomization-section');
    const cameraToggle = document.getElementById('randomize-camera-toggle');
    cameraSection.style.opacity = useCurrentPano ? '1' : '0.5';
    cameraToggle.disabled = !useCurrentPano;
    if (!useCurrentPano) { cameraToggle.checked = false; }
    const cameraToggleLabel = cameraToggle.closest('label.flex.items-center.gap-2');
    if (cameraToggleLabel) cameraToggleLabel.style.cursor = useCurrentPano ? 'pointer' : 'not-allowed';

    const randomizeCamera = useCurrentPano && cameraToggle.checked;
    gsap.to("#camera-sliders", { maxHeight: randomizeCamera ? 150 : 0, opacity: randomizeCamera ? 1 : 0, paddingTop: randomizeCamera ? '0.75rem' : 0, duration: 0.4 });

    const task = document.querySelector('input[name="dataset-task"]:checked').value;
    const segGroup = document.getElementById('segmentation-method-group');
    const segMethod = document.querySelector('input[name="segmentation-method"]:checked').value;
    const concaveOptions = document.getElementById('concave-hull-options');
    const renderMaskOptions = document.getElementById('render-mask-options');

    const showSegGroup = task === 'segmentation';
    gsap.to(segGroup, {
        maxHeight: showSegGroup ? 300 : 0,
        opacity: showSegGroup ? 1 : 0,
        paddingTop: showSegGroup ? '0.5rem' : 0,
        marginTop: showSegGroup ? '1rem' : 0,
        duration: 0.4
    });

    const showConcaveOptions = showSegGroup && segMethod === 'concave';
    gsap.to(concaveOptions, {
        maxHeight: showConcaveOptions ? 100 : 0,
        opacity: showConcaveOptions ? 1 : 0,
        paddingTop: showConcaveOptions ? '0.75rem' : 0,
        duration: 0.4
    });

    const includeBackgrounds = document.getElementById('include-background-toggle').checked;
    gsap.to("#background-ratio-container", {
        maxHeight: includeBackgrounds ? 100 : 0,
        opacity: includeBackgrounds ? 1 : 0,
        paddingTop: includeBackgrounds ? '0.5rem' : 0,
        duration: 0.4
    });

    const showRenderMaskOptions = showSegGroup && segMethod === 'render';
    gsap.to(renderMaskOptions, {
        maxHeight: showRenderMaskOptions ? 100 : 0,
        opacity: showRenderMaskOptions ? 1 : 0,
        paddingTop: showRenderMaskOptions ? '0.75rem' : 0,
        duration: 0.4
    });

    updateSplitRatios();
}

function updateSplitRatios() {
    const valToggle = document.getElementById('include-val-set');
    const testToggle = document.getElementById('include-test-set');
    const trainSlider = document.getElementById('split-train-ratio');
    const valSlider = document.getElementById('split-val-ratio');

    valSlider.disabled = !valToggle.checked;
    testToggle.disabled = !valToggle.checked;
    if (!valToggle.checked) testToggle.checked = false;

    valSlider.style.visibility = testToggle.checked ? 'visible' : 'hidden';

    valSlider.parentElement.style.opacity = valToggle.checked ? '1' : '0.4';

    const trainPctRaw = parseInt(trainSlider.value);
    let valEndPctRaw = valToggle.checked ? parseInt(valSlider.value) : trainPctRaw;

    const trainPct = Math.max(0, Math.min(100, trainPctRaw));
    let valEndPct = valToggle.checked ? Math.max(trainPct, Math.min(100, valEndPctRaw)) : trainPct;

    if (!testToggle.checked) {
        valEndPct = 100;
    }

    const valPct = valToggle.checked ? valEndPct - trainPct : 0;
    const testPct = testToggle.checked ? 100 - valEndPct : 0;
    const trainDisplayPct = 100 - valPct - testPct;

    document.getElementById('split-train-display').textContent = `${trainDisplayPct}%`;
    document.getElementById('split-val-display').textContent = `${valPct}%`;
    document.getElementById('split-test-display').textContent = `${testPct}%`;

    document.getElementById('split-val-display').style.visibility = valToggle.checked ? 'visible' : 'hidden';
    document.getElementById('split-test-display').style.visibility = testToggle.checked ? 'visible' : 'hidden';

    const total = trainDisplayPct + valPct + testPct;
    document.getElementById('split-ratio-sum').textContent = total;

    document.getElementById('split-bg-train').style.width = `${trainDisplayPct}%`;
    const valBg = document.getElementById('split-bg-val');
    valBg.style.left = `${trainDisplayPct}%`;
    valBg.style.width = `${valPct}%`;
    valBg.style.display = valToggle.checked ? 'block' : 'none';

    const testBg = document.getElementById('split-bg-test');
    testBg.style.left = `${trainDisplayPct + valPct}%`;
    testBg.style.width = `${testPct}%`;
    testBg.style.display = testToggle.checked ? 'block' : 'none';
}

async function handleDatasetGeneration() {
    const startBtn = document.getElementById('start-dataset');
    const cancelBtn = document.getElementById('cancel-dataset');
    const btnText = startBtn.querySelector('.btn-text');
    const btnSpinner = startBtn.querySelector('.btn-spinner');

    btnText.classList.add('hidden');
    btnSpinner.classList.remove('hidden');
    startBtn.disabled = true;
    cancelBtn.disabled = true;

    try {
        const trainSlider = document.getElementById('split-train-ratio');
        const valSlider = document.getElementById('split-val-ratio');

        const options = {
            task: document.querySelector('input[name="dataset-task"]:checked').value,
            segmentationMethod: document.querySelector('input[name="segmentation-method"]:checked').value,
            concavity: parseFloat(document.getElementById('concave-hull-concavity').value),
            rdpEpsilon: parseFloat(document.getElementById('mask-simplification-slider').value),
            samples: parseInt(document.getElementById('dataset-samples').value),
            name: document.getElementById('dataset-name').value.trim(),

            includeBackgrounds: document.getElementById('include-background-toggle').checked,
            backgroundRatio: parseInt(document.getElementById('background-ratio-slider').value) / 100,

            useCurrentModel: document.querySelector('input[name="model-source"]:checked').value === 'current',
            randomizeModel: document.getElementById('randomize-model-toggle').checked,
            posVar: parseInt(document.getElementById('position-variance').value) / 100,
            rotVar: parseInt(document.getElementById('rotation-variance').value) / 100,
            scaleVar: parseInt(document.getElementById('scale-variance').value) / 100,
            useCurrentPanorama: document.querySelector('input[name="pano-source"]:checked').value === 'current',
            randomizeCamera: document.getElementById('randomize-camera-toggle').checked,
            hVar: THREE.MathUtils.degToRad(parseInt(document.getElementById('horizontal-variance').value)),
            vVar: THREE.MathUtils.degToRad(parseInt(document.getElementById('vertical-variance').value)),
            includeVal: document.getElementById('include-val-set').checked,
            includeTest: document.getElementById('include-test-set').checked,
            trainSplit: parseInt(trainSlider.value),
            valSplit: parseInt(valSlider.value),
            initial: {
                model: {
                    position: selectedObject ? selectedObject.position.clone() : new THREE.Vector3(),
                    rotation: selectedObject ? selectedObject.rotation.clone() : new THREE.Euler(),
                    scale: selectedObject ? selectedObject.scale.x : 1
                },
                camera: { position: camera.position.clone() }
            }
        };

        if (!options.name) { showNotification({ text: "Dataset name cannot be empty.", type: "error" }); return; }
        if ((options.task === 'detection' || options.task === 'segmentation') && !options.includeVal) {
            showNotification({ text: "Validation set is required for YOLO formats.", type: "error" });
            document.getElementById('include-val-set').checked = true;
            updateDatasetModalUI();
            return;
        }

        document.getElementById('dataset-modal').classList.add('hidden');
        const progressContainer = document.getElementById('progress-container');
        progressContainer.classList.remove('hidden'); progressContainer.classList.add('flex');

        const zip = new JSZip();
        const rootFolderName = options.name;
        const usedClasses = new Set();

        const originalModelAsset = selectedObject ? modelAssets.find(m => m.name === selectedObject.name) : null;
        const originalPanoramaAsset = activePanorama;
        const transformWasVisible = transformControls.visible;
        if (transformWasVisible) transformControls.visible = false;

        for (let i = 0; i < options.samples; i++) {
            updateProgress(i + 1, options.samples);
            const currentModelAsset = await randomizeSceneForGeneration(options, usedClasses);

            renderer.render(scene, camera);
            const imageName = `sample_${String(i).padStart(5, '0')}.jpg`;
            const labelName = imageName.replace('.jpg', '.txt');

            const modelClassName = currentModelAsset ? currentModelAsset.name.replace(/\.[^/.]+$/, "") : 'background_only';
            const classIndex = currentModelAsset ? modelAssets.findIndex(m => m.name === currentModelAsset.name) : -1;

            const { imageBlob, labelData } = await generateSampleData(options, classIndex);

            if (imageBlob) {
                const randomVal = Math.random() * 100;
                let splitFolder = 'train';

                const valPct = options.includeVal ? options.valSplit - options.trainSplit : 0;
                const testPct = options.includeTest ? 100 - options.valSplit : 0;
                const trainPct = 100 - valPct - testPct;

                if (options.includeTest && randomVal > (trainPct + valPct)) {
                    splitFolder = 'test';
                } else if (options.includeVal && randomVal > trainPct) {
                    splitFolder = 'val';
                }

                if (options.task === 'classification') {
                    const classFolder = currentModelAsset && selectedObject.visible ? modelClassName : 'background_only';
                    zip.folder(rootFolderName).folder(splitFolder).folder(classFolder).file(imageName, imageBlob);
                } else {
                    zip.folder(rootFolderName).folder('images').folder(splitFolder).file(imageName, imageBlob);
                    if (labelData) { zip.folder(rootFolderName).folder('labels').folder(splitFolder).file(labelName, labelData); }
                }
            }
            if (i % 10 === 0) await new Promise(resolve => setTimeout(resolve, 1));
        }

        if (options.task !== 'classification') {
            const classMap = {};
            // Find the original array index for each model that was actually used.
            for (const modelName of usedClasses) {
                const index = modelAssets.findIndex(m => m.name === modelName);
                if (index !== -1) {
                    // Map the index to the clean class name.
                    classMap[index] = modelName.replace(/\.[^/.]+$/, "");
                }
            }
        
            // Sort the classes by their original index to create the final list.
            const sortedKeys = Object.keys(classMap).map(Number).sort((a, b) => a - b);
            const classNames = sortedKeys.map(key => `  ${key}: ${classMap[key]}`).join('\n');
        
            // Build the final YAML content.
            let yamlContent = `path: ${rootFolderName}\ntrain: images/train\nval: images/val\n`;
            if (options.includeTest) yamlContent += `test: images/test\n`;
            yamlContent += `\nnames:\n${classNames}`;
            zip.folder(rootFolderName).file(`${options.name}.yaml`, yamlContent);
        }

        if (transformWasVisible) transformControls.visible = true;

        updateProgress(options.samples, options.samples, "Compressing ZIP...");
        const content = await zip.generateAsync({ type: "blob" });
        const link = document.createElement('a');
        link.href = URL.createObjectURL(content);
        link.download = `${rootFolderName}.zip`;
        link.click();
        URL.revokeObjectURL(link.href);
        progressContainer.classList.add('hidden');
        showNotification({ text: 'Dataset generation complete!', type: 'success' });

        await loadAsset(originalPanoramaAsset || panoramaAssets[0], 'panorama');
        await loadAsset(originalModelAsset || modelAssets[0], 'model');

    } finally {
        btnText.classList.remove('hidden');
        btnSpinner.classList.add('hidden');
        startBtn.disabled = false;
        cancelBtn.disabled = false;
    }
}

async function randomizeSceneForGeneration(options, usedClasses) {
    let currentModelAsset;
    if (!options.useCurrentPanorama) {
        const randomPano = panoramaAssets[Math.floor(Math.random() * panoramaAssets.length)];
        await loadAsset(randomPano, 'panorama');
    }
    if (!options.useCurrentModel) {
        currentModelAsset = modelAssets[Math.floor(Math.random() * modelAssets.length)];
        await loadAsset(currentModelAsset, 'model');
    } else {
        currentModelAsset = modelAssets.find(m => m.name === selectedObject.name);
    }

    if (!currentModelAsset || !selectedObject) return null;

    if (options.includeBackgrounds && Math.random() < options.backgroundRatio) {
        selectedObject.visible = false;
    } else {
        selectedObject.visible = true;
    }

    if (selectedObject.visible) {
        usedClasses.add(currentModelAsset.name);
    }

    if (!selectedObject.visible) return null;

    if (options.randomizeModel) {
        if (options.useCurrentModel) {
            const { position: basePos, rotation: baseRot, scale: baseScale } = options.initial.model;
            selectedObject.position.copy(basePos).add(new THREE.Vector3().randomDirection().multiplyScalar(Math.random() * 5));

            const baseQuaternion = new THREE.Quaternion().setFromEuler(baseRot);
            const rotOffset = new THREE.Euler((Math.random() - 0.5) * 2 * Math.PI * options.rotVar, (Math.random() - 0.5) * 2 * Math.PI * options.rotVar, (Math.random() - 0.5) * 2 * Math.PI * options.rotVar);
            baseQuaternion.multiply(new THREE.Quaternion().setFromEuler(rotOffset));
            selectedObject.quaternion.copy(baseQuaternion);

            selectedObject.scale.setScalar(baseScale * (1 + (Math.random() - 0.5) * 2 * options.scaleVar));
        } else {
            selectedObject.position.set((Math.random() - 0.5) * 10, Math.random() * 5, (Math.random() - 0.5) * 10);
            selectedObject.rotation.set(Math.random() * Math.PI * 2, Math.random() * Math.PI * 2, Math.random() * Math.PI * 2);
            const box = new THREE.Box3().setFromObject(selectedObject);
            const size = box.getSize(new THREE.Vector3());
            const maxDim = Math.max(size.x, size.y, size.z);
            const baseScale = maxDim > 0 ? 4.0 / maxDim : 1.0;
            selectedObject.scale.setScalar(baseScale * (0.5 + Math.random()));
        }
    }

    if (options.useCurrentPanorama && options.randomizeCamera) {
        const spherical = new THREE.Spherical().setFromCartesianCoords(options.initial.camera.position.x, options.initial.camera.position.y, options.initial.camera.position.z);
        spherical.theta += (Math.random() - 0.5) * 2 * options.hVar;
        spherical.phi += (Math.random() - 0.5) * 2 * options.vVar;
        spherical.phi = Math.max(0.1, Math.min(Math.PI - 0.1, spherical.phi));
        camera.position.setFromSpherical(spherical);
    } else if (!options.useCurrentPanorama) {
        camera.position.set((Math.random() - 0.5) * 20, Math.random() * 8 + 1, (Math.random() - 0.5) * 15 + 5);
    }

    // --- START: FIX FOR MODEL CENTERING ---
    const lookAtTarget = selectedObject.position.clone();

    // If Position Variance is > 0, apply a 2D-like offset to the camera's target point.
    if (options.randomizeModel && options.posVar > 0) {
        const toObject = new THREE.Vector3().subVectors(selectedObject.position, camera.position);
        const distance = toObject.length();

        // Get camera's local right and up vectors to define the offset plane
        const right = new THREE.Vector3().crossVectors(camera.up, toObject).normalize();
        const up = new THREE.Vector3().crossVectors(toObject, right).normalize();

        // Calculate the dimensions of the visible area (frustum) at the object's distance
        const fovInRadians = THREE.MathUtils.degToRad(camera.fov);
        const frustumHeight = 2.0 * distance * Math.tan(fovInRadians / 2.0);
        const frustumWidth = frustumHeight * camera.aspect;

        // Convert the user's percentage (e.g., 90% -> 0.9) to a max shift ratio (e.g., 0.45)
        const maxShiftRatio = options.posVar / 2.0;

        // Calculate random offsets in world units
        const offsetX = (Math.random() - 0.5) * 2 * (frustumWidth * maxShiftRatio);
        const offsetY = (Math.random() - 0.5) * 2 * (frustumHeight * maxShiftRatio);

        // Apply the calculated offsets to the target point
        lookAtTarget.addScaledVector(right, offsetX);
        lookAtTarget.addScaledVector(up, offsetY);
    }

    // Point the camera at the final target (which may be offset)
    camera.lookAt(lookAtTarget);
    // --- END: FIX FOR MODEL CENTERING ---
    
    selectedObject.updateMatrixWorld(true);
    return currentModelAsset;
}

function findTopLevelGroup(object) { let current = object; while (current.parent && current.parent !== scene && current.parent !== maskScene) { current = current.parent; } return current; }
function getCanvasBlob() { return new Promise(resolve => renderer.domElement.toBlob(resolve, 'image/jpeg', 0.9)); }
function updateProgress(current, total, text = 'Processing...') {
    document.getElementById('progress-bar').style.width = `${total > 0 ? (current / total) * 100 : 0}%`;
    document.getElementById('progress-label').textContent = text;
    document.getElementById('progress-count').textContent = `(${current}/${total})`;
}
function setActiveCard(containerId, name) { document.querySelectorAll(`#${containerId} .asset-card`).forEach(c => c.classList.toggle('active', c.dataset.name === name)); }

function showNotification({ text, type = 'success', duration = 3000, id = null }) {
    const container = document.getElementById('notification-container');
    if (id) { const existing = document.getElementById(id); if (existing) { existing.querySelector('span').textContent = text; return; } }
    const el = document.createElement('div');
    el.id = id || `notif-${Date.now()}`;
    el.className = `notification glass-ui p-3 px-4 rounded-lg flex items-center gap-3`;
    const colors = { success: 'bg-green-500/50', error: 'bg-red-500/50', info: 'bg-blue-500/50', loading: 'bg-stone-500/50' };
    const icons = { success: 'check-circle', error: 'alert-triangle', info: 'info', loading: null };
    el.classList.add(colors[type]);
    let iconHtml = type === 'loading' ? `<div class="notification-spinner"></div>` : `<i data-feather="${icons[type]}" class="w-5 h-5"></i>`;
    el.innerHTML = `${iconHtml}<span>${text}</span>`;
    container.appendChild(el);
    if (icons[type]) feather.replace();
    setTimeout(() => el.classList.add('show'), 10);
    if (type !== 'loading') { setTimeout(() => hideNotification(el.id), duration); }
}

function hideNotification(id) {
    const el = document.getElementById(id);
    if (el) { el.classList.remove('show'); setTimeout(() => el.remove(), 500); }
}

function animate() {
    requestAnimationFrame(animate);
    orbitControls.update();
    renderer.render(scene, camera);
}
window.addEventListener('resize', () => { camera.aspect = window.innerWidth / window.innerHeight; camera.updateProjectionMatrix(); renderer.setSize(window.innerWidth, window.innerHeight); });

async function generateSampleData(options, classIndex = 0) {
    let labelData = null;
    if (selectedObject && selectedObject.visible && classIndex !== -1 && options.task !== 'classification') {
        if (options.task === 'segmentation' && options.segmentationMethod === 'render') {
            labelData = await generateRenderMaskData(classIndex, options);
        } else {
            const points = getProjected2dVertices();
            if (points && points.length > 3) {
                if (options.task === 'detection') {
                    const hull = computeConvexHull(points);
                    if (hull.length > 0) labelData = getDetectionLabel(hull, classIndex);
                } else if (options.task === 'segmentation') {
                    if (options.segmentationMethod === 'concave') {
                        labelData = generateConcaveHullData(points, classIndex, options.concavity);
                    } else {
                        const hull = computeConvexHull(points);
                        if (hull.length > 0) labelData = getSegmentationLabel(hull, classIndex);
                    }
                }
            }
        }
    }
    renderer.render(scene, camera);
    const imageBlob = await getCanvasBlob();
    return { imageBlob, labelData };
}

function getDetectionLabel(hull, classIndex) {
    let minX = 1, maxX = -1, minY = 1, maxY = -1;
    for (const p of hull) {
        minX = Math.min(minX, p.x); maxX = Math.max(maxX, p.x);
        minY = Math.min(minY, p.y); maxY = Math.max(maxY, p.y);
    }
    if (maxX <= minX || maxY <= minY) return null;
    const normCenterX = (((minX + maxX) / 2) + 1) / 2;
    const normCenterY = (-((minY + maxY) / 2) + 1) / 2;
    const normWidth = (maxX - minX) / 2;
    const normHeight = (maxY - minY) / 2;
    if (normWidth < 0.001 || normHeight < 0.001) return null;
    return `${classIndex} ${normCenterX.toFixed(6)} ${normCenterY.toFixed(6)} ${normWidth.toFixed(6)} ${normHeight.toFixed(6)}`;
}

function getSegmentationLabel(points, classIndex) {
    if (points.length < 3) return null;
    const yoloPoints = points.map(p => `${((p.x + 1) / 2).toFixed(6)} ${((-p.y + 1) / 2).toFixed(6)}`).join(' ');
    return `${classIndex} ${yoloPoints}`;
}

function generateConcaveHullData(points, classIndex, concavity) {
    const screenPoints = points.map(p => [p.x, p.y]);
    const hull = concaveman(screenPoints, concavity);
    if (hull.length < 3) return getSegmentationLabel(computeConvexHull(points), classIndex);
    const mappedHull = hull.map(p => ({ x: p[0], y: p[1] }));
    return getSegmentationLabel(mappedHull, classIndex);
}

async function generateRenderMaskData(classIndex, options) {
    const assetInfo = modelAssets.find(m => m.name === selectedObject.name);
    if (!assetInfo || !assetInfo.originalMaterials) return null;

    const originalBackground = scene.background;

    panoramaSphere.visible = false;

    scene.background = new THREE.Color(0x000000);
    selectedObject.traverse(node => { if (node.isMesh) node.material = maskMaterial; });
    renderer.render(scene, camera);

    const { width, height } = renderer.domElement;
    const context = renderer.getContext();
    const pixelData = new Uint8Array(width * height * 4);
    context.readPixels(0, 0, width, height, context.RGBA, context.UNSIGNED_BYTE, pixelData);

    scene.background = originalBackground;
    selectedObject.traverse(node => { if (node.isMesh) node.material = assetInfo.originalMaterials.get(node); });

    panoramaSphere.visible = true;

    const contours = findContours(pixelData, width, height);
    if (contours.length === 0) return null;

    contours.sort((a, b) => b.length - a.length);
    const mainContour = contours[0];

    const simplifiedContour = rdp(mainContour, options.rdpEpsilon);
    if (simplifiedContour.length < 3) return null;

    const yoloPoints = simplifiedContour.map(p => `${(p.x / width).toFixed(6)} ${((height - p.y) / height).toFixed(6)}`).join(' ');
    return `${classIndex} ${yoloPoints}`;
}

function computeConvexHull(points) {
    if (points.length <= 3) return points;
    points.sort((a, b) => a.x === b.x ? a.y - b.y : a.x - b.x);
    const cross_product = (o, a, b) => (a.x - o.x) * (b.y - o.y) - (a.y - o.y) * (b.x - o.x);
    const lower = [];
    for (const p of points) {
        while (lower.length >= 2 && cross_product(lower[lower.length - 2], lower[lower.length - 1], p) <= 0) lower.pop();
        lower.push(p);
    }
    const upper = [];
    for (let i = points.length - 1; i >= 0; i--) {
        const p = points[i];
        while (upper.length >= 2 && cross_product(upper[upper.length - 2], upper[upper.length - 1], p) <= 0) upper.pop();
        upper.push(p);
    }
    upper.pop(); lower.pop();
    return lower.concat(upper);
}

function getProjected2dVertices() {
    if (!selectedObject || !selectedObject.visible) return null;
    const projectedVertices = [];
    selectedObject.updateMatrixWorld(true);
    selectedObject.traverse(node => {
        if (node.isMesh) {
            const geometry = node.geometry;
            const positionAttribute = geometry.attributes.position;
            if (!positionAttribute) return;
            const vertex = new THREE.Vector3();
            for (let i = 0; i < positionAttribute.count; i++) {
                vertex.fromBufferAttribute(positionAttribute, i);
                vertex.applyMatrix4(node.matrixWorld);
                vertex.project(camera);
                if (vertex.z < 1) projectedVertices.push({ x: vertex.x, y: vertex.y });
            }
        }
    });
    return projectedVertices.length > 0 ? projectedVertices : null;
}

function rdp(points, epsilon) {
    if (points.length < 3) return points;
    let dmax = 0; let index = 0;
    const end = points.length - 1;
    for (let i = 1; i < end; i++) {
        const d = perpendicularDistance(points[i], points[0], points[end]);
        if (d > dmax) { index = i; dmax = d; }
    }
    if (dmax > epsilon) {
        const recResults1 = rdp(points.slice(0, index + 1), epsilon);
        const recResults2 = rdp(points.slice(index, end + 1), epsilon);
        return recResults1.slice(0, recResults1.length - 1).concat(recResults2);
    } else { return [points[0], points[end]]; }
}
function perpendicularDistance(point, lineStart, lineEnd) {
    let dx = lineEnd.x - lineStart.x; let dy = lineEnd.y - lineStart.y;
    if (dx === 0 && dy === 0) return Math.sqrt(Math.pow(point.x - lineStart.x, 2) + Math.pow(point.y - lineStart.y, 2));
    const t = ((point.x - lineStart.x) * dx + (point.y - lineStart.y) * dy) / (dx * dx + dy * dy);
    const projectionX = lineStart.x + t * dx; const projectionY = lineStart.y + t * dy;
    return Math.sqrt(Math.pow(point.x - projectionX, 2) + Math.pow(point.y - projectionY, 2));
}

function findContours(data, width, height) {
    const visited = new Uint8Array(width * height);
    const contours = [];
    const isWhite = (x, y) => {
        if (x < 0 || x >= width || y < 0 || y >= height) return false;
        return data[(y * width + x) * 4] > 128;
    };
    const directions = [[0, -1], [1, 0], [0, 1], [-1, 0]]; // N, E, S, W

    // Corrected loops to scan the entire image from (0,0).
    for (let y = 0; y < height; y++) {
        for (let x = 0; x < width; x++) {
            // The condition `!isWhite(x, y - 1)` finds the top edge of a shape.
            if (isWhite(x, y) && !isWhite(x, y - 1) && !visited[y * width + x]) {
                const path = [];
                let cx = x, cy = y;
                // Start by looking East (1) from the initial point.
                let dir = 1; 
                
                while (true) {
                    path.push({ x: cx, y: cy });
                    visited[cy * width + cx] = 1;
                    
                    let foundNext = false;
                    // Check directions starting from the one to the left of the current direction.
                    for (let i = 0; i < 4; i++) {
                        const nextDir = (dir + 3 + i) % 4; // Turn left, then straight, then right
                        const [dx, dy] = directions[nextDir];
                        const nx = cx + dx, ny = cy + dy;

                        if (isWhite(nx, ny)) {
                            cx = nx;
                            cy = ny;
                            dir = nextDir;
                            foundNext = true;
                            break;
                        }
                    }
                    if (!foundNext || (cx === x && cy === y)) {
                        // Break if we're stuck or returned to the start
                        break;
                    }
                }
                if (path.length > 10) { // Keep a minimum length to avoid noise
                    contours.push(path);
                }
            }
        }
    }
    return contours;
}

init();