static/script.js

// static/script.js
document.addEventListener('DOMContentLoaded', function() {
    const messageForm = document.getElementById('message-form');
    const userInput = document.getElementById('user-input');
    const chatMessages = document.getElementById('chat-messages');
    let websocket;

    // Add welcome message
    addMessage('Welcome! I can help you search the web, calculate expressions, and generate image prompts. What would you like to do?', 'ai');

    // Connect WebSocket
    function connectWebSocket() {
        // Get the correct WebSocket URL based on the current page URL
        const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
        const wsUrl = `${protocol}//${window.location.host}/ws`;
        
        websocket = new WebSocket(wsUrl);
        
        websocket.onopen = function(event) {
            console.log('WebSocket connected');
        };
        
        websocket.onmessage = function(event) {
            const data = JSON.parse(event.data);
            
            if (data.type === 'ai_message') {
                addMessage(data.content, 'ai');
            } else if (data.type === 'tool_message') {
                addMessage(data.content, 'tool');
            }
        };
        
        websocket.onclose = function(event) {
            console.log('WebSocket disconnected');
            // Try to reconnect after a delay
            setTimeout(connectWebSocket, 3000);
        };
        
        websocket.onerror = function(error) {
            console.error('WebSocket error:', error);
        };
    }

    // Connect to WebSocket when page loads
    connectWebSocket();

    // Handle message submission
    messageForm.addEventListener('submit', function(event) {
        event.preventDefault();
        const message = userInput.value.trim();
        
        if (message) {
            // Add user message to chat
            addMessage(message, 'user');
            
            // Send message to server via WebSocket
            if (websocket && websocket.readyState === WebSocket.OPEN) {
                websocket.send(message);
                
                // Add a "thinking" message
                const thinkingId = addThinkingMessage();
                
                // Clear input
                userInput.value = '';
                
                // Focus on input
                userInput.focus();
            } else {
                addMessage('Connection lost. Trying to reconnect...', 'system');
                connectWebSocket();
            }
        }
    });

    // Function to add a message to the chat
    function addMessage(content, type) {
        // Remove any thinking message first
        removeThinkingMessage();
        
        const messageDiv = document.createElement('div');
        messageDiv.classList.add('message');
        
        switch (type) {
            case 'user':
                messageDiv.classList.add('user-message');
                break;
            case 'ai':
                messageDiv.classList.add('ai-message');
                break;
            case 'tool':
                messageDiv.classList.add('tool-message');
                break;
            default:
                messageDiv.classList.add('system-message');
        }
        
        // Process markdown-like code blocks
        content = content.replace(/```([\s\S]*?)```/g, '<pre><code>$1</code></pre>');
        
        messageDiv.innerHTML = content;
        chatMessages.appendChild(messageDiv);
        
        // Scroll to bottom
        chatMessages.scrollTop = chatMessages.scrollHeight;
    }

    // Function to add a "thinking" message
    function addThinkingMessage() {
        const thinkingDiv = document.createElement('div');
        thinkingDiv.classList.add('message', 'ai-message', 'thinking');
        thinkingDiv.id = 'thinking-message';
        thinkingDiv.textContent = 'Thinking...';
        chatMessages.appendChild(thinkingDiv);
        
        // Scroll to bottom
        chatMessages.scrollTop = chatMessages.scrollHeight;
        
        return 'thinking-message';
    }

    // Function to remove the "thinking" message
    function removeThinkingMessage() {
        const thinkingMessage = document.getElementById('thinking-message');
        if (thinkingMessage) {
            thinkingMessage.remove();
        }
    }
});