document.addEventListener('DOMContentLoaded', () => {
    // State management
    let currentFile = null;
    let isDrawing = false;
    let startX, startY;
    let roi = { x1: 0, y1: 0, x2: 100, y2: 100 };
    let previewImage = new Image();

    // Elements
    const modelDropZone = document.getElementById('model-drop-zone');
    const modelInput = document.getElementById('model-input');
    const modelStatus = document.getElementById('model-status');
    const statusText = document.getElementById('status-text');
    const statusIcon = modelStatus.querySelector('i');

    const mediaDropZone = document.getElementById('media-drop-zone');
    const mediaInput = document.getElementById('media-input');
    
    const previewSection = document.getElementById('preview-section');
    const roiCanvas = document.getElementById('roi-canvas');
    const ctx = roiCanvas.getContext('2d');
    
    const thresholdInput = document.getElementById('threshold-input');
    const confMaxInput = document.getElementById('conf-max-input');
    const confRangeVal = document.getElementById('conf-range-val');
    const roiX1 = document.getElementById('roi-x1');
    const roiY1 = document.getElementById('roi-y1');
    const roiX2 = document.getElementById('roi-x2');
    const roiY2 = document.getElementById('roi-y2');
    const resetRoiBtn = document.getElementById('reset-roi-btn');
    
    const progressCard = document.getElementById('progress-card');
    const loading = document.getElementById('loading');
    const videoProgressContainer = document.getElementById('video-progress-container');
    const videoProgressBar = document.getElementById('video-progress-bar');
    const videoStatusMsg = document.getElementById('video-status-msg');
    const videoPercentage = document.getElementById('video-percentage');

    const analyzeBtn = document.getElementById('analyze-btn');

    const resultSection = document.getElementById('result-section');
    const resultImage = document.getElementById('result-image');
    const resultCount = document.getElementById('result-count');
    const downloadBtn = document.getElementById('download-btn');

    const videoResultSection = document.getElementById('video-result-section');
    const resultVideo = document.getElementById('result-video');
    const videoDownloadBtn = document.getElementById('video-download-btn');

    // Drag and Drop Setup
    [modelDropZone, mediaDropZone].forEach(zone => {
        ['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
            zone.addEventListener(eventName, e => {
                e.preventDefault();
                e.stopPropagation();
            });
        });

        ['dragenter', 'dragover'].forEach(eventName => {
            zone.addEventListener(eventName, () => zone.classList.add('dragover'));
        });

        ['dragleave', 'drop'].forEach(eventName => {
            zone.addEventListener(eventName, () => zone.classList.remove('dragover'));
        });
    });

    // Clicks
    modelDropZone.addEventListener('click', () => modelInput.click());
    mediaDropZone.addEventListener('click', () => mediaInput.click());

    // File Handlers
    modelInput.addEventListener('change', e => handleModelUpload(e.target.files[0]));
    modelDropZone.addEventListener('drop', e => handleModelUpload(e.dataTransfer.files[0]));

    mediaInput.addEventListener('change', e => handleMediaSelection(e.target.files[0]));
    mediaDropZone.addEventListener('drop', e => handleMediaSelection(e.dataTransfer.files[0]));

    // Settings
    const updateConfLabel = () => {
        const min = Math.round(thresholdInput.value * 100);
        const max = Math.round(confMaxInput.value * 100);
        confRangeVal.innerText = `${min}% - ${max}%`;
    };

    thresholdInput.addEventListener('input', updateConfLabel);
    confMaxInput.addEventListener('input', updateConfLabel);

    [roiX1, roiY1, roiX2, roiY2].forEach(input => {
        input.addEventListener('change', updateROIFromInputs);
    });

    resetRoiBtn.addEventListener('click', () => {
        roi = { x1: 0, y1: 0, x2: 100, y2: 100 };
        updateInputsFromROI();
        drawROI();
    });

    analyzeBtn.addEventListener('click', startInference);

    // --- Functions ---

    async function handleModelUpload(file) {
        if (!file || !file.name.endsWith('.pt')) {
            showToast('Please upload a valid YOLO .pt model.', 'error');
            return;
        }

        const formData = new FormData();
        formData.append('file', file);

        statusText.innerText = 'Uploading model...';
        modelStatus.classList.remove('loaded');
        statusIcon.className = 'fas fa-spinner fa-spin';

        try {
            const resp = await fetch('/upload-model', { method: 'POST', body: formData });
            const data = await resp.json();
            if (data.status === 'success') {
                statusText.innerText = `Model: ${file.name}`;
                modelStatus.classList.add('loaded');
                statusIcon.className = 'fas fa-check-circle';
                showToast('Model loaded successfully!', 'success');
            } else {
                throw new Error(data.detail);
            }
        } catch (err) {
            statusText.innerText = 'Error loading model';
            statusIcon.className = 'fas fa-exclamation-circle';
            showToast(err.message, 'error');
        }
    }

    async function handleMediaSelection(file) {
        if (!file) return;
        currentFile = file;

        // Reset state
        resultSection.classList.add('hidden');
        videoResultSection.classList.add('hidden');
        progressCard.classList.add('hidden');
        
        if (file.type.startsWith('image/')) {
            const reader = new FileReader();
            reader.onload = e => {
                previewImage.onload = () => initCanvas();
                previewImage.src = e.target.result;
            };
            reader.readAsDataURL(file);
        } else if (file.type.startsWith('video/')) {
            extractVideoFrame(file);
        } else {
            showToast('Unsupported file type.', 'error');
        }
    }

    function extractVideoFrame(file) {
        const video = document.createElement('video');
        video.preload = 'metadata';
        video.src = URL.createObjectURL(file);
        video.onloadedmetadata = () => {
            video.currentTime = 0.1; // Seek a bit in to avoid black frames
        };
        video.onseeked = () => {
            const tempCanvas = document.createElement('canvas');
            tempCanvas.width = video.videoWidth;
            tempCanvas.height = video.videoHeight;
            const tempCtx = tempCanvas.getContext('2d');
            tempCtx.drawImage(video, 0, 0);
            previewImage.onload = () => initCanvas();
            previewImage.src = tempCanvas.toDataURL('image/jpeg');
            URL.revokeObjectURL(video.src);
        };
    }

    function initCanvas() {
        previewSection.classList.remove('hidden');
        previewSection.scrollIntoView({ behavior: 'smooth' });
        
        // Scale canvas to fit container but keep aspect ratio
        const containerWidth = roiCanvas.parentElement.clientWidth;
        const scale = containerWidth / previewImage.width;
        
        roiCanvas.width = previewImage.width * scale;
        roiCanvas.height = previewImage.height * scale;
        
        drawROI();
    }

    function drawROI() {
        ctx.clearRect(0, 0, roiCanvas.width, roiCanvas.height);
        ctx.drawImage(previewImage, 0, 0, roiCanvas.width, roiCanvas.height);
        
        // Darken outside
        ctx.fillStyle = 'rgba(0, 0, 0, 0.4)';
        
        const x1 = (roi.x1 / 100) * roiCanvas.width;
        const y1 = (roi.y1 / 100) * roiCanvas.height;
        const x2 = (roi.x2 / 100) * roiCanvas.width;
        const y2 = (roi.y2 / 100) * roiCanvas.height;
        
        const w = x2 - x1;
        const h = y2 - y1;

        // Draw overlay path with a "hole" for the ROI
        ctx.beginPath();
        ctx.rect(0, 0, roiCanvas.width, roiCanvas.height);
        ctx.rect(x1, y1, w, h);
        ctx.fill('evenodd');

        // Draw border
        ctx.strokeStyle = '#f59e0b';
        ctx.lineWidth = 3;
        ctx.setLineDash([5, 5]);
        ctx.strokeRect(x1, y1, w, h);
        
        // Optional: corner handles design look
        ctx.fillStyle = '#f59e0b';
        ctx.fillRect(x1-4, y1-4, 8, 8);
        ctx.fillRect(x2-4, y1-4, 8, 8);
        ctx.fillRect(x1-4, y2-4, 8, 8);
        ctx.fillRect(x2-4, y2-4, 8, 8);
    }

    // Canvas Events
    roiCanvas.addEventListener('mousedown', e => {
        isDrawing = true;
        const rect = roiCanvas.getBoundingClientRect();
        startX = e.clientX - rect.left;
        startY = e.clientY - rect.top;
        
        roi.x1 = (startX / roiCanvas.width) * 100;
        roi.y1 = (startY / roiCanvas.height) * 100;
    });

    roiCanvas.addEventListener('mousemove', e => {
        if (!isDrawing) return;
        const rect = roiCanvas.getBoundingClientRect();
        const curX = e.clientX - rect.left;
        const curY = e.clientY - rect.top;
        
        roi.x2 = (curX / roiCanvas.width) * 100;
        roi.y2 = (curY / roiCanvas.height) * 100;
        
        updateInputsFromROI();
        drawROI();
    });

    roiCanvas.addEventListener('mouseup', () => {
        isDrawing = false;
        // Normalize coordinates (ensure x1 < x2, y1 < y2)
        if (roi.x1 > roi.x2) [roi.x1, roi.x2] = [roi.x2, roi.x1];
        if (roi.y1 > roi.y2) [roi.y1, roi.y2] = [roi.y2, roi.y1];
        updateInputsFromROI();
        drawROI();
    });

    function updateInputsFromROI() {
        roiX1.value = Math.round(roi.x1);
        roiY1.value = Math.round(roi.y1);
        roiX2.value = Math.round(roi.x2);
        roiY2.value = Math.round(roi.y2);
    }

    function updateROIFromInputs() {
        roi.x1 = parseInt(roiX1.value);
        roi.y1 = parseInt(roiY1.value);
        roi.x2 = parseInt(roiX2.value);
        roi.y2 = parseInt(roiY2.value);
        drawROI();
    }

    async function startInference() {
        if (!currentFile) return;

        progressCard.classList.remove('hidden');
        progressCard.scrollIntoView({ behavior: 'smooth' });
        
        const isVideo = currentFile.type.startsWith('video/');
        const formData = new FormData();
        formData.append('file', currentFile);
        formData.append('conf_min', thresholdInput.value);
        formData.append('conf_max', confMaxInput.value);
        formData.append('roi', JSON.stringify(roi));

        if (isVideo) {
            handleVideoInference(formData);
        } else {
            handleImageInference(formData);
        }
    }

    async function handleImageInference(formData) {
        loading.classList.remove('hidden');
        videoProgressContainer.classList.add('hidden');
        
        try {
            const resp = await fetch('/inference', { method: 'POST', body: formData });
            const data = await resp.json();
            if (data.status === 'success') {
                resultImage.src = data.image;
                resultCount.innerText = `${data.count} Detections`;
                resultSection.classList.remove('hidden');
                resultSection.scrollIntoView({ behavior: 'smooth' });
            } else {
                throw new Error(data.detail);
            }
        } catch (err) {
            showToast(err.message, 'error');
        } finally {
            progressCard.classList.add('hidden');
        }
    }

    async function handleVideoInference(formData) {
        loading.classList.add('hidden');
        videoProgressContainer.classList.remove('hidden');
        videoProgressBar.style.width = '0%';
        videoPercentage.innerText = '0%';
        videoStatusMsg.innerText = 'Uploading video...';

        try {
            const resp = await fetch('/inference-video', { method: 'POST', body: formData });
            const data = await resp.json();
            if (data.status === 'success') {
                videoStatusMsg.innerText = 'Processing frames...';
                pollVideoProgress(data.task_id);
            } else {
                throw new Error(data.detail);
            }
        } catch (err) {
            showToast(err.message, 'error');
            progressCard.classList.add('hidden');
        }
    }

    function pollVideoProgress(taskId) {
        const interval = setInterval(async () => {
            try {
                const resp = await fetch(`/video-progress/${taskId}`);
                const data = await resp.json();

                if (data.status === 'processing') {
                    videoProgressBar.style.width = `${data.progress}%`;
                    videoPercentage.innerText = `${data.progress}%`;
                } else if (data.status === 'completed') {
                    clearInterval(interval);
                    videoProgressBar.style.width = '100%';
                    videoPercentage.innerText = '100%';
                    videoStatusMsg.innerText = 'Processing complete!';
                    
                    showVideoResult(taskId);
                } else if (data.status === 'error') {
                    clearInterval(interval);
                    showToast(data.message, 'error');
                    progressCard.classList.add('hidden');
                }
            } catch (err) {
                console.error(err);
            }
        }, 1000);
    }

    function showVideoResult(taskId) {
        const url = `/video-result/${taskId}`;
        resultVideo.src = url;
        videoDownloadBtn.href = url;
        videoResultSection.classList.remove('hidden');
        videoResultSection.scrollIntoView({ behavior: 'smooth' });
        progressCard.classList.add('hidden');
    }

    function showToast(message, type = 'info') {
        const toast = document.createElement('div');
        toast.className = `toast ${type}`;
        Object.assign(toast.style, {
            position: 'fixed', bottom: '20px', right: '20px', padding: '1rem 1.5rem',
            borderRadius: '10px', color: 'white', zIndex: '1000',
            background: type === 'error' ? '#ef4444' : '#10b981',
            boxShadow: '0 4px 15px rgba(0,0,0,0.3)', animation: 'slideIn 0.3s ease forwards'
        });
        toast.innerText = message;
        document.body.appendChild(toast);
        setTimeout(() => {
            toast.style.animation = 'slideOut 0.3s ease forwards';
            setTimeout(() => toast.remove(), 300);
        }, 3000);
    }
});