index.js

import { Wllama } from 'https://cdn.jsdelivr.net/npm/@wllama/wllama@2.3.1/esm/index.js';

// Constants
const MODEL_MAX_CONTEXT = 512;
const MODEL_FILE = 'model_f16.gguf';
const textDecoder = new TextDecoder();

// State
let wllama = null;
let engineReady = false;
let isGenerating = false;
let userStopped = false;
let abortController = null;
let progressTimer = null;  // delay timer — prevents flicker on fast connections

// DOM
const deviceBadge       = document.getElementById('deviceBadge');
const promptInput       = document.getElementById('promptInput');
const generateBtn       = document.getElementById('generateBtn');
const generateBtnText   = document.getElementById('generateBtnText');
const stopBtn           = document.getElementById('stopBtn');
const copyBtn           = document.getElementById('copyBtn');
const clearBtn          = document.getElementById('clearBtn');
const outputBox         = document.getElementById('outputBox');
const statusText        = document.getElementById('statusText');
const tempInput         = document.getElementById('temperature');
const tempVal           = document.getElementById('tempValue');
const progressOverlay   = document.getElementById('progressOverlay');
const progressFill      = document.getElementById('progressFill');
const progressPct       = document.getElementById('progressPct');
const progressLabel     = document.getElementById('progressLabel');
const progressSub       = document.getElementById('progressSub');
const promptTokenMetric = document.getElementById('promptTokenMetric');
const generatedTokenMetric = document.getElementById('generatedTokenMetric');
const speedMetric       = document.getElementById('speedMetric');
const timeMetric        = document.getElementById('timeMetric');

// Slider UI
tempInput.addEventListener('input', (e) =>
    tempVal.textContent = parseFloat(e.target.value).toFixed(2));

// ──────────────────────────────────────────
// Progress overlay helpers
// ──────────────────────────────────────────
function showProgress(label = 'Downloading model…', sub = 'FP16 · 43 MB · first run only') {
    progressLabel.textContent = label;
    progressSub.textContent = sub;
    setProgress(0);
    // Delay by 400ms — if download finishes faster, overlay never appears (no flicker)
    clearTimeout(progressTimer);
    progressTimer = setTimeout(() => {
        progressOverlay.classList.remove('hidden');
    }, 400);
}

function setProgress(pct, loaded, total) {
    const p = Math.min(100, Math.max(0, pct));
    progressFill.style.width = `${p}%`;
    progressPct.textContent = `${Math.round(p)}%`;
    if (typeof loaded === 'number' && typeof total === 'number' && total > 0) {
        const mb = (total / 1024 / 1024).toFixed(0);
        const done = (loaded / 1024 / 1024).toFixed(1);
        progressSub.textContent = `${done} / ${mb} MB downloaded`;
    }
}

function hideProgress() {
    clearTimeout(progressTimer);  // cancel delayed show if download was fast
    progressTimer = null;
    progressOverlay.classList.add('hidden');
    setProgress(0);
}

// ──────────────────────────────────────────
// Engine initialisation (lazy — called on first Generate click)
// ──────────────────────────────────────────
async function initializeEngine() {
    deviceBadge.textContent = 'Loading…';
    deviceBadge.className = 'badge loading';
    showProgress('Downloading model…', 'FP16 · 43 MB · cached after first run');

    const CONFIG_PATHS = {
        'single-thread/wllama.wasm': 'https://cdn.jsdelivr.net/npm/@wllama/wllama@2.3.1/src/single-thread/wllama.wasm',
        'multi-thread/wllama.wasm':  'https://cdn.jsdelivr.net/npm/@wllama/wllama@2.3.1/src/multi-thread/wllama.wasm',
    };

    wllama = new Wllama(CONFIG_PATHS);

    const nThreads = Math.min(8, navigator.hardwareConcurrency || 4);
    const modelUrl = new URL(`./${MODEL_FILE}`, import.meta.url).href;

    await wllama.loadModelFromUrl(modelUrl, {
        n_ctx: MODEL_MAX_CONTEXT,
        n_threads: nThreads,
        cache_type_k: 'f32',
        cache_type_v: 'f32',
        progressCallback: ({ loaded, total }) => {
            if (total > 0) setProgress((loaded / total) * 100, loaded, total);
        },
    });

    hideProgress();
    engineReady = true;
    deviceBadge.className = 'badge webgpu';
    deviceBadge.textContent = 'FP16 · WASM';
    console.log('Wllama engine ready.');
}

// ──────────────────────────────────────────
// Update prompt-token metric
// ──────────────────────────────────────────
async function updatePromptMetric() {
    if (!wllama) return;
    try {
        const tokens = await wllama.tokenize(promptInput.value || '');
        promptTokenMetric.textContent = `${tokens.length} prompt tok`;
    } catch (_) {}
}

promptInput.addEventListener('input', updatePromptMetric);

