index.js

// index.js content here
class AIChatApp {
    constructor() {
        this.messagesContainer = document.getElementById('messages-container');
        this.messageInput = document.getElementById('message-input');
        this.sendButton = document.getElementById('send-button');
        this.loadingOverlay = document.getElementById('loading-overlay');
        this.statusBar = document.getElementById('status-bar');
        this.statusText = document.getElementById('status-text');
        this.modelStatus = document.getElementById('model-status');
        this.deviceSelect = document.getElementById('device-select');
        this.deviceToggle = document.getElementById('device-toggle');

        this.generator = null;
        this.currentDevice = 'cpu';
        this.isGenerating = false;
        this.messages = [];

        this.initializeEventListeners();
        this.initializeAutoResize();
        this.loadModel();
    }

    initializeEventListeners() {
        // Send message
        this.sendButton.addEventListener('click', () => this.sendMessage());
        this.messageInput.addEventListener('keydown', (e) => {
            if (e.key === 'Enter' && !e.shiftKey) {
                e.preventDefault();
                this.sendMessage();
            }
        });

        // Device selection
        this.deviceSelect.addEventListener('change', (e) => {
            this.currentDevice = e.target.value;
            this.updateDeviceUI();
            if (this.generator) {
                this.loadModel();
            }
        });

        this.deviceToggle.addEventListener('click', () => {
            this.currentDevice = this.currentDevice === 'cpu' ? 'webgpu' : 'cpu';
            this.deviceSelect.value = this.currentDevice;
            this.updateDeviceUI();
            if (this.generator) {
                this.loadModel();
            }
        });

        // Clear messages
        document.addEventListener('keydown', (e) => {
            if (e.ctrlKey && e.key === 'l') {
                e.preventDefault();
                this.clearChat();
            }
        });
    }

    initializeAutoResize() {
        this.messageInput.addEventListener('input', () => {
            this.messageInput.style.height = 'auto';
            this.messageInput.style.height = Math.min(this.messageInput.scrollHeight, 120) + 'px';
        });
    }

    async loadModel() {
        try {
            this.setStatus('Loading AI model...', 'Loading...');
            this.showLoading(true);

            const device = this.currentDevice === 'webgpu' ? { device: 'webgpu' } : {};
            
            // Check WebGPU support if selected
            if (this.currentDevice === 'webgpu') {
                if (!navigator.gpu) {
                    throw new Error('WebGPU is not supported in this browser. Falling back to CPU.');
                }
            }

            this.generator = await pipeline(
                "text-generation",
                "onnx-community/MobileLLM-R1-360M-ONNX",
                {
                    dtype: "fp32",
                    ...device
                }
            );

            this.setStatus('Model loaded successfully!', `Model: MobileLLM-R1-360M (${this.currentDevice.toUpperCase()})`);
            this.modelStatus.textContent = `Model: MobileLLM-R1-360M (${this.currentDevice.toUpperCase()})`;
            this.updateDeviceUI();
            this.showLoading(false);

        } catch (error) {
            console.error('Failed to load model:', error);
            this.setStatus('Failed to load model. Retrying with CPU...', 'Error');
            
            if (this.currentDevice === 'webgpu') {
                // Fallback to CPU
                this.currentDevice = 'cpu';
                this.deviceSelect.value = 'cpu';
                this.updateDeviceUI();
                this.loadModel();
            } else {
                this.showLoading(false);
                this.setStatus('Model loading failed. Please refresh the page.', 'Error');
                this.messagesContainer.innerHTML += `
                    <div class="error-message">
                        <div class="message system-message">
                            <div class="message-content">
                                <i class="fas fa-exclamation-triangle"></i>
                                <strong>Error:</strong> Failed to load AI model. Please check your internet connection and refresh the page.
                                <br><small>Details: ${error.message}</small>
                            </div>
                        </div>
                    </div>
                `;
            }
        }
    }

    async sendMessage() {
        if (this.isGenerating || !this.generator) {
            return;
        }

        const input = this.messageInput.value.trim();
        if (!input) return;

        // Add user message
        const userMessage = {
            role: "user",
            content: input
        };
        this.messages.push(userMessage);
        this.addMessage('user', input, true);
        this.messageInput.value = '';
        this.messageInput.style.height = 'auto';

        // Generate response
        await this.generateResponse(userMessage);
    }

