Update static/script.js

static/script.js

@@ -4,12 +4,14 @@
     let currentUser = null;
     let token = localStorage.getItem('token');
     let ws = null;
-    let messages = {};           // id -> msg object
-    let messageIds = [];         // ordered
+    let wsReconnectTimer = null;
+    let messages = {};
+    let messageIds = [];
     let oldestMessageId = null;
     let replyToId = null;
     let typingTimer = null;
-    let pendingFile = null;     // {file_path, file_type, file_name}
+    let pendingFile = null;
+    let isConnected = false;
 
     const API = window.location.origin;
 
@@ -70,6 +72,7 @@
             document.documentElement.style.setProperty('--bubble-text-other', txt);
         }
     }
+    
     function loadTheme() {
         const saved = localStorage.getItem('theme') || 'theme-light';
         themeSelect.value = saved;
@@ -78,39 +81,91 @@
     }
 
     // --- Auth ---
-    loginTab.addEventListener('click', ()=>{
-        isLogin=true; loginTab.classList.add('active'); signupTab.classList.remove('active');
-        displayNameGroup.style.display='none'; authSubmit.textContent='Login';
+    loginTab.addEventListener('click', () => {
+        isLogin = true;
+        loginTab.classList.add('active');
+        signupTab.classList.remove('active');
+        displayNameGroup.style.display = 'none';
+        authSubmit.textContent = 'Login';
     });
-    signupTab.addEventListener('click', ()=>{
-        isLogin=false; signupTab.classList.add('active'); loginTab.classList.remove('active');
-        displayNameGroup.style.display='block'; authSubmit.textContent='Sign Up';
+    
+    signupTab.addEventListener('click', () => {
+        isLogin = false;
+        signupTab.classList.add('active');
+        loginTab.classList.remove('active');
+        displayNameGroup.style.display = 'block';
+        authSubmit.textContent = 'Sign Up';
     });
-    authForm.addEventListener('submit', async (e)=>{
+    
+    authForm.addEventListener('submit', async (e) => {
         e.preventDefault();
-        const u = authUsername.value.trim(), p = authPassword.value, d = authDisplayName.value.trim();
-        if(!u||!p){ authError.textContent='Fill all fields.'; return; }
-        try{
-            const url = isLogin ? `${API}/login?username=${encodeURIComponent(u)}&password=${encodeURIComponent(p)}`
-                                : `${API}/signup?username=${encodeURIComponent(u)}&password=${encodeURIComponent(p)}&display_name=${encodeURIComponent(d)}`;
-            const res = await fetch(url, {method:'POST'});
-            if(!res.ok) throw new Error((await res.json()).detail||'Auth failed');
+        const u = authUsername.value.trim();
+        const p = authPassword.value;
+        const d = authDisplayName.value.trim();
+        
+        if (!u || !p) {
+            authError.textContent = 'Please fill all fields.';
+            return;
+        }
+        
+        try {
+            const url = isLogin 
+                ? `${API}/login?username=${encodeURIComponent(u)}&password=${encodeURIComponent(p)}`
+                : `${API}/signup?username=${encodeURIComponent(u)}&password=${encodeURIComponent(p)}&display_name=${encodeURIComponent(d)}`;
+            
+            const res = await fetch(url, { method: 'POST' });
+            
+            if (!res.ok) {
+                const err = await res.json();
+                throw new Error(err.detail || 'Authentication failed');
+            }
+            
             const data = await res.json();
             token = data.token;
-            currentUser = { id: data.id, username: data.username, display_name: data.display_name, avatar_path: data.avatar_path };
+            currentUser = {
+                id: data.id,
+                username: data.username,
+                display_name: data.display_name,
+                avatar_path: data.avatar_path
+            };
+            
             localStorage.setItem('token', token);
+            authError.textContent = '';
             showChat();
-        }catch(err){ authError.textContent=err.message; }
+        } catch (err) {
+            authError.textContent = err.message;
+        }
     });
 
     async function autoLogin() {
-        if(!token) return false;
-        try{
-            const res = await fetch(`${API}/me`, {headers:{token}});
-            if(!res.ok){ token=null; localStorage.removeItem('token'); return false; }
+        if (!token) {
+            return false;
+        }
+        
+        try {
+            const res = await fetch(`${API}/me`, {
+                headers: { 'token': token }
+            });
+            
+            if (!res.ok) {
+                token = null;
+                localStorage.removeItem('token');
+                return false;
+            }
+            
             currentUser = await res.json();
             return true;
-        }catch{ return false; }
+        } catch (e) {
+            token = null;
+            localStorage.removeItem('token');
+            return false;
+        }
+    }
+
+    function showAuth() {
+        authScreen.classList.remove('hidden');
+        chatScreen.classList.add('hidden');
+        disconnectWS();
     }
 
     function showChat() {
@@ -120,59 +175,131 @@
         loadMessages();
     }
 
-    // --- WebSocket ---
+    // --- WebSocket (fixed – never send null as token) ---
     function connectWS() {
-        if(ws && ws.readyState===WebSocket.OPEN) return;
-        const proto = location.protocol==='https:'?'wss:':'ws:';
-        ws = new WebSocket(`${proto}//${location.host}/ws?token=${encodeURIComponent(token)}`);
-        ws.onopen = ()=>{ console.log('WS open'); };
-        ws.onmessage = (ev)=> handleWSMessage(JSON.parse(ev.data));
-        ws.onclose = ()=> setTimeout(connectWS, 3000);
-        ws.onerror = ()=> {};
+        // Don't connect if we don't have a valid token
+        if (!token || token === 'null' || token === 'undefined') {
+            console.warn('No valid token – not connecting WebSocket');
+            return;
+        }
+        
+        // Don't create duplicate connections
+        if (ws && ws.readyState === WebSocket.OPEN) {
+            return;
+        }
+        
+        // Close any existing connection
+        if (ws) {
+            ws.close();
+            ws = null;
+        }
+        
+        const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
+        const wsUrl = `${proto}//${location.host}/ws?token=${encodeURIComponent(token)}`;
+        
+        console.log('Connecting WebSocket...');
+        ws = new WebSocket(wsUrl);
+        
+        ws.onopen = () => {
+            console.log('WebSocket connected');
+            isConnected = true;
+            // Clear any reconnect timer
+            if (wsReconnectTimer) {
+                clearTimeout(wsReconnectTimer);
+                wsReconnectTimer = null;
+            }
+            // Request initial message load
+            ws.send(JSON.stringify({ type: 'load_messages', limit: 50 }));
+        };
+        
+        ws.onmessage = (ev) => {
+            try {
+                handleWSMessage(JSON.parse(ev.data));
+            } catch (e) {
+                console.error('Error parsing WS message:', e);
+            }
+        };
+        
+        ws.onclose = () => {
+            console.log('WebSocket closed');
+            isConnected = false;
+            ws = null;
+            // Only reconnect if we have a valid token
+            if (token && token !== 'null' && token !== 'undefined' && currentUser) {
+                wsReconnectTimer = setTimeout(connectWS, 3000);
+            }
+        };
+        
+        ws.onerror = (err) => {
+            console.error('WebSocket error:', err);
+        };
+    }
+
+    function disconnectWS() {
+        if (wsReconnectTimer) {
+            clearTimeout(wsReconnectTimer);
+            wsReconnectTimer = null;
+        }
+        if (ws) {
+            ws.close();
+            ws = null;
+        }
+        isConnected = false;
     }
 
     function handleWSMessage(msg) {
-        switch(msg.type){
-            case 'online_users': onlineCountSpan.textContent = `(${msg.users.length} online)`; break;
-            case 'user_online': break;
-            case 'user_offline': break;
+        switch (msg.type) {
+            case 'online_users':
+                onlineCountSpan.textContent = `(${msg.users.length} online)`;
+                break;
             case 'new_message':
-                if(messages[msg.id]) return;
+                if (messages[msg.id]) return;
                 messages[msg.id] = msg;
                 messageIds.push(msg.id);
                 renderMessage(msg);
                 scrollBottom();
-                if(document.visibilityState==='visible') ws.send(JSON.stringify({type:'mark_read', up_to_message_id:msg.id}));
+                // Mark as read if visible
+                if (document.visibilityState === 'visible' && isConnected) {
+                    ws.send(JSON.stringify({ type: 'mark_read', up_to_message_id: msg.id }));
+                }
+                // Browser notification
+                if (document.visibilityState !== 'visible' && Notification.permission === 'granted') {
+                    new Notification(`New message from ${msg.display_name || msg.username}`, {
+                        body: msg.content || 'Sent a file'
+                    });
+                }
                 break;
             case 'messages_batch':
-                if(msg.messages.length>0){
-                    for(const m of msg.messages){
-                        if(!messages[m.id]){ messages[m.id]=m; messageIds.unshift(m.id); }
+                if (msg.messages.length > 0) {
+                    for (const m of msg.messages) {
+                        if (!messages[m.id]) {
+                            messages[m.id] = m;
+                            messageIds.unshift(m.id);
+                        }
                     }
                     renderAll();
                     scrollBottom();
                 }
-                if(msg.messages.length<50) loadOlderBtn.classList.add('hidden');
-                else loadOlderBtn.classList.remove('hidden');
+                loadOlderBtn.classList.toggle('hidden', msg.messages.length < 50);
                 break;
             case 'typing_update':
-                typingIndicator.textContent = msg.is_typing ? `${msg.username} is typing…` : ''; break;
+                typingIndicator.textContent = msg.is_typing ? `${msg.username} is typing...` : '';
+                break;
             case 'edit_message':
-                if(messages[msg.message_id]){
+                if (messages[msg.message_id]) {
                     messages[msg.message_id].content = msg.content;
                     messages[msg.message_id].is_edited = true;
                     updateBubble(msg.message_id);
                 }
                 break;
             case 'delete_message':
-                if(messages[msg.message_id]){
+                if (messages[msg.message_id]) {
                     messages[msg.message_id].is_deleted = true;
                     updateBubble(msg.message_id);
                 }
                 break;
             case 'read_receipt':
-                if(messages[msg.message_id] && messages[msg.message_id].sender_id === currentUser.id){
-                    // Update status to "read"
+                if (messages[msg.message_id] && messages[msg.message_id].sender_id === currentUser?.id) {
                     messages[msg.message_id].status = 'read';
                     updateBubble(msg.message_id);
                 }
@@ -180,173 +307,281 @@
         }
     }
 
-    // --- Render ---
+    // --- Message rendering (same as before) ---
     function renderMessage(msg) {
         const el = document.getElementById(`msg-${msg.id}`);
-        if(el){
+        if (el) {
             updateBubble(msg.id);
             return;
         }
+        
         const div = document.createElement('div');
         div.id = `msg-${msg.id}`;
-        div.className = `message-bubble ${msg.sender_id===currentUser.id?'message-self':'message-other'}`;
+        div.className = `message-bubble ${msg.sender_id === currentUser?.id ? 'message-self' : 'message-other'}`;
         div.innerHTML = buildHTML(msg);
         messagesList.appendChild(div);
     }
+
     function buildHTML(msg) {
-        if(msg.is_deleted) return `<em>This message was deleted</em>`;
+        if (msg.is_deleted) return '<em>This message was deleted</em>';
+        
         let html = '';
-        if(msg.reply_to_id && messages[msg.reply_to_id]){
+        if (msg.reply_to_id && messages[msg.reply_to_id]) {
             const reply = messages[msg.reply_to_id];
-            html += `<div class="reply-indicator">↩ ${reply.display_name||reply.username}: ${truncate(reply.content,40)}</div>`;
+            html += `<div class="reply-indicator">↩ ${reply.display_name || reply.username}: ${truncate(reply.content, 40)}</div>`;
         }
-        html += `<div class="message-info"><span>${msg.display_name||msg.username}</span><span class="timestamp">${formatTime(msg.timestamp_ms)}</span></div>`;
+        
+        html += `<div class="message-info">
+                    <span>${msg.display_name || msg.username}</span>
+                    <span class="timestamp">${formatTime(msg.timestamp_ms)}</span>
+                 </div>`;
         html += `<div class="message-text">${escapeHtml(msg.content)}</div>`;
-        if(msg.is_edited) html += ` <em class="edited-label">[edited]</em>`;
-        if(msg.file_path) html += renderAttachment(msg);
-        if(msg.sender_id===currentUser.id && !msg.is_deleted){
-            const statusIcon = msg.status==='read' ? '✓✓' : '✓';
-            html += `<div class="read-status double-check">${statusIcon}</div>`;
+        
+        if (msg.is_edited) html += ` <em class="edited-label">[edited]</em>`;
+        if (msg.file_path) html += renderAttachment(msg);
+        
+        // Status indicators for own messages
+        if (msg.sender_id === currentUser?.id && !msg.is_deleted) {
+            const statusIcon = msg.status === 'read' ? '✓✓' : '✓';
+            html += `<div class="message-status">${statusIcon}</div>`;
             html += `<div class="message-actions">
                         <button onclick="window.editMsg(${msg.id})">✏️</button>
                         <button onclick="window.deleteMsg(${msg.id})">🗑️</button>
                         <button onclick="window.replyTo(${msg.id})">↩</button>
                      </div>`;
         }
+        
         return html;
     }
+
     function renderAttachment(msg) {
         const url = `${API}/download/${msg.file_path}?token=${encodeURIComponent(token)}`;
-        if(msg.file_type.startsWith('image/')) return `<div class="attachment"><img src="${url}" alt="${msg.file_name}" onerror="this.parentElement.innerHTML='Failed to decode'"></div>`;
-        if(msg.file_type.startsWith('video/')) return `<div class="attachment"><video controls><source src="${url}" type="${msg.file_type}"></video></div>`;
-        if(msg.file_type.startsWith('audio/')) return `<div class="attachment"><audio controls src="${url}"></audio></div>`;
+        
+        if (msg.file_type.startsWith('image/')) {
+            return `<div class="attachment"><img src="${url}" alt="${msg.file_name}" onerror="this.parentElement.innerHTML='<em>Failed to decode</em>'"></div>`;
+        }
+        if (msg.file_type.startsWith('video/')) {
+            return `<div class="attachment"><video controls><source src="${url}" type="${msg.file_type}">Your browser doesn't support video</video></div>`;
+        }
+        if (msg.file_type.startsWith('audio/')) {
+            return `<div class="attachment"><audio controls src="${url}"></audio></div>`;
+        }
         return `<div class="attachment"><div class="file-preview" onclick="window.open('${url}')">📄 ${escapeHtml(msg.file_name)}</div></div>`;
     }
+
     function updateBubble(id) {
         const el = document.getElementById(`msg-${id}`);
-        if(!el) return;
+        if (!el) return;
         const msg = messages[id];
-        el.className = `message-bubble ${msg.sender_id===currentUser.id?'message-self':'message-other'}`;
-        if(msg.is_deleted) el.classList.add('deleted-bubble');
+        if (!msg) return;
+        
+        el.className = `message-bubble ${msg.sender_id === currentUser?.id ? 'message-self' : 'message-other'}`;
+        if (msg.is_deleted) el.classList.add('deleted-bubble');
         el.innerHTML = buildHTML(msg);
     }
+
     function renderAll() {
         messagesList.innerHTML = '';
-        messageIds.forEach(id => { if(messages[id]) { renderMessage(messages[id]); } });
+        messageIds.forEach(id => {
+            if (messages[id]) renderMessage(messages[id]);
+        });
     }
+
     function scrollBottom() {
-        document.getElementById('messages-container').scrollTop = document.getElementById('messages-container').scrollHeight;
+        const container = document.getElementById('messages-container');
+        container.scrollTop = container.scrollHeight;
     }
 
-    // --- Actions ---
-    window.replyTo = function(id){ replyToId=id; const m=messages[id]; if(m){ replyPreviewText.textContent=`Replying to ${m.display_name||m.username}: ${truncate(m.content,30)}`; replyPreview.classList.remove('hidden'); } };
-    cancelReply.addEventListener('click', ()=>{ replyToId=null; replyPreview.classList.add('hidden'); });
-    window.editMsg = function(id){
-        const newContent = prompt('Edit:', messages[id]?.content);
-        if(newContent!==null && newContent.trim()!=='') ws.send(JSON.stringify({type:'edit_message', message_id:id, content:newContent}));
+    // --- Message actions ---
+    window.replyTo = function(id) {
+        replyToId = id;
+        const m = messages[id];
+        if (m) {
+            replyPreviewText.textContent = `Replying to ${m.display_name || m.username}: ${truncate(m.content, 30)}`;
+            replyPreview.classList.remove('hidden');
+        }
+    };
+    
+    cancelReply.addEventListener('click', () => {
+        replyToId = null;
+        replyPreview.classList.add('hidden');
+    });
+    
+    window.editMsg = function(id) {
+        const newContent = prompt('Edit message:', messages[id]?.content || '');
+        if (newContent !== null && newContent.trim() !== '' && isConnected) {
+            ws.send(JSON.stringify({ type: 'edit_message', message_id: id, content: newContent }));
+        }
     };
-    window.deleteMsg = function(id){
-        if(confirm('Delete permanently?')) ws.send(JSON.stringify({type:'delete_message', message_id:id}));
+    
+    window.deleteMsg = function(id) {
+        if (confirm('Delete this message permanently?') && isConnected) {
+            ws.send(JSON.stringify({ type: 'delete_message', message_id: id }));
+        }
     };
 
-    // --- Sending ---
-    composeForm.addEventListener('submit', (e)=>{
+    // --- Sending messages ---
+    composeForm.addEventListener('submit', (e) => {
         e.preventDefault();
+        if (!isConnected) {
+            console.warn('Not connected – cannot send');
+            return;
+        }
+        
         const content = messageInput.value.trim();
-        if(!content && !pendingFile) return;
-        const msg = { type:'chat_message', content, reply_to_id:replyToId, client_id:Date.now().toString(36)+Math.random().toString(36).slice(2) };
-        if(pendingFile){
+        if (!content && !pendingFile) return;
+        
+        const msg = {
+            type: 'chat_message',
+            content: content || '',
+            reply_to_id: replyToId,
+            client_id: Date.now().toString(36) + Math.random().toString(36).slice(2)
+        };
+        
+        if (pendingFile) {
             msg.file_path = pendingFile.file_path;
             msg.file_type = pendingFile.file_type;
             msg.file_name = pendingFile.file_name;
             pendingFile = null;
         }
+        
         ws.send(JSON.stringify(msg));
         messageInput.value = '';
         clearReply();
     });
-    function clearReply(){ replyToId=null; replyPreview.classList.add('hidden'); }
 
-    // -- Typing --
-    messageInput.addEventListener('input', ()=>{
-        ws.send(JSON.stringify({type:'typing', is_typing:true}));
+    function clearReply() {
+        replyToId = null;
+        replyPreview.classList.add('hidden');
+    }
+
+    // --- Typing indicator ---
+    messageInput.addEventListener('input', () => {
+        if (!isConnected) return;
+        ws.send(JSON.stringify({ type: 'typing', is_typing: true }));
         clearTimeout(typingTimer);
-        typingTimer = setTimeout(()=> ws.send(JSON.stringify({type:'typing', is_typing:false})), 2000);
+        typingTimer = setTimeout(() => {
+            if (isConnected) ws.send(JSON.stringify({ type: 'typing', is_typing: false }));
+        }, 2000);
     });
 
-    // -- File upload (chunked) --
-    attachBtn.addEventListener('click', ()=> fileInput.click());
-    fileInput.addEventListener('change', async (e)=>{
+    // --- File upload ---
+    attachBtn.addEventListener('click', () => fileInput.click());
+    
+    fileInput.addEventListener('change', async (e) => {
         const file = e.target.files[0];
-        if(!file) return;
-        const CHUNK = 5*1024*1024;
+        if (!file) return;
+        
+        const CHUNK = 5 * 1024 * 1024; // 5MB
         const total = Math.ceil(file.size / CHUNK);
-        const uploadId = Date.now().toString(36)+Math.random().toString(36).slice(2);
+        const uploadId = Date.now().toString(36) + Math.random().toString(36).slice(2);
+        
         uploadProgress.classList.remove('hidden');
         uploadFilename.textContent = file.name;
-        uploadBar.value=0; uploadBar.max=100;
-        for(let i=0;i<total;i++){
-            const start = i*CHUNK, end = Math.min(start+CHUNK, file.size);
-            const blob = file.slice(start,end);
-            const fd = new FormData(); fd.append('file', blob);
-            const url = `${API}/upload_chunk?chunk_index=${i}&total_chunks=${total}&file_name=${encodeURIComponent(file.name)}&file_type=${encodeURIComponent(file.type||'application/octet-stream')}&upload_id=${uploadId}`;
-            try{
-                const res = await fetch(url, {method:'POST', body:fd, headers:{token}});
-                if(!res.ok) throw new Error('Upload failed');
+        uploadBar.value = 0;
+        uploadBar.max = 100;
+        
+        for (let i = 0; i < total; i++) {
+            const start = i * CHUNK;
+            const end = Math.min(start + CHUNK, file.size);
+            const blob = file.slice(start, end);
+            
+            const fd = new FormData();
+            fd.append('file', blob);
+            
+            const url = `${API}/upload_chunk?chunk_index=${i}&total_chunks=${total}&file_name=${encodeURIComponent(file.name)}&file_type=${encodeURIComponent(file.type || 'application/octet-stream')}&upload_id=${uploadId}`;
+            
+            try {
+                const res = await fetch(url, {
+                    method: 'POST',
+                    body: fd,
+                    headers: { 'token': token }
+                });
+                
+                if (!res.ok) throw new Error('Upload failed');
+                
                 const data = await res.json();
-                if(data.status==='complete'){
-                    pendingFile = {file_path:data.file_path, file_type:data.file_type, file_name:data.file_name};
-                    alert('File ready! Add a message and send.');
+                
+                if (data.status === 'complete') {
+                    pendingFile = {
+                        file_path: data.file_path,
+                        file_type: data.file_type,
+                        file_name: data.file_name
+                    };
                     uploadProgress.classList.add('hidden');
                     break;
                 }
-                uploadBar.value = ((i+1)/total)*100;
-            }catch(err){
-                alert('Upload error: '+err.message);
+                
+                uploadBar.value = ((i + 1) / total) * 100;
+            } catch (err) {
+                alert('Upload error: ' + err.message);
                 uploadProgress.classList.add('hidden');
                 break;
             }
         }
-        fileInput.value='';
+        
+        fileInput.value = '';
     });
 
-    // -- Load older --
-    loadOlderBtn.addEventListener('click', ()=>{
-        if(ws) ws.send(JSON.stringify({type:'load_messages', before_id: oldestMessageId, limit:50}));
+    // --- Load older messages ---
+    loadOlderBtn.addEventListener('click', () => {
+        if (isConnected && messageIds.length > 0) {
+            const oldestId = messageIds[0];
+            ws.send(JSON.stringify({ type: 'load_messages', before_id: oldestId, limit: 50 }));
+        }
     });
 
-    function loadMessages(){
-        if(ws && ws.readyState===WebSocket.OPEN) ws.send(JSON.stringify({type:'load_messages', limit:50}));
+    function loadMessages() {
+        if (isConnected) {
+            ws.send(JSON.stringify({ type: 'load_messages', limit: 50 }));
+        }
     }
 
     // --- Settings ---
-    settingsBtn.addEventListener('click', ()=>{
+    settingsBtn.addEventListener('click', () => {
         settingDisplayName.value = currentUser?.display_name || '';
-        if(currentUser?.avatar_path){
-            avatarPreview.src = `${API}/download/${currentUser.avatar_path}?token=${token}`;
-            avatarPreview.style.display='block';
-        }else avatarPreview.style.display='none';
+        if (currentUser?.avatar_path) {
+            avatarPreview.src = `${API}/download/${currentUser.avatar_path}?token=${encodeURIComponent(token)}`;
+            avatarPreview.style.display = 'block';
+        } else {
+            avatarPreview.style.display = 'none';
+        }
         settingsModal.classList.remove('hidden');
     });
-    closeSettingsBtn.addEventListener('click', ()=> settingsModal.classList.add('hidden'));
-    saveSettingsBtn.addEventListener('click', async ()=>{
+    
+    closeSettingsBtn.addEventListener('click', () => settingsModal.classList.add('hidden'));
+    
+    saveSettingsBtn.addEventListener('click', async () => {
         const newName = settingDisplayName.value.trim();
-        if(newName !== currentUser.display_name){
-            await fetch(`${API}/update_profile?display_name=${encodeURIComponent(newName)}`, {method:'POST', headers:{token}});
+        if (newName && newName !== currentUser?.display_name) {
+            await fetch(`${API}/update_profile?display_name=${encodeURIComponent(newName)}`, {
+                method: 'POST',
+                headers: { 'token': token }
+            });
             currentUser.display_name = newName;
         }
+        
         const avatarFile = settingAvatar.files[0];
-        if(avatarFile){
-            const fd = new FormData(); fd.append('file', avatarFile);
-            const res = await fetch(`${API}/upload_avatar`, {method:'POST', body:fd, headers:{token}});
-            if(res.ok){
+        if (avatarFile) {
+            const fd = new FormData();
+            fd.append('file', avatarFile);
+            
+            const res = await fetch(`${API}/upload_avatar`, {
+                method: 'POST',
+                body: fd,
+                headers: { 'token': token }
+            });
+            
+            if (res.ok) {
                 const data = await res.json();
                 currentUser.avatar_path = data.avatar_path;
-                avatarPreview.src = `${API}/download/${data.avatar_path}?token=${token}`;
+                avatarPreview.src = `${API}/download/${data.avatar_path}?token=${encodeURIComponent(token)}`;
+                avatarPreview.style.display = 'block';
             }
         }
+        
         const theme = themeSelect.value;
-        if(theme==='theme-custom'){
+        if (theme === 'theme-custom') {
             localStorage.setItem('customSelfColor', customSelfColor.value);
             localStorage.setItem('customOtherColor', customOtherColor.value);
             localStorage.setItem('customTextColor', customTextColor.value);
@@ -354,34 +589,69 @@
         applyTheme(theme);
         settingsModal.classList.add('hidden');
     });
-    themeSelect.addEventListener('change', (e)=>{
-        if(e.target.value==='theme-custom') customThemeSettings.classList.remove('hidden');
-        else customThemeSettings.classList.add('hidden');
+    
+    themeSelect.addEventListener('change', (e) => {
+        customThemeSettings.classList.toggle('hidden', e.target.value !== 'theme-custom');
     });
 
-    logoutBtn.addEventListener('click', async ()=>{
-        if(ws) ws.close();
-        if(token) await fetch(`${API}/logout?token=${encodeURIComponent(token)}`, {method:'GET'});
+    // --- Logout ---
+    logoutBtn.addEventListener('click', async () => {
+        if (token) {
+            try {
+                await fetch(`${API}/logout?token=${encodeURIComponent(token)}`, { method: 'GET' });
+            } catch (e) {
+                console.error('Logout error:', e);
+            }
+        }
+        
+        disconnectWS();
         localStorage.removeItem('token');
-        token=null; currentUser=null; messages={}; messageIds=[];
-        authScreen.classList.remove('hidden'); chatScreen.classList.add('hidden');
+        token = null;
+        currentUser = null;
+        messages = {};
+        messageIds = [];
+        messagesList.innerHTML = '';
+        typingIndicator.textContent = '';
+        onlineCountSpan.textContent = '0';
+        
+        showAuth();
     });
 
     // --- Notification permission ---
-    if(window.Notification && Notification.permission==='default') Notification.requestPermission();
+    if (window.Notification && Notification.permission === 'default') {
+        Notification.requestPermission();
+    }
+
+    // --- Utilities ---
+    function truncate(s, n) {
+        if (!s) return '';
+        return s.length > n ? s.slice(0, n - 1) + '…' : s;
+    }
 
-    // --- Utils ---
-    function truncate(s,n){ return s.length>n ? s.slice(0,n-1)+'…' : s; }
-    function escapeHtml(t){ const d=document.createElement('div'); d.textContent=t; return d.innerHTML; }
-    function formatTime(ms){
-        const d = new Date(ms), now = new Date();
-        if(d.toDateString()===now.toDateString()) return `Today at ${d.toLocaleTimeString([],{hour:'2-digit',minute:'2-digit'})}`;
-        return d.toLocaleDateString() +' '+ d.toLocaleTimeString([],{hour:'2-digit',minute:'2-digit'});
+    function escapeHtml(text) {
+        const div = document.createElement('div');
+        div.textContent = text;
+        return div.innerHTML;
     }
 
-    // -- Init --
+    function formatTime(ms) {
+        const d = new Date(ms);
+        const now = new Date();
+        if (d.toDateString() === now.toDateString()) {
+            return `Today at ${d.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}`;
+        }
+        return d.toLocaleDateString() + ' ' + d.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' });
+    }
+
+    // --- Initialize ---
     loadTheme();
-    (async ()=>{
-        if(token && await autoLogin()) showChat();
+    
+    (async () => {
+        const loggedIn = await autoLogin();
+        if (loggedIn) {
+            showChat();
+        } else {
+            showAuth();
+        }
     })();
 })();