// ──────────────────────────────────────────
// Main Generate handler
// ──────────────────────────────────────────
async function handleGenerate() {
    if (isGenerating) return;

    const promptText = promptInput.value.trim();
    if (!promptText) { promptInput.focus(); return; }

    // First click: load the engine, then immediately generate
    if (!engineReady) {
        generateBtn.disabled = true;
        generateBtnText.textContent = 'Loading…';
        statusText.textContent = '';
        try {
            await initializeEngine();
        } catch (err) {
            console.error('Engine load failed:', err);
            deviceBadge.className = 'badge wasm';
            deviceBadge.textContent = 'Error';
            statusText.textContent = 'Failed to load model.';
            generateBtn.disabled = false;
            generateBtnText.textContent = 'Retry';
            hideProgress();
            return;
        }
        generateBtn.disabled = false;
        generateBtnText.textContent = 'Generate';
    }

    await runGeneration(promptText);
}

// ──────────────────────────────────────────
// Token-by-token generation loop
// ──────────────────────────────────────────
async function runGeneration(promptText) {
    const temperature = parseFloat(tempInput.value);
    const maxNewTokens = MODEL_MAX_CONTEXT - 20; // use almost full context

    isGenerating = true;
    userStopped = false;
    abortController = new AbortController();

    generateBtn.disabled = true;
    generateBtnText.textContent = 'Generating…';
    stopBtn.disabled = false;
    statusText.textContent = '';

    // Build output DOM
    outputBox.innerHTML = '';
    const promptSpan = document.createElement('span');
    promptSpan.className = 'prompt-text';
    promptSpan.textContent = promptText;
    const generatedNode = document.createTextNode('');
    const cursorSpan = document.createElement('span');
    cursorSpan.className = 'cursor';
    cursorSpan.textContent = '|';
    outputBox.appendChild(promptSpan);
    outputBox.appendChild(generatedNode);
    outputBox.appendChild(cursorSpan);

    let count = 0;
    const t0 = performance.now();
    const eosId = wllama.getEOS();

    try {
        await wllama.kvClear();

        // BOS fix: wllama tokenizer never adds BOS, so we prepend manually
        const rawTokens   = await wllama.tokenize(promptText, false);
        const promptTokens = [wllama.getBOS(), ...rawTokens];

        await wllama.decode(promptTokens, { skipLogits: false });
        await wllama.samplingInit({
            temp: temperature,
            top_k: 40,
            top_p: 0.90,
            penalty_repeat: 1.15,
            penalty_last_n: 256,
            penalty_present: 0.5,
            penalty_freq: 0.5,
        }, promptTokens);

        for (let i = 0; i < maxNewTokens; i++) {
            if (userStopped || abortController.signal.aborted) break;

            const sampled = await wllama.samplingSample();
            const tok = sampled.token;
            if (tok === eosId || tok === 3 || tok === 0) break;

            await wllama.samplingAccept([tok]);

            const piece = typeof sampled.piece === 'string'
                ? sampled.piece
                : textDecoder.decode(sampled.piece, { stream: true });

            generatedNode.data += piece;
            count++;

            await wllama.decode([tok], { skipLogits: false });

            const elapsed = (performance.now() - t0) / 1000;
            if (elapsed > 0) {
                generatedTokenMetric.textContent = `${count} tok generated`;
                speedMetric.textContent = `${(count / elapsed).toFixed(1)} tok/s`;
                timeMetric.textContent = `${elapsed.toFixed(1)}s`;
            }
            outputBox.scrollTop = outputBox.scrollHeight;
        }

        statusText.textContent = userStopped ? 'Stopped' : 'Done';
        await updatePromptMetric();

    } catch (err) {
        statusText.textContent = userStopped ? 'Stopped' : 'Error';
        if (!userStopped) console.error(err);
    } finally {
        isGenerating = false;
        generateBtn.disabled = false;
        generateBtnText.textContent = 'Generate';
        stopBtn.disabled = true;
        cursorSpan.remove();
    }
}

// ──────────────────────────────────────────
// Event Listeners
// ──────────────────────────────────────────
generateBtn.addEventListener('click', handleGenerate);

stopBtn.addEventListener('click', () => {
    userStopped = true;
    abortController?.abort();
});

copyBtn.addEventListener('click', () => {
    const text = outputBox.innerText.replace(/\|$/, '');
    navigator.clipboard.writeText(text).then(() => {
        const orig = copyBtn.textContent;
        copyBtn.textContent = 'Copied!';
        setTimeout(() => copyBtn.textContent = orig, 1800);
    });
});

clearBtn.addEventListener('click', () => {
    if (isGenerating) return;
    promptInput.value = '';
    outputBox.innerHTML = '<span class="placeholder-text">आपकी कहानी यहाँ दिखाई देगी…</span>';
    promptTokenMetric.textContent = '—';
    generatedTokenMetric.textContent = '—';
    speedMetric.textContent = '—';
    timeMetric.textContent = '—';
    statusText.textContent = '';
});

document.querySelectorAll('.sample-btn').forEach(btn => {
    btn.addEventListener('click', () => {
        if (isGenerating) return;
        const text = btn.getAttribute('data-prompt');
        promptInput.value = text;
        outputBox.innerHTML = `<span class="prompt-text">${text}</span><span class="cursor">|</span>`;
        promptInput.focus();
        if (engineReady) updatePromptMetric();
    });
});