    async generateResponse(userMessage) {
        this.isGenerating = true;
        this.showLoading(true);
        this.setStatus('AI is generating response...', 'Generating...');

        try {
            // Prepare messages for the model (just the latest user message for simplicity)
            const messages = [userMessage];

            const streamer = new TextStreamer(this.generator.tokenizer, {
                skip_prompt: true,
                skip_special_tokens: true,
            });

            const output = await this.generator(messages, {
                max_new_tokens: 200,
                do_sample: false,
                temperature: 0.7,
                streamer: streamer
            });

            // Get the generated response
            const responseContent = output[0].generated_text.at(-1).content;
            
            // Add assistant response
            this.addMessage('assistant', responseContent, false);
            this.messages.push({ role: "assistant", content: responseContent });

            this.setStatus('Response generated successfully!', 'Ready');
            this.isGenerating = false;
            this.showLoading(false);

        } catch (error) {
            console.error('Generation error:', error);
            this.addMessage('assistant', `I'm sorry, I encountered an error while generating a response: ${error.message}. Please try again.`, false);
            this.setStatus('Generation failed. Please try again.', 'Error');
            this.isGenerating = false;
            this.showLoading(false);
        }
    }

    addMessage(type, content, isUser = false) {
        const messageId = `msg-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
        const time = new Date().toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' });
        const icon = type === 'user' ? 'fa-user' : 'fa-robot';
        const author = type === 'user' ? 'You' : 'AI Assistant';
        const messageClass = type === 'user' ? 'user-message' : 'assistant-message';

        const messageHTML = `
            <div class="message ${messageClass}" role="log" id="${messageId}" aria-label="${author} message">
                <div class="message-header">
                    <i class="fas ${icon}"></i>
                    <span class="message-author">${author}</span>
                    <span class="message-time">${time}</span>
                </div>
                <div class="message-content">
                    ${this.formatMessageContent(content, type)}
                </div>
            </div>
        `;

        this.messagesContainer.insertAdjacentHTML('beforeend', messageHTML);
        this.messagesContainer.scrollTop = this.messagesContainer.scrollHeight;

        // Add math rendering if needed
        this.renderMath();
    }

    formatMessageContent(content, type) {
        // Basic markdown-like formatting
        let formatted = content
            .replace(/\*\*(.*?)\*\*/g, '<strong>$1</strong>')
            .replace(/\*(.*?)\*/g, '<em>$1</em>')
            .replace(/`(.*?)`/g, '<code>$1</code>')
            .replace(/\n/g, '<br>');

        // Math equation detection and formatting
        if (type === 'user' || content.includes('=') || content.match(/[a-zA-Z]\^[0-9]/)) {
            formatted = formatted.replace(
                /\$([^\$]+)\$/g,
                '<span class="math-equation">$1</span>'
            );
        }

        return formatted;
    }

    renderMath() {
        // Simple math rendering using MathJax or basic CSS
        const mathElements = this.messagesContainer.querySelectorAll('.math-equation');
        mathElements.forEach(el => {
            // Basic superscript/subscript formatting
            el.innerHTML = el.innerHTML
                .replace(/([a-zA-Z0-9])(\^[0-9]+)/g, '$1<sup>$2</sup>')
                .replace(/([a-zA-Z0-9])(_[0-9]+)/g, '$1<sub>$2</sub>');
        });
    }

    clearChat() {
        if (confirm('Clear all messages?')) {
            this.messages = [];
            this.messagesContainer.innerHTML = `
                <div class="welcome-message">
                    <div class="message assistant-message">
                        <div class="message-header">
                            <i class="fas fa-robot"></i>
                            <span class="message-author">AI Assistant</span>
                            <span class="message-time"></span>
                        </div>
                        <div class="message-content">
                            Chat cleared! How can I help you today? Try asking me to solve a math problem like "Solve x^2 - x - 6 = 0".
                        </div>
                    </div>
                </div>
            `;
            this.setStatus('Chat cleared', 'Ready');
        }
    }

    updateDeviceUI() {
        const icon = this.currentDevice === 'webgpu' ? 'fa-microchip' : 'fa-cpu';
        const text = this.currentDevice === 'webgpu' ? 'GPU' : 'CPU';
        this.deviceToggle.innerHTML = `<i class="fas ${icon}"></i> ${text}`;
        this.deviceToggle.title = `Switch to ${this.currentDevice === 'webgpu' ? 'CPU' : 'GPU (WebGPU)'}`;
    }

    setStatus(text, modelText = '') {
        this.statusText.textContent = text;
        if (modelText) {
            this.modelStatus.textContent = modelText;
        }
    }

    showLoading(show) {
        this.loadingOverlay.style.display = show ? 'flex' : 'none';
        if (!show) {
            this.sendButton.disabled = false;
            this.sendButton.innerHTML = '<i class="fas fa-paper-plane"></i>';
        } else {
            this.sendButton.disabled = true;
            this.sendButton.innerHTML = '<i class="fas fa-spinner fa-spin"></i>';
        }
    }
}

// Initialize the app when DOM is loaded
document.addEventListener('DOMContentLoaded', () => {
    new AIChatApp();
});

// Service Worker for offline capability (optional)
if ('serviceWorker' in navigator) {
    window.addEventListener('load', () => {
        navigator.serviceWorker.register('/sw.js').catch(() => {
            // Ignore service worker registration failure
        });
    });
}