Fix login flows

public/city-select.html

@@ -265,8 +265,15 @@
         }
 
         function selectCity(cityName) {
+            const urlParams = new URLSearchParams(window.location.search);
+            const fromDashboard = urlParams.get('from') === 'dashboard';
+            
+            if (user || fromDashboard) {
+                window.location.href = `/city/${cityName}/dashboard`;
+            } else {
                 window.location.href = `/city/${cityName}`;
             }
+        }
 
         function showMessage(message, type = 'success') {
             const messageDiv = document.getElementById('message');

public/dashboard.html

@@ -322,7 +322,7 @@
                 <h1>💧 Water Stations - <span id="cityName">Loading...</span></h1>
                 Welcome, <span id="username"></span>
             </div>
-            <a href="/city-select" class="city-selector">Change City</a>
+            <a href="/city-select?from=dashboard" class="city-selector">Change City</a>
             <button class="logout-btn" onclick="logout()">Logout</button>
         </div>
     </div>

public/index.html

@@ -133,8 +133,10 @@
         <p id="cityName">Loading...</p>
     </div>
     
-    <a href="/login" class="auth-button">Login</a>
+    <div id="authButtons">
-    <a href="/city-select" class="auth-button" style="top: 1rem; right: 6rem;">All Cities</a>
+        <a href="/login" class="auth-button" id="loginBtn">Login</a>
+        <a href="/city-select" class="auth-button" style="top: 1rem; right: 6rem;" id="changeCityBtn">All Cities</a>
+    </div>
     
     <div id="map"></div>
     
@@ -163,6 +165,36 @@
         let map;
         let stations = [];
         let currentCity = null;
+        let user = null;
+
+        async function checkAuth() {
+            try {
+                const response = await fetch('/api/user');
+                const data = await response.json();
+                user = data.user;
+                updateAuthButtons();
+            } catch (error) {
+                console.error('Auth check failed:', error);
+                updateAuthButtons();
+            }
+        }
+
+        function updateAuthButtons() {
+            const loginBtn = document.getElementById('loginBtn');
+            const changeCityBtn = document.getElementById('changeCityBtn');
+            
+            if (user) {
+                loginBtn.textContent = 'Dashboard';
+                loginBtn.href = `/city/${currentCity}/dashboard`;
+                changeCityBtn.textContent = 'Change City';
+                changeCityBtn.href = '/city-select';
+            } else {
+                loginBtn.textContent = 'Login';
+                loginBtn.href = `/login?redirect=${encodeURIComponent(window.location.pathname)}`;
+                changeCityBtn.textContent = 'All Cities';
+                changeCityBtn.href = '/city-select';
+            }
+        }
 
         function initMap() {
             // Get city name from URL
@@ -184,6 +216,7 @@
                 attribution: '© OpenStreetMap contributors'
             }).addTo(map);
 
+            checkAuth();
             loadStations();
         }
 
@@ -197,6 +230,7 @@
                     }
                     displayStations();
                     fitMapToStations();
+                    updateAuthButtons();
                 })
                 .catch(error => {
                     console.error('Error loading stations:', error);

public/login.html

@@ -211,8 +211,8 @@
                 <span>or</span>
             </div>
 
-            <a href="/auth/google" class="btn btn-google">Google</a>
+            <a href="/auth/google" class="btn btn-google" id="googleBtn">Google</a>
-            <a href="/auth/instagram" class="btn btn-instagram">Instagram</a>
+            <a href="/auth/instagram" class="btn btn-instagram" id="instagramBtn">Instagram</a>
 
             <div class="toggle-form">
                 <a href="#" onclick="toggleForm()">Don't have an account? Register</a>
@@ -240,8 +240,8 @@
                 <span>or</span>
             </div>
 
-            <a href="/auth/google" class="btn btn-google">Sign up with Google</a>
+            <a href="/auth/google" class="btn btn-google" id="googleRegBtn">Sign up with Google</a>
-            <a href="/auth/instagram" class="btn btn-instagram">Sign up with Instagram</a>
+            <a href="/auth/instagram" class="btn btn-instagram" id="instagramRegBtn">Sign up with Instagram</a>
 
             <div class="toggle-form">
                 <a href="#" onclick="toggleForm()">Already have an account? Login</a>
@@ -292,7 +292,18 @@
                 });
                 
                 if (response.ok) {
+                    const urlParams = new URLSearchParams(window.location.search);
+                    const redirectUrl = urlParams.get('redirect');
+                    if (redirectUrl) {
+                        const cityMatch = redirectUrl.match(/\/city\/([^\/]+)/);
+                        if (cityMatch) {
+                            window.location.href = `/city/${cityMatch[1]}/dashboard`;
+                        } else {
                             window.location.href = '/dashboard';
+                        }
+                    } else {
+                        window.location.href = '/dashboard';
+                    }
                 } else {
                     showMessage('Invalid username or password');
                 }
