// 1. Update the showResults function to better handle NER specifics

function showResults(result) {
    resultsSection.classList.remove('d-none');
    
    // Hide all result sections first
    document.getElementById('sentiment-results').classList.add('d-none');
    document.getElementById('intent-results').classList.add('d-none');
    document.getElementById('topic-results').classList.add('d-none');
    document.getElementById('ner-results').classList.add('d-none');
    
    const task = taskSelect.value;
    const method = methodSelect.value;
    const resultSection = document.getElementById(`${task}-results`);
    resultSection.classList.remove('d-none');

    if (task === 'ner') {
        // For NER, we only handle the display of entities
        displayNERResults(result);
    } else {
        // Common fields for other tasks
        const labelElement = document.getElementById(`${task}-label`);
        const scoreElement = document.getElementById(`${task}-score`);
        
        // Set label and score
        labelElement.textContent = result.label;
        const scorePercentage = (result.score * 100).toFixed(2);
        scoreElement.style.width = `${scorePercentage}%`;
        scoreElement.textContent = `${scorePercentage}%`;

        // Set appropriate badge colors
        setBadgeColor(labelElement, result.label);

        // Reset all explanation containers
        ['sentiment', 'intent', 'topic'].forEach(taskType => {
            const container = document.getElementById(`${taskType}-explanation-container`);
            const element = document.getElementById(`${taskType}-explanation`);
            if (container) container.classList.add('d-none');
            if (element) element.innerHTML = '';
        });

        // Task-specific handling remains the same...
    }
}

// 2. Update the displayNERResults function for better error handling and display

function displayNERResults(result) {
    const entityGroups = document.getElementById('ner-entity-groups');
    const highlightedText = document.getElementById('ner-highlighted-text');
    
    // Clear previous results
    entityGroups.innerHTML = '';
    highlightedText.innerHTML = '';

    // Handle errors or empty results
    if (!result || !result.entities) {
        entityGroups.innerHTML = '<p class="text-muted">Error processing entities.</p>';
        highlightedText.innerHTML = inputText.value;
        return;
    }

    if (result.entities.length === 0) {
        entityGroups.innerHTML = '<p class="text-muted">No entities found in the text.</p>';
        highlightedText.innerHTML = `<p class="text-muted">${inputText.value}</p>`;
        return;
    }

    // Group entities by type
    const groupedEntities = {};
    result.entities.forEach(entity => {
        if (!groupedEntities[entity.entity]) {
            groupedEntities[entity.entity] = [];
        }
        groupedEntities[entity.entity].push(entity);
    });

    // Display entity groups
    Object.entries(groupedEntities).forEach(([type, entities]) => {
        const typeColor = entityColors[type] || entityColors.DEFAULT;
        const typeDiv = document.createElement('div');
        typeDiv.className = 'mb-3 p-2 border rounded';
        typeDiv.innerHTML = `
            <div class="d-flex align-items-center mb-2">
                <div class="fw-bold">${type}</div>
                <div class="ms-2 badge bg-secondary">${entities.length}</div>
            </div>
            <div class="d-flex flex-wrap gap-2">
                ${entities.map(entity => `
                    <span class="badge px-2 py-1" 
                          style="background-color: ${typeColor}; color: #000;"
                          title="Confidence: ${(entity.score * 100).toFixed(1)}%">
                        ${entity.word}
                    </span>
                `).join('')}
            </div>
        `;
        entityGroups.appendChild(typeDiv);
    });

    // Create highlighted text with improved styling
    let text = inputText.value;
    let lastIndex = 0;
    const spans = [];

    // Sort entities by start position
    const sortedEntities = [...result.entities].sort((a, b) => a.start - b.start);

    sortedEntities.forEach(entity => {
        if (entity.start > lastIndex) {
            spans.push(escapeHtml(text.substring(lastIndex, entity.start)));
        }
        const typeColor = entityColors[entity.entity] || entityColors.DEFAULT;
        spans.push(`<span class="px-1 rounded" 
                         style="background-color: ${typeColor}; cursor: help;" 
                         title="${entity.entity} (${(entity.score * 100).toFixed(1)}% confident)">
                    ${escapeHtml(text.substring(entity.start, entity.end))}</span>`);
        lastIndex = entity.end;
    });

    if (lastIndex < text.length) {
        spans.push(escapeHtml(text.substring(lastIndex)));
    }

    highlightedText.innerHTML = `<div class="p-3 border rounded">${spans.join('')}</div>`;
}

// 3. Add a helper function for HTML escaping
function escapeHtml(text) {
    const div = document.createElement('div');
    div.textContent = text;
    return div.innerHTML;
}