// index.js
import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.3';

// DOM elements
const taskSelect = document.getElementById('task-select');
const deviceSelect = document.getElementById('device-select');
const inputText = document.getElementById('input-text');
const runBtn = document.getElementById('run-btn');
const outputContainer = document.getElementById('output');

// Model mapping
const modelMap = {
    'text-classification': 'distilbert-base-uncased-finetuned-sst-2-english',
    'sentiment-analysis': 'distilbert-base-uncased-finetuned-sst-2-english',
    'question-answering': 'deepset/roberta-base-squad2',
    'text-generation': 'gpt2',
    'translation': 't5-small'
};

// Initialize pipeline
let currentPipeline = null;
let currentTask = null;

// Check WebGPU support
function checkWebGPUSupport() {
    return 'gpu' in navigator;
}

// Update device options based on support
if (!checkWebGPUSupport()) {
    const option = document.querySelector('#device-select option[value="webgpu"]');
    option.disabled = true;
    option.textContent = 'GPU (WebGPU) - Not Supported';
    deviceSelect.value = 'cpu';
}

// Function to create pipeline
async function createPipeline(task, device) {
    const modelName = modelMap[task];
    
    // Show loading state
    outputContainer.innerHTML = '<p class="loading">Loading model...</p>';
    
    try {
        // Create pipeline with selected device
        if (device === 'webgpu' && checkWebGPUSupport()) {
            return await pipeline(task, modelName, { device: 'webgpu' });
        } else {
            return await pipeline(task, modelName);
        }
    } catch (error) {
        console.error('Error creating pipeline:', error);
        outputContainer.innerHTML = `<p class="error">Error loading model: ${error.message}</p>`;
        return null;
    }
}

// Function to run the model
async function runModel() {
    const task = taskSelect.value;
    const device = deviceSelect.value;
    const text = inputText.value.trim();
    
    if (!text) {
        outputContainer.innerHTML = '<p class="error">Please enter some text to process.</p>';
        return;
    }
    
    // Check if pipeline needs to be recreated
    if (!currentPipeline || currentTask !== task) {
        currentPipeline = await createPipeline(task, device);
        currentTask = task;
    }
    
    if (!currentPipeline) {
        return;
    }
    
    // Show processing state
    outputContainer.innerHTML = '<p class="loading">Processing...</p>';
    
    try {
        let result;
        
        switch (task) {
            case 'text-classification':
            case 'sentiment-analysis':
                result = await currentPipeline(text);
                break;
            case 'question-answering':
                // For QA, we need context and question
                const [context, question] = text.split('\n').length > 1 
                    ? [text.split('\n').slice(0, -1).join(' '), text.split('\n').pop()] 
                    : ['The sky is blue.', text];
                result = await currentPipeline(question, context);
                break;
            case 'text-generation':
                result = await currentPipeline(text, { max_new_tokens: 50 });
                break;
            case 'translation':
                result = await currentPipeline(text, { 
                    src_lang: 'en', 
                    tgt_lang: 'de' 
                });
                break;
            default:
                throw new Error('Unsupported task');
        }
        
        // Display results
        displayResults(result, task);
    } catch (error) {
        console.error('Error running model:', error);
        outputContainer.innerHTML = `<p class="error">Error processing text: ${error.message}</p>`;
    }
}

// Function to display results based on task
function displayResults(result, task) {
    let content = '';
    
    switch (task) {
        case 'text-classification':
        case 'sentiment-analysis':
            content = `
                <h4>Classification Results:</h4>
                <ul>
                    ${result.map(item => `
                        <li>
                            <strong>${item.label}:</strong> ${(item.score * 100).toFixed(2)}%
                        </li>
                    `).join('')}
                </ul>
            `;
            break;
        case 'question-answering':
            content = `
                <h4>Answer:</h4>
                <p>${result.answer}</p>
                <p><strong>Confidence:</strong> ${(result.score * 100).toFixed(2)}%</p>
            `;
            break;
        case 'text-generation':
            content = `
                <h4>Generated Text:</h4>
                <p>${result[0].generated_text}</p>
            `;
            break;
        case 'translation':
            content = `
                <h4>Translation:</h4>
                <p>${result[0].translation_text}</p>
            `;
            break;
        default:
            content = `<pre>${JSON.stringify(result, null, 2)}</pre>`;
    }
    
    outputContainer.innerHTML = content;
}

// Event listeners
runBtn.addEventListener('click', runModel);

// Handle task change
taskSelect.addEventListener('change', () => {
    // Reset pipeline when task changes
    currentPipeline = null;
    currentTask = null;
    
    // Update placeholder text based on task
    switch (taskSelect.value) {
        case 'question-answering':
            inputText.placeholder = "Enter context and question separated by a new line\nContext: The sky is blue.\nQuestion: What color is the sky?";
            break;
        case 'text-generation':
            inputText.placeholder = "Enter a prompt to generate text from...";
            break;
        case 'translation':
            inputText.placeholder = "Enter text to translate from English to German...";
            break;
        default:
            inputText.placeholder = "Enter your text here...";
    }
});

// Initialize with default task
taskSelect.dispatchEvent(new Event('change'));