@@ -321,7 +332,18 @@
                 if (response.ok) {
                     showMessage('Registration successful! Redirecting...', 'success');
                     setTimeout(() => {
+                        const urlParams = new URLSearchParams(window.location.search);
+                        const redirectUrl = urlParams.get('redirect');
+                        if (redirectUrl) {
+                            const cityMatch = redirectUrl.match(/\/city\/([^\/]+)/);
+                            if (cityMatch) {
+                                window.location.href = `/city/${cityMatch[1]}/dashboard`;
+                            } else {
                                 window.location.href = '/dashboard';
+                            }
+                        } else {
+                            window.location.href = '/dashboard';
+                        }
                     }, 1500);
                 } else {
                     showMessage(result.error || 'Registration failed');
@@ -330,6 +352,20 @@
                 showMessage('Registration failed. Please try again.');
             }
         });
+
+        // Update OAuth links with redirect parameter
+        document.addEventListener('DOMContentLoaded', function() {
+            const urlParams = new URLSearchParams(window.location.search);
+            const redirectUrl = urlParams.get('redirect');
+            
+            if (redirectUrl) {
+                const encodedRedirect = encodeURIComponent(redirectUrl);
+                document.getElementById('googleBtn').href = `/auth/google?redirect=${encodedRedirect}`;
+                document.getElementById('instagramBtn').href = `/auth/instagram?redirect=${encodedRedirect}`;
+                document.getElementById('googleRegBtn').href = `/auth/google?redirect=${encodedRedirect}`;
+                document.getElementById('instagramRegBtn').href = `/auth/instagram?redirect=${encodedRedirect}`;
+            }
+        });
     </script>
 </body>
 </html>

server.js

@@ -185,25 +185,35 @@ app.get('/city-select', (req, res) => {
 });
 
 // Authentication routes
-app.get('/auth/google',
+app.get('/auth/google', (req, res, next) => {
-  passport.authenticate('google', { scope: ['profile', 'email'] })
+  if (req.query.redirect) {
-);
+    req.session.redirectUrl = req.query.redirect;
+  }
+  passport.authenticate('google', { scope: ['profile', 'email'] })(req, res, next);
+});
 
 app.get('/auth/google/callback',
   passport.authenticate('google', { failureRedirect: '/login' }),
   (req, res) => {
-    res.redirect('/city/salem/dashboard');
+    const redirectUrl = req.session.redirectUrl || '/city/salem/dashboard';
+    delete req.session.redirectUrl;
+    res.redirect(redirectUrl);
   }
 );
 
-app.get('/auth/instagram',
+app.get('/auth/instagram', (req, res, next) => {
-  passport.authenticate('instagram')
+  if (req.query.redirect) {
-);
+    req.session.redirectUrl = req.query.redirect;
+  }
+  passport.authenticate('instagram')(req, res, next);
+});
 
 app.get('/auth/instagram/callback',
   passport.authenticate('instagram', { failureRedirect: '/login' }),
   (req, res) => {
-    res.redirect('/city/salem/dashboard');
+    const redirectUrl = req.session.redirectUrl || '/city/salem/dashboard';
+    delete req.session.redirectUrl;
+    res.redirect(redirectUrl);
   }
 );
 
@@ -420,6 +430,9 @@ app.get('/city/:cityName/dashboard', (req, res) => {
 });
 
 app.get('/login', (req, res) => {
+  if (req.query.redirect) {
+    req.session.redirectUrl = req.query.redirect;
+  }
   res.sendFile(path.join(__dirname, 'public', 'login.html'));
 });