let session = null;
let stoi = null;
let itos = null;

const chatContainer = document.getElementById('chat-container');
const userInput = document.getElementById('user-input');
const sendBtn = document.getElementById('send-btn');
const loadingMsg = document.getElementById('loading-msg');

function addMessage(text, sender) {
    const msgDiv = document.createElement('div');
    msgDiv.classList.add('message', sender === 'user' ? 'user-message' : 'ai-message');
    msgDiv.innerText = text;
    chatContainer.appendChild(msgDiv);
    chatContainer.scrollTop = chatContainer.scrollHeight;
}

async function initONNX() {
    try {
        const vocabRes = await fetch('vocab.json');
        const vocab = await vocabRes.json();
        stoi = vocab.stoi;
        itos = vocab.itos;

        ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/";
        ort.env.wasm.numThreads = 1;

        // The critical fix: ensuring the ONNX model is fetched correctly
        session = await ort.InferenceSession.create('./ares.onnx', { executionProviders: ['wasm'] });
        
        loadingMsg.innerHTML = '<p>Ares ONNX Model successfully loaded into browser VRAM! Ready.</p>';
        userInput.disabled = false;
        sendBtn.disabled = false;

    } catch (e) {
        console.error(e);
        loadingMsg.innerHTML = `<p style="color:red">Failed to load ONNX Model: ${e.message}</p>`;
    }
}

async function generateAresResponse(query) {
    const prompt = `<|user|>${query}<|end|><|assistant|>`;
    let tokens = prompt.split('').map(c => stoi[c] || 0);

    const max_new_tokens = 40;
    const vocab_size = Object.keys(stoi).length;

    for (let i = 0; i < max_new_tokens; i++) {
        // Use Int32Array to avoid BigInt/WASM crash
        const int32Tokens = Int32Array.from(tokens);
        const tensor = new ort.Tensor('int32', int32Tokens, [1, tokens.length]);
        
        const feeds = { input_tokens: tensor };
        const results = await session.run(feeds);
        
        const logits = results.logits.data; 
        const offset = (tokens.length - 1) * vocab_size;
        
        let max_val = -Infinity;
        let best_idx = 0;
        for (let v = 0; v < vocab_size; v++) {
            let val = logits[offset + v];
            if (val > max_val) {
                max_val = val;
                best_idx = v;
            }
        }

        tokens.push(best_idx);
        let char = itos[best_idx];

        if (char === "<|end|>") break;
    }

    const outputString = tokens.map(t => itos[t]).join('');
    let reply = outputString.substring(prompt.length);
    if(reply.includes("<|end|>")) reply = reply.split("<|end|>")[0];
    
    return reply;
}

sendBtn.addEventListener('click', async () => {
    const text = userInput.value.trim();
    if (!text) return;
    
    addMessage(text, 'user');
    userInput.value = '';

    const placeholder = "Thinking...";
    addMessage(placeholder, 'ai');
    
    try {
        const reply = await generateAresResponse(text);
        chatContainer.lastChild.innerText = reply;
    } catch(e) {
        console.error(e);
        chatContainer.lastChild.innerText = "Inference Error: " + e.message;
    }
});

userInput.addEventListener('keypress', (e) => {
    if (e.key === 'Enter') {
        sendBtn.click();
    }
});

initONNX();