Merge pull request #6 from mattintech/feature/optimization

Feature/optimization

TODO.md

@@ -15,21 +15,21 @@
 - [x] Convert singletons to proper DI
 - [x] Inject ViewModels using Hilt
 
-### 1.3 Room Database Setup
+### 1.3 Room Database Setup ✅
-- [ ] Add Room dependencies
+- [x] Add Room dependencies
-- [ ] Create Message and User entities
+- [x] Create Message and User entities
-- [ ] Implement DAOs for data access
+- [x] Implement DAOs for data access
-- [ ] Create database migrations
+- [x] Create database migrations
-- [ ] Store messages in Room database
+- [x] Store messages in Room database
-- [ ] Load message history on app restart
+- [x] Load message history on app restart
-- [ ] Implement message sync logic
+- [x] Implement message sync logic
 
-### 1.4 Coroutines & Flow Optimization
+### 1.4 Coroutines & Flow Optimization ✅
-- [ ] Convert callbacks to coroutines
+- [x] Convert callbacks to coroutines
-- [ ] Use Flow for reactive data streams
+- [x] Use Flow for reactive data streams
-- [ ] Implement proper scope management
+- [x] Implement proper scope management
-- [ ] Replace GlobalScope with proper lifecycle scopes
+- [x] Replace GlobalScope with proper lifecycle scopes
-- [ ] Add proper error handling with coroutines
+- [x] Add proper error handling with coroutines
 
 ## Phase 2: Core UX Improvements
 
@@ -47,6 +47,16 @@
 
 ### 2.3 Enhanced Messaging Features
 - [ ] Message status indicators (sent/delivered/read)
+  - [ ] Add status field to MessageEntity (pending/sent/delivered/failed)
+  - [ ] Show status icons in message bubbles
+  - [ ] Update status when delivery confirmed
+- [ ] Store-and-forward messaging pattern
+  - [ ] Save messages with "pending" status initially
+  - [ ] Implement acknowledgment protocol in WifiAwareManager
+  - [ ] Update to "sent" only after confirmation received
+  - [ ] Queue messages when offline/disconnected
+  - [ ] Auto-retry failed messages with exponential backoff
+  - [ ] Mark messages as failed after max retries
 - [ ] User presence indicators (online/offline/typing)
 - [ ] Message timestamps with proper formatting
 - [ ] Offline message queue
@@ -160,18 +170,22 @@
 ## Current Status
 - ✅ Phase 1.1 - MVVM Architecture - COMPLETED
 - ✅ Phase 1.2 - Dependency Injection with Hilt - COMPLETED
+- ✅ Phase 1.3 - Room Database Setup - COMPLETED
+- ✅ Phase 1.4 - Coroutines & Flow Optimization - COMPLETED
 - ✅ Phase 2.1 - Connection Status Management - COMPLETED
 - 🚀 Next Priority Options:
-  - Phase 1.3 - Room Database (Foundation for persistence)
   - Phase 2.2 - User List Feature (Core UX)
   - Phase 2.3 - Enhanced Messaging (Better UX)
   - Phase 3.1 - Material 3 Update (Modern UI)
+  - Phase 3.2 - Dark Theme & Theming
 
 ## Completed Work Summary
 1. **MVVM Architecture**: ViewModels, Repository pattern, proper separation of concerns
 2. **Dependency Injection**: Hilt integration with proper scoping and lifecycle management
-3. **Connection Status**: Visual indicator with real-time updates, activity-based detection
+3. **Room Database**: Message persistence with proper DAOs and entity mapping
-4. **Sleep/Wake Handling**: Auto-recovery when messages resume after device sleep
+4. **Coroutines & Flow**: Converted callbacks to coroutines, implemented Flow for reactive streams, proper scope management
+5. **Connection Status**: Visual indicator with real-time updates, activity-based detection
+6. **Sleep/Wake Handling**: Auto-recovery when messages resume after device sleep
 
 ## Development Notes
 - Architecture foundation (Phase 1) should be completed before moving to advanced features

app/src/main/java/com/mattintech/lchat/network/WifiAwareManager.kt

@@ -9,7 +9,13 @@ import android.os.Build
 import android.util.Log
 import androidx.annotation.RequiresApi
 import com.mattintech.lchat.utils.LOG_PREFIX
+import kotlinx.coroutines.*
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.*
 import java.util.concurrent.ConcurrentHashMap
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+import kotlin.coroutines.suspendCoroutine
 
 @RequiresApi(Build.VERSION_CODES.O)
 class WifiAwareManager(private val context: Context) {
@@ -18,6 +24,12 @@ class WifiAwareManager(private val context: Context) {
         private const val TAG = LOG_PREFIX + "WifiAwareManager:"
         private const val SERVICE_NAME = "lchat"
         private const val PORT = 8888
+        
+        // Keep-alive constants
+        private const val KEEP_ALIVE_INTERVAL_MS = 15000L // Send keep-alive every 15 seconds
+        private const val KEEP_ALIVE_TIMEOUT_MS = 30000L // Consider connection lost after 30 seconds
+        private const val MESSAGE_TYPE_KEEP_ALIVE = "KEEP_ALIVE"
+        private const val MESSAGE_TYPE_KEEP_ALIVE_ACK = "KEEP_ALIVE_ACK"
     }
     
     private var wifiAwareManager: android.net.wifi.aware.WifiAwareManager? = null
@@ -26,31 +38,82 @@ class WifiAwareManager(private val context: Context) {
     private var subscribeDiscoverySession: SubscribeDiscoverySession? = null
     
     private val peerHandles = ConcurrentHashMap<String, PeerHandle>()
+    private val peerRoomMapping = ConcurrentHashMap<String, String>() // PeerHandle.toString() -> roomName
+    private var currentRoom: String? = null
+    
+    // Replace callbacks with Flows
+    private val _messageFlow = MutableSharedFlow<Triple<String, String, String>>()
+    val messageFlow: SharedFlow<Triple<String, String, String>> = _messageFlow.asSharedFlow()
+    
+    private val _connectionFlow = MutableSharedFlow<Pair<String, Boolean>>()
+    val connectionFlow: SharedFlow<Pair<String, Boolean>> = _connectionFlow.asSharedFlow()
+    
+    // Keep legacy callbacks for backward compatibility
     private var messageCallback: ((String, String, String) -> Unit)? = null
     private var connectionCallback: ((String, Boolean) -> Unit)? = null
     
-    private val attachCallback = object : AttachCallback() {
+    // Exception handler for coroutine errors
-        override fun onAttached(session: WifiAwareSession) {
+    private val exceptionHandler = CoroutineExceptionHandler { _, throwable ->
-            Log.d(TAG, "Wi-Fi Aware attached")
+        Log.e(TAG, "Coroutine exception: ", throwable)
-            wifiAwareSession = session
+    }
-        }
+    
-        
+    private val coroutineScope = CoroutineScope(SupervisorJob() + Dispatchers.IO + exceptionHandler)
-        override fun onAttachFailed() {
+    
-            Log.e(TAG, "Wi-Fi Aware attach failed")
+    // Keep-alive tracking
+    private val lastPeerActivity = ConcurrentHashMap<String, Long>()
+    private var keepAliveJob: Job? = null
+    
+    fun initialize() {
+        coroutineScope.launch {
+            val result = initializeAsync()
+            if (result.isFailure) {
+                Log.e(TAG, "Failed to initialize Wi-Fi Aware: ${result.exceptionOrNull()?.message}")
+            }
         }
     }
     
-    fun initialize() {
+    suspend fun initializeAsync(): Result<Unit> = withContext(Dispatchers.IO) {
         wifiAwareManager = context.getSystemService(Context.WIFI_AWARE_SERVICE) as? android.net.wifi.aware.WifiAwareManager
         
-        if (wifiAwareManager?.isAvailable == true) {
+        // Always check if Wi-Fi Aware is available
-            wifiAwareManager?.attach(attachCallback, null)
+        if (wifiAwareManager?.isAvailable != true) {
-        } else {
             Log.e(TAG, "Wi-Fi Aware is not available")
+            wifiAwareSession = null
+            return@withContext Result.failure(Exception("Wi-Fi Aware is not available"))
+        }
+        
+        // If we already have a session, verify it's still valid
+        if (wifiAwareSession != null) {
+            Log.d(TAG, "Wi-Fi Aware already initialized - verifying session is still valid")
+            // Session is likely still valid
+            return@withContext Result.success(Unit)
+        }
+        
+        // Need to attach
+        attachToWifiAware()
+    }
+    
+    private suspend fun attachToWifiAware(): Result<Unit> = suspendCancellableCoroutine { continuation ->
+        wifiAwareManager?.attach(object : AttachCallback() {
+            override fun onAttached(session: WifiAwareSession) {
+                Log.d(TAG, "Wi-Fi Aware attached")
+                wifiAwareSession = session
+                continuation.resume(Result.success(Unit))
+            }
+            
+            override fun onAttachFailed() {
+                Log.e(TAG, "Wi-Fi Aware attach failed")
+                continuation.resume(Result.failure(Exception("Wi-Fi Aware attach failed")))
+            }
+        }, null)
+        
+        continuation.invokeOnCancellation {
+            Log.d(TAG, "Wi-Fi Aware attach cancelled")
         }
     }
     
     fun startHostMode(roomName: String) {
+        currentRoom = roomName
         val config = PublishConfig.Builder()
             .setServiceName(SERVICE_NAME)
             .setServiceSpecificInfo(roomName.toByteArray())
@@ -66,17 +129,56 @@ class WifiAwareManager(private val context: Context) {
                 val messageStr = String(message)
                 Log.d(TAG, "Host: Received message: $messageStr")
                 
-                if (messageStr == "CONNECT_REQUEST") {
+                when (messageStr) {
-                    Log.d(TAG, "Host: Received connection request")
+                    "CONNECT_REQUEST" -> {
-                    acceptConnection(peerHandle)
+                        Log.d(TAG, "Host: Received connection request from peer")
-                } else {
+                        coroutineScope.launch {
-                    handleIncomingMessage(peerHandle, message)
+                            val result = acceptConnectionAsync(peerHandle)
+                            if (result.isSuccess) {
+                                Log.d(TAG, "Host: Successfully accepted connection")
+                                startKeepAlive()
+                            } else {
+                                Log.e(TAG, "Host: Failed to accept connection: ${result.exceptionOrNull()?.message}")
+                            }
+                        }
+                    }
+                    MESSAGE_TYPE_KEEP_ALIVE -> {
+                        Log.d(TAG, "Host: Received keep-alive from peer")
+                        handleKeepAlive(peerHandle, true)
+                    }
+                    MESSAGE_TYPE_KEEP_ALIVE_ACK -> {
+                        Log.d(TAG, "Host: Received keep-alive ACK from peer")
+                        handleKeepAlive(peerHandle, false)
+                    }
+                    else -> {
+                        handleIncomingMessage(peerHandle, message)
+                    }
+                }
+            }
+            
+            override fun onSessionTerminated() {
+                Log.w(TAG, "Host publish session terminated")
+                publishDiscoverySession = null
+                stopKeepAlive()
+                // Emit disconnection event
+                coroutineScope.launch {
+                    _connectionFlow.emit(Pair("", false))
                 }
             }
         }, null)
     }
     
     fun startClientMode() {
+        // Close any existing subscribe session before starting a new one
+        if (subscribeDiscoverySession != null) {
+            Log.d(TAG, "Closing existing subscribe session before starting new one")
+            subscribeDiscoverySession?.close()
+            subscribeDiscoverySession = null
+        }
+        
+        // Clear any stale peer handles
+        peerHandles.clear()
+        
         val config = SubscribeConfig.Builder()
             .setServiceName(SERVICE_NAME)
             .build()
@@ -97,24 +199,66 @@ class WifiAwareManager(private val context: Context) {
                 
                 // Store peer handle for this room
                 peerHandles[roomName] = peerHandle
+                currentRoom = roomName
                 
                 // Send connection request to host
                 Log.d(TAG, "Sending connection request to room: $roomName")
                 subscribeDiscoverySession?.sendMessage(peerHandle, 0, "CONNECT_REQUEST".toByteArray())
                 
+                // Update peer activity when discovered
+                val peerId = peerHandle.toString()
+                peerRoomMapping[peerId] = roomName
+                lastPeerActivity[roomName] = System.currentTimeMillis()
+                
                 // Wait a bit for host to prepare, then connect
-                android.os.Handler(android.os.Looper.getMainLooper()).postDelayed({
+                coroutineScope.launch {
-                    connectToPeer(peerHandle, roomName)
+                    delay(500)
-                }, 500)
+                    val result = connectToPeerAsync(peerHandle, roomName)
+                    if (result.isFailure) {
+                        Log.e(TAG, "Failed to connect to peer: ${result.exceptionOrNull()?.message}")
+                        lastPeerActivity.remove(peerId)
+                    }
+                }
             }
             
             override fun onMessageReceived(peerHandle: PeerHandle, message: ByteArray) {
-                handleIncomingMessage(peerHandle, message)
+                val messageStr = String(message)
+                when (messageStr) {
+                    MESSAGE_TYPE_KEEP_ALIVE -> {
+                        Log.d(TAG, "Client: Received keep-alive from host")
+                        handleKeepAlive(peerHandle, true)
+                    }
+                    MESSAGE_TYPE_KEEP_ALIVE_ACK -> {
+                        Log.d(TAG, "Client: Received keep-alive ACK from host")
+                        handleKeepAlive(peerHandle, false)
+                    }
+                    else -> {
+                        handleIncomingMessage(peerHandle, message)
+                    }
+                }
+            }
+            
+            override fun onSessionTerminated() {
+                Log.w(TAG, "Client subscribe session terminated")
+                subscribeDiscoverySession = null
+                stopKeepAlive()
+                // Don't clear wifiAwareSession - it's still valid
+                // Emit disconnection event
+                coroutineScope.launch {
+                    _connectionFlow.emit(Pair("", false))
+                }
             }
         }, null)
     }
     
     private fun connectToPeer(peerHandle: PeerHandle, roomName: String) {
+        coroutineScope.launch {
+            connectToPeerAsync(peerHandle, roomName)
+        }
+    }
+    
+    private suspend fun connectToPeerAsync(peerHandle: PeerHandle, roomName: String): Result<Unit> = withContext(Dispatchers.IO) {
+        suspendCancellableCoroutine { continuation ->
         Log.d(TAG, "connectToPeer: Starting connection to room: $roomName")
         val networkSpecifier = WifiAwareNetworkSpecifier.Builder(subscribeDiscoverySession!!, peerHandle)
             .setPskPassphrase("lchat-secure-key")
@@ -129,69 +273,155 @@ class WifiAwareManager(private val context: Context) {
             
         val connectivityManager = context.getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager
         
-        try {
+            try {
-            connectivityManager.requestNetwork(networkRequest, object : ConnectivityManager.NetworkCallback() {
+                var isResumed = false
-                override fun onAvailable(network: android.net.Network) {
+                val callback = object : ConnectivityManager.NetworkCallback() {
-                    Log.d(TAG, "onAvailable: Network connected for room: $roomName")
+                    override fun onAvailable(network: android.net.Network) {
-                    connectionCallback?.invoke(roomName, true)
+                        Log.d(TAG, "onAvailable: Network connected for room: $roomName")
+                        if (!isResumed) {
+                            isResumed = true
+                            continuation.resume(Result.success(Unit))
+                        }
+                        // Emit to flow and legacy callback
+                        coroutineScope.launch {
+                            _connectionFlow.emit(Pair(roomName, true))
+                        }
+                        connectionCallback?.invoke(roomName, true)
+                        // Start keep-alive for client connections
+                        startKeepAlive()
+                    }
+                    
+                    override fun onLost(network: android.net.Network) {
+                        Log.d(TAG, "onLost: Network lost for room: $roomName")
+                        // Check if we still have active keep-alive for this room
+                        val lastActivity = lastPeerActivity[roomName]
+                        val currentTime = System.currentTimeMillis()
+                        
+                        if (lastActivity != null && (currentTime - lastActivity) < KEEP_ALIVE_TIMEOUT_MS) {
+                            Log.d(TAG, "Network lost but keep-alive still active for room: $roomName - ignoring disconnect")
+                            // Don't disconnect if keep-alive is still active
+                            return
+                        }
+                        
+                        // Clear peer handles when connection is truly lost
+                        peerHandles.remove(roomName)
+                        lastPeerActivity.remove(roomName)
+                        coroutineScope.launch {
+                            _connectionFlow.emit(Pair(roomName, false))
+                        }
+                        connectionCallback?.invoke(roomName, false)
+                    }
+                    
+                    override fun onUnavailable() {
+                        Log.e(TAG, "onUnavailable: Network request failed for room: $roomName")
+                        if (!isResumed) {
+                            isResumed = true
+                            continuation.resume(Result.failure(Exception("Network unavailable for room: $roomName")))
+                        }
+                        coroutineScope.launch {
+                            _connectionFlow.emit(Pair(roomName, false))
+                        }
+                        connectionCallback?.invoke(roomName, false)
+                    }
+                    
+                    override fun onCapabilitiesChanged(network: android.net.Network, networkCapabilities: NetworkCapabilities) {
+                        Log.d(TAG, "onCapabilitiesChanged: Capabilities changed for room: $roomName")
+                    }
+                    
+                    override fun onLinkPropertiesChanged(network: android.net.Network, linkProperties: android.net.LinkProperties) {
+                        Log.d(TAG, "onLinkPropertiesChanged: Link properties changed for room: $roomName")
+                    }
                 }
                 
-                override fun onLost(network: android.net.Network) {
+                connectivityManager.requestNetwork(networkRequest, callback, android.os.Handler(android.os.Looper.getMainLooper()))
-                    Log.d(TAG, "onLost: Network lost for room: $roomName")
-                    connectionCallback?.invoke(roomName, false)
-                }
                 
-                override fun onUnavailable() {
+                Log.d(TAG, "connectToPeer: Network request submitted for room: $roomName")
-                    Log.e(TAG, "onUnavailable: Network request failed for room: $roomName")
-                    connectionCallback?.invoke(roomName, false)
-                }
                 
-                override fun onCapabilitiesChanged(network: android.net.Network, networkCapabilities: NetworkCapabilities) {
+                continuation.invokeOnCancellation {
-                    Log.d(TAG, "onCapabilitiesChanged: Capabilities changed for room: $roomName")
+                    connectivityManager.unregisterNetworkCallback(callback)
                 }
-                
+            } catch (e: Exception) {
-                override fun onLinkPropertiesChanged(network: android.net.Network, linkProperties: android.net.LinkProperties) {
+                Log.e(TAG, "connectToPeer: Failed to request network", e)
-                    Log.d(TAG, "onLinkPropertiesChanged: Link properties changed for room: $roomName")
+                continuation.resume(Result.failure(e))
+                coroutineScope.launch {
+                    _connectionFlow.emit(Pair(roomName, false))
                 }
-            }, android.os.Handler(android.os.Looper.getMainLooper()), 30000) // 30 second timeout
+                connectionCallback?.invoke(roomName, false)
-            
+            }
-            Log.d(TAG, "connectToPeer: Network request submitted for room: $roomName")
-        } catch (e: Exception) {
-            Log.e(TAG, "connectToPeer: Failed to request network", e)
-            connectionCallback?.invoke(roomName, false)
         }
     }
     
-    private fun acceptConnection(peerHandle: PeerHandle) {
+    
-        Log.d(TAG, "acceptConnection: Accepting connection from client")
+    private suspend fun acceptConnectionAsync(peerHandle: PeerHandle): Result<Unit> = withContext(Dispatchers.IO) {
-        val networkSpecifier = WifiAwareNetworkSpecifier.Builder(publishDiscoverySession!!, peerHandle)
+        suspendCancellableCoroutine { continuation ->
-            .setPskPassphrase("lchat-secure-key")
+            Log.d(TAG, "acceptConnection: Accepting connection from client")
-            .setPort(PORT)
-            .build()
             
-        val networkRequest = NetworkRequest.Builder()
+            try {
-            .addTransportType(NetworkCapabilities.TRANSPORT_WIFI_AWARE)
+                val networkSpecifier = WifiAwareNetworkSpecifier.Builder(publishDiscoverySession!!, peerHandle)
-            .setNetworkSpecifier(networkSpecifier)
+                    .setPskPassphrase("lchat-secure-key")
-            .build()
+                    .setPort(PORT)
-            
+                    .build()
-        val connectivityManager = context.getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager
+                
-        connectivityManager.requestNetwork(networkRequest, object : ConnectivityManager.NetworkCallback() {
+                val networkRequest = NetworkRequest.Builder()
-            override fun onAvailable(network: android.net.Network) {
+                    .addTransportType(NetworkCapabilities.TRANSPORT_WIFI_AWARE)
-                Log.d(TAG, "Client connected")
+                    .setNetworkSpecifier(networkSpecifier)
-                peerHandles[peerHandle.toString()] = peerHandle
+                    .build()
+                
+                val connectivityManager = context.getSystemService(Context.CONNECTIVITY_SERVICE) as ConnectivityManager
+                var isResumed = false
+                
+                val callback = object : ConnectivityManager.NetworkCallback() {
+                    override fun onAvailable(network: android.net.Network) {
+                        Log.d(TAG, "Client connected")
+                        val peerId = peerHandle.toString()
+                        val roomName = currentRoom ?: "host"
+                        peerHandles[roomName] = peerHandle
+                        peerRoomMapping[peerId] = roomName
+                        lastPeerActivity[roomName] = System.currentTimeMillis()
+                        if (!isResumed) {
+                            isResumed = true
+                            continuation.resume(Result.success(Unit))
+                        }
+                    }
+                    
+                    override fun onUnavailable() {
+                        Log.e(TAG, "Failed to accept client connection - Check if Wi-Fi is enabled")
+                        if (!isResumed) {
+                            isResumed = true
+                            continuation.resume(Result.failure(Exception("Failed to accept client connection")))
+                        }
+                    }
+                }
+                
+                connectivityManager.requestNetwork(networkRequest, callback)
+                
+                continuation.invokeOnCancellation {
+                    connectivityManager.unregisterNetworkCallback(callback)
+                }
+            } catch (e: Exception) {
+                Log.e(TAG, "acceptConnection: Failed to accept connection", e)
+                continuation.resume(Result.failure(e))
             }
-            
+        }
-            override fun onUnavailable() {
-                Log.e(TAG, "Failed to accept client connection - Check if Wi-Fi is enabled")
-            }
-        })
     }
     
     private fun handleIncomingMessage(peerHandle: PeerHandle, message: ByteArray) {
         try {
+            // Update peer activity on any message
+            val peerId = peerHandle.toString()
+            val roomName = peerRoomMapping[peerId] ?: currentRoom
+            
+            if (roomName != null) {
+                lastPeerActivity[roomName] = System.currentTimeMillis()
+            }
+            
             val messageStr = String(message)
             val parts = messageStr.split("|", limit = 3)
             if (parts.size == 3) {
+                // Emit to flow and legacy callback
+                coroutineScope.launch {
+                    _messageFlow.emit(Triple(parts[0], parts[1], parts[2]))
+                }
                 messageCallback?.invoke(parts[0], parts[1], parts[2])
             }
         } catch (e: Exception) {
@@ -199,18 +429,35 @@ class WifiAwareManager(private val context: Context) {
         }
     }
     
-    fun sendMessage(userId: String, userName: String, content: String) {
+    fun sendMessage(userId: String, userName: String, content: String): Boolean {
         val message = "$userId|$userName|$content".toByteArray()
+        var messagesSent = 0
         
-        if (publishDiscoverySession != null) {
+        if (publishDiscoverySession != null && peerHandles.isNotEmpty()) {
             peerHandles.values.forEach { peerHandle ->
-                publishDiscoverySession?.sendMessage(peerHandle, 0, message)
+                try {
+                    publishDiscoverySession?.sendMessage(peerHandle, messagesSent, message)
+                    messagesSent++
+                } catch (e: Exception) {
+                    Log.e(TAG, "Failed to send message to peer", e)
+                }
             }
-        } else if (subscribeDiscoverySession != null) {
+        } else if (subscribeDiscoverySession != null && peerHandles.isNotEmpty()) {
             peerHandles.values.forEach { peerHandle ->
-                subscribeDiscoverySession?.sendMessage(peerHandle, 0, message)
+                try {
+                    subscribeDiscoverySession?.sendMessage(peerHandle, messagesSent, message)
+                    messagesSent++
+                } catch (e: Exception) {
+                    Log.e(TAG, "Failed to send message to peer", e)
+                }
             }
         }
+        
+        if (messagesSent == 0) {
+            Log.w(TAG, "No messages sent - no active session or no peer handles")
+        }
+        
+        return messagesSent > 0
     }
     
     fun setMessageCallback(callback: (String, String, String) -> Unit) {
@@ -221,10 +468,116 @@ class WifiAwareManager(private val context: Context) {
         connectionCallback = callback
     }
     
+    private fun startKeepAlive() {
+        keepAliveJob?.cancel()
+        keepAliveJob = coroutineScope.launch {
+            while (isActive) {
+                delay(KEEP_ALIVE_INTERVAL_MS)
+                sendKeepAlive()
+                checkPeerActivity()
+            }
+        }
+    }
+    
+    private fun stopKeepAlive() {
+        keepAliveJob?.cancel()
+        keepAliveJob = null
+        lastPeerActivity.clear()
+    }
+    
+    private fun sendKeepAlive() {
+        val message = MESSAGE_TYPE_KEEP_ALIVE.toByteArray()
+        var messagesSent = 0
+        
+        if (publishDiscoverySession != null && peerHandles.isNotEmpty()) {
+            peerHandles.forEach { (roomName, peerHandle) ->
+                try {
+                    publishDiscoverySession?.sendMessage(peerHandle, messagesSent, message)
+                    messagesSent++
+                    Log.d(TAG, "Sent keep-alive to room: $roomName")
+                } catch (e: Exception) {
+                    Log.e(TAG, "Failed to send keep-alive to room: $roomName", e)
+                }
+            }
+        } else if (subscribeDiscoverySession != null && peerHandles.isNotEmpty()) {
+            peerHandles.forEach { (roomName, peerHandle) ->
+                try {
+                    subscribeDiscoverySession?.sendMessage(peerHandle, messagesSent, message)
+                    messagesSent++
+                    Log.d(TAG, "Sent keep-alive to room: $roomName")
+                } catch (e: Exception) {
+                    Log.e(TAG, "Failed to send keep-alive to room: $roomName", e)
+                }
+            }
+        }
+    }
+    
+    private fun handleKeepAlive(peerHandle: PeerHandle, shouldReply: Boolean) {
+        // Update last activity for this peer
+        val peerId = peerHandle.toString()
+        val roomName = peerRoomMapping[peerId] ?: currentRoom
+        
+        if (roomName != null) {
+            lastPeerActivity[roomName] = System.currentTimeMillis()
+            Log.d(TAG, "Updated keep-alive activity for room: $roomName")
+        }
+        
+        // Send acknowledgment if requested
+        if (shouldReply) {
+            val ackMessage = MESSAGE_TYPE_KEEP_ALIVE_ACK.toByteArray()
+            try {
+                if (publishDiscoverySession != null) {
+                    publishDiscoverySession?.sendMessage(peerHandle, 0, ackMessage)
+                } else if (subscribeDiscoverySession != null) {
+                    subscribeDiscoverySession?.sendMessage(peerHandle, 0, ackMessage)
+                }
+                Log.d(TAG, "Sent keep-alive ACK to peer")
+            } catch (e: Exception) {
+                Log.e(TAG, "Failed to send keep-alive ACK", e)
+            }
+        }
+    }
+    
+    private fun checkPeerActivity() {
+        val currentTime = System.currentTimeMillis()
+        val inactiveRooms = mutableListOf<String>()
+        
+        lastPeerActivity.forEach { (roomName, lastActivity) ->
+            if (currentTime - lastActivity > KEEP_ALIVE_TIMEOUT_MS) {
+                Log.w(TAG, "Room $roomName inactive for too long, considering disconnected")
+                inactiveRooms.add(roomName)
+            }
+        }
+        
+        // Remove inactive rooms
+        inactiveRooms.forEach { roomName ->
+            lastPeerActivity.remove(roomName)
+            peerHandles.remove(roomName)
+            
+            // Clean up peer room mapping
+            peerRoomMapping.entries.removeIf { it.value == roomName }
+            
+            // Emit disconnection event for this room
+            coroutineScope.launch {
+                _connectionFlow.emit(Pair(roomName, false))
+            }
+        }
+    }
+    
     fun stop() {
+        Log.d(TAG, "Stopping WifiAwareManager")
+        stopKeepAlive()
         publishDiscoverySession?.close()
+        publishDiscoverySession = null
         subscribeDiscoverySession?.close()
+        subscribeDiscoverySession = null
+        // Close and clear the wifiAwareSession to force re-attachment
         wifiAwareSession?.close()
+        wifiAwareSession = null
         peerHandles.clear()
+        peerRoomMapping.clear()
+        currentRoom = null
+        // Don't cancel the coroutine scope - we need it for future operations
+        Log.d(TAG, "WifiAwareManager stopped - session cleared for fresh start")
     }
 }

app/src/main/java/com/mattintech/lchat/repository/ChatRepository.kt

@@ -19,6 +19,14 @@ class ChatRepository @Inject constructor(
     private val wifiAwareManager: WifiAwareManager,
     private val messageDao: MessageDao
 ) {
+    // Exception handler for repository operations
+    private val exceptionHandler = CoroutineExceptionHandler { _, throwable ->
+        android.util.Log.e("ChatRepository", "Repository coroutine exception: ", throwable)
+        _connectionState.value = ConnectionState.Error(throwable.message ?: "Unknown error")
+    }
+    
+    // Repository scope for background operations
+    private val repositoryScope = CoroutineScope(SupervisorJob() + Dispatchers.IO + exceptionHandler)
     
     private var currentRoomName: String = ""
     
@@ -41,93 +49,145 @@ class ChatRepository @Inject constructor(
     private var messageCallback: ((String, String, String) -> Unit)? = null
     private var connectionCallback: ((String, Boolean) -> Unit)? = null
     
-    private var lastActivityTime = System.currentTimeMillis()
+    // Keep-alive is now handled by WifiAwareManager
-    private var connectionCheckJob: Job? = null
-    private val connectionTimeout = 30000L // 30 seconds
     
     init {
         wifiAwareManager.initialize()
-        setupWifiAwareCallbacks()
+        // Only use Flow-based collection, not callbacks
+        collectWifiAwareFlows()
     }
     
-    private fun setupWifiAwareCallbacks() {
+    private fun collectWifiAwareFlows() {
-        wifiAwareManager.setMessageCallback { userId, userName, content ->
+        // Collect messages from Flow
-            val message = Message(
+        repositoryScope.launch {
-                id = UUID.randomUUID().toString(),
+            try {
-                userId = userId,
+                wifiAwareManager.messageFlow.collect { (userId, userName, content) ->
-                userName = userName,
+                val message = Message(
-                content = content,
+                    id = UUID.randomUUID().toString(),
-                timestamp = System.currentTimeMillis(),
+                    userId = userId,
-                isOwnMessage = false
+                    userName = userName,
-            )
+                    content = content,
-            addMessage(message)
+                    timestamp = System.currentTimeMillis(),
-            
+                    isOwnMessage = false
-            // Update last activity time
+                )
-            lastActivityTime = System.currentTimeMillis()
+                addMessage(message)
-            
+                
-            // If we're receiving messages, we must be connected
+                // Don't change connection state based on messages - let WifiAwareManager handle it
-            if (_connectionState.value !is ConnectionState.Connected && 
-                _connectionState.value !is ConnectionState.Hosting) {
-                when (_connectionState.value) {
-                    is ConnectionState.Hosting -> {} // Keep hosting state
-                    else -> _connectionState.value = ConnectionState.Connected("Active")
                 }
+            } catch (e: Exception) {
+                android.util.Log.e("ChatRepository", "Error collecting message flow", e)
+            }
+        }
+        
+        // Collect connection state from Flow
+        repositoryScope.launch {
+            try {
+                wifiAwareManager.connectionFlow.collect { (roomName, isConnected) ->
+                    android.util.Log.d("ChatRepository", "Connection flow update - room: $roomName, connected: $isConnected, current state: ${_connectionState.value}")
+                    
+                    if (isConnected) {
+                        currentRoomName = roomName
+                        // Only change to connected if we're not already hosting
+                        if (_connectionState.value !is ConnectionState.Hosting) {
+                            _connectionState.value = ConnectionState.Connected(roomName)
+                        }
+                        loadMessagesFromDatabase(roomName)
+                    } else {
+                        // Only disconnect if the room matches our current room
+                        if (roomName.isEmpty() || roomName == currentRoomName) {
+                            _connectionState.value = ConnectionState.Disconnected
+                        }
+                    }
+                    // Call the legacy callback if set
+                    connectionCallback?.invoke(roomName, isConnected)
+                }
+            } catch (e: Exception) {
+                android.util.Log.e("ChatRepository", "Error collecting connection flow", e)
+                _connectionState.value = ConnectionState.Error(e.message ?: "Connection error")
             }
-            
-            messageCallback?.invoke(userId, userName, content)
         }
     }
     
+    // Removed setupWifiAwareCallbacks - now using Flow-based collection only
+    
     fun startHostMode(roomName: String) {
         currentRoomName = roomName
+        // Ensure WifiAwareManager is initialized before starting
+        wifiAwareManager.initialize()
         wifiAwareManager.startHostMode(roomName)
         _connectionState.value = ConnectionState.Hosting(roomName)
-        startConnectionMonitoring()
         loadMessagesFromDatabase(roomName)
     }
     
     private fun loadMessagesFromDatabase(roomName: String) {
-        CoroutineScope(Dispatchers.IO).launch {
+        repositoryScope.launch {
-            val storedMessages = messageDao.getMessagesForRoomOnce(roomName)
+            try {
-                .map { it.toMessage() }
+                val storedMessages = messageDao.getMessagesForRoomOnce(roomName)
-            _messages.value = storedMessages
+                    .map { it.toMessage() }
+                _messages.value = storedMessages
+            } catch (e: Exception) {
+                android.util.Log.e("ChatRepository", "Error loading messages from database", e)
+                // Don't crash, just continue with empty messages
+                _messages.value = emptyList()
+            }
         }
     }
     
     fun startClientMode() {
-        wifiAwareManager.startClientMode()
+        // Reset state for fresh start
-        _connectionState.value = ConnectionState.Searching
+        _messages.value = emptyList()
-        startConnectionMonitoring()
+        currentRoomName = ""
+        
+        // Ensure WifiAwareManager is initialized before starting
+        repositoryScope.launch {
+            // Give Wi-Fi Aware time to stabilize after network changes
+            delay(500)
+            wifiAwareManager.initialize()
+            // Small delay to ensure initialization completes
+            delay(100)
+            wifiAwareManager.startClientMode()
+            _connectionState.value = ConnectionState.Searching
+        }
     }
     
     fun sendMessage(userId: String, userName: String, content: String) {
-        val message = Message(
+        // Only allow sending messages if connected or hosting
-            id = UUID.randomUUID().toString(),
+        when (_connectionState.value) {
-            userId = userId,
+            is ConnectionState.Connected, is ConnectionState.Hosting -> {
-            userName = userName,
+                val message = Message(
-            content = content,
+                    id = UUID.randomUUID().toString(),
-            timestamp = System.currentTimeMillis(),
+                    userId = userId,
-            isOwnMessage = true
+                    userName = userName,
-        )
+                    content = content,
-        addMessage(message)
+                    timestamp = System.currentTimeMillis(),
-        wifiAwareManager.sendMessage(userId, userName, content)
+                    isOwnMessage = true
-        
+                )
-        // Update last activity time
+                val sent = wifiAwareManager.sendMessage(userId, userName, content)
-        lastActivityTime = System.currentTimeMillis()
+                if (sent) {
-        
+                    addMessage(message)
-        // If we can send messages, update connection state if needed
+                } else {
-        if (_connectionState.value is ConnectionState.Disconnected || 
+                    android.util.Log.e("ChatRepository", "Failed to send message - no active connection")
-            _connectionState.value is ConnectionState.Error) {
+                    _connectionState.value = ConnectionState.Disconnected
-            _connectionState.value = ConnectionState.Connected("Active")
+                }
+            }
+            else -> {
+                android.util.Log.w("ChatRepository", "Cannot send message - not connected. State: ${_connectionState.value}")
+            }
         }
     }
     
     private fun addMessage(message: Message) {
-        _messages.value = _messages.value + message
+        // Add message and sort by timestamp to ensure proper order
+        _messages.value = (_messages.value + message).sortedBy { it.timestamp }
         
         // Save to database
-        CoroutineScope(Dispatchers.IO).launch {
+        repositoryScope.launch {
-            messageDao.insertMessage(message.toEntity(currentRoomName))
+            try {
+                messageDao.insertMessage(message.toEntity(currentRoomName))
+            } catch (e: Exception) {
+                android.util.Log.e("ChatRepository", "Error saving message to database", e)
+                // Don't crash, message is already in memory
+            }
         }
     }
     
@@ -141,49 +201,17 @@ class ChatRepository @Inject constructor(
     
     fun setConnectionCallback(callback: (String, Boolean) -> Unit) {
         connectionCallback = callback
-        wifiAwareManager.setConnectionCallback { roomName, isConnected ->
+        // Connection state is now handled by Flow collection
-            if (isConnected) {
-                currentRoomName = roomName
-                _connectionState.value = ConnectionState.Connected(roomName)
-                loadMessagesFromDatabase(roomName)
-            } else {
-                _connectionState.value = ConnectionState.Disconnected
-            }
-            callback(roomName, isConnected)
-        }
     }
     
     fun stop() {
-        stopConnectionMonitoring()
         wifiAwareManager.stop()
         _connectionState.value = ConnectionState.Disconnected
         _connectedUsers.value = emptyList()
-    }
+        // Don't cancel the repository scope - we need it for future operations
-    
+        // Clear messages when stopping
-    private fun startConnectionMonitoring() {
+        _messages.value = emptyList()
-        connectionCheckJob?.cancel()
+        currentRoomName = ""
-        connectionCheckJob = CoroutineScope(Dispatchers.IO).launch {
-            while (isActive) {
-                delay(5000) // Check every 5 seconds
-                val timeSinceLastActivity = System.currentTimeMillis() - lastActivityTime
-                
-                // If no activity for 30 seconds and we think we're connected, mark as disconnected
-                if (timeSinceLastActivity > connectionTimeout) {
-                    when (_connectionState.value) {
-                        is ConnectionState.Connected,
-                        is ConnectionState.Hosting -> {
-                            _connectionState.value = ConnectionState.Disconnected
-                        }
-                        else -> {} // Keep current state
-                    }
-                }
-            }
-        }
-    }
-    
-    private fun stopConnectionMonitoring() {
-        connectionCheckJob?.cancel()
-        connectionCheckJob = null
     }
     
     sealed class ConnectionState {

app/src/main/java/com/mattintech/lchat/ui/ChatFragment.kt

@@ -138,6 +138,8 @@ class ChatFragment : Fragment() {
     override fun onDestroyView() {
         super.onDestroyView()
         Log.d(TAG, "onDestroyView")
+        // Disconnect when leaving the chat screen
+        viewModel.disconnect()
         _binding = null
     }
 }

app/src/main/java/com/mattintech/lchat/ui/LobbyFragment.kt

@@ -18,6 +18,10 @@ import com.mattintech.lchat.viewmodel.LobbyState
 import com.mattintech.lchat.viewmodel.LobbyViewModel
 import com.mattintech.lchat.utils.LOG_PREFIX
 import dagger.hilt.android.AndroidEntryPoint
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import kotlinx.coroutines.launch
 
 @AndroidEntryPoint
 class LobbyFragment : Fragment() {
@@ -94,48 +98,60 @@ class LobbyFragment : Fragment() {
     }
     
     private fun observeViewModel() {
-        viewModel.savedUserName.observe(viewLifecycleOwner) { savedName ->
+        // Collect saved user name
-            if (!savedName.isNullOrEmpty() && binding.nameInput.text.isNullOrEmpty()) {
+        viewLifecycleOwner.lifecycleScope.launch {
-                binding.nameInput.setText(savedName)
+            viewLifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-            }
+                viewModel.savedUserName.collect { savedName ->
-        }
+                    if (!savedName.isNullOrEmpty() && binding.nameInput.text.isNullOrEmpty()) {
-        
+                        binding.nameInput.setText(savedName)
-        viewModel.state.observe(viewLifecycleOwner) { state ->
-            when (state) {
-                is LobbyState.Idle -> {
-                    binding.noRoomsText.visibility = View.GONE
-                }
-                is LobbyState.Connecting -> {
-                    if (binding.modeRadioGroup.checkedRadioButtonId == R.id.clientRadio) {
-                        binding.noRoomsText.visibility = View.VISIBLE
-                        binding.noRoomsText.text = getString(R.string.connecting)
                     }
                 }
-                is LobbyState.Connected -> {
-                    if (binding.modeRadioGroup.checkedRadioButtonId == R.id.clientRadio) {
-                        val userName = binding.nameInput.text?.toString()?.trim() ?: ""
-                        viewModel.onConnectedToRoom(state.roomName, userName)
-                    }
-                }
-                is LobbyState.Error -> {
-                    binding.noRoomsText.visibility = View.VISIBLE
-                    binding.noRoomsText.text = state.message
-                    Toast.makeText(context, state.message, Toast.LENGTH_LONG).show()
-                }
             }
         }
         
-        viewModel.events.observe(viewLifecycleOwner) { event ->
+        // Collect state
-            when (event) {
+        viewLifecycleOwner.lifecycleScope.launch {
-                is LobbyEvent.NavigateToChat -> {
+            viewLifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                    navigateToChat(event.roomName, event.userName, event.isHost)
+                viewModel.state.collect { state ->
-                    viewModel.clearEvent()
+                    when (state) {
+                        is LobbyState.Idle -> {
+                            binding.noRoomsText.visibility = View.GONE
+                        }
+                        is LobbyState.Connecting -> {
+                            if (binding.modeRadioGroup.checkedRadioButtonId == R.id.clientRadio) {
+                                binding.noRoomsText.visibility = View.VISIBLE
+                                binding.noRoomsText.text = getString(R.string.connecting)
+                            }
+                        }
+                        is LobbyState.Connected -> {
+                            if (binding.modeRadioGroup.checkedRadioButtonId == R.id.clientRadio) {
+                                val userName = binding.nameInput.text?.toString()?.trim() ?: ""
+                                viewModel.onConnectedToRoom(state.roomName, userName)
+                            }
+                        }
+                        is LobbyState.Error -> {
+                            binding.noRoomsText.visibility = View.VISIBLE
+                            binding.noRoomsText.text = state.message
+                            Toast.makeText(context, state.message, Toast.LENGTH_LONG).show()
+                        }
+                    }
                 }
-                is LobbyEvent.ShowError -> {
+            }
-                    Toast.makeText(context, event.message, Toast.LENGTH_SHORT).show()
+        }
-                    viewModel.clearEvent()
+        
+        // Collect events
+        viewLifecycleOwner.lifecycleScope.launch {
+            viewLifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                viewModel.events.collect { event ->
+                    when (event) {
+                        is LobbyEvent.NavigateToChat -> {
+                            navigateToChat(event.roomName, event.userName, event.isHost)
+                        }
+                        is LobbyEvent.ShowError -> {
+                            Toast.makeText(context, event.message, Toast.LENGTH_SHORT).show()
+                        }
+                    }
                 }
-                null -> {}
             }
         }
     }

app/src/main/java/com/mattintech/lchat/viewmodel/ChatViewModel.kt

@@ -1,7 +1,5 @@
 package com.mattintech.lchat.viewmodel
 
-import androidx.lifecycle.LiveData
-import androidx.lifecycle.MutableLiveData
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.mattintech.lchat.data.Message
@@ -25,8 +23,8 @@ class ChatViewModel @Inject constructor(
     private val chatRepository: ChatRepository
 ) : ViewModel() {
     
-    private val _state = MutableLiveData<ChatState>(ChatState.Connected)
+    private val _state = MutableStateFlow<ChatState>(ChatState.Connected)
-    val state: LiveData<ChatState> = _state
+    val state: StateFlow<ChatState> = _state.asStateFlow()
     
     private val _messagesFlow = MutableStateFlow<Flow<List<Message>>>(flowOf(emptyList()))
     
@@ -85,10 +83,8 @@ class ChatViewModel @Inject constructor(
     }
     
     fun disconnect() {
-        viewModelScope.launch {
+        chatRepository.stop()
-            chatRepository.stop()
+        _state.value = ChatState.Disconnected
-            _state.value = ChatState.Disconnected
-        }
     }
     
     override fun onCleared() {

app/src/main/java/com/mattintech/lchat/viewmodel/LobbyViewModel.kt

@@ -1,12 +1,16 @@
 package com.mattintech.lchat.viewmodel
 
-import androidx.lifecycle.LiveData
-import androidx.lifecycle.MutableLiveData
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.mattintech.lchat.repository.ChatRepository
 import com.mattintech.lchat.utils.PreferencesManager
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharedFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -32,14 +36,14 @@ class LobbyViewModel @Inject constructor(
     private val preferencesManager: PreferencesManager
 ) : ViewModel() {
     
-    private val _state = MutableLiveData<LobbyState>(LobbyState.Idle)
+    private val _state = MutableStateFlow<LobbyState>(LobbyState.Idle)
-    val state: LiveData<LobbyState> = _state
+    val state: StateFlow<LobbyState> = _state.asStateFlow()
     
-    private val _events = MutableLiveData<LobbyEvent?>()
+    private val _events = MutableSharedFlow<LobbyEvent>()
-    val events: LiveData<LobbyEvent?> = _events
+    val events: SharedFlow<LobbyEvent> = _events.asSharedFlow()
     
-    private val _savedUserName = MutableLiveData<String?>()
+    private val _savedUserName = MutableStateFlow<String?>(null)
-    val savedUserName: LiveData<String?> = _savedUserName
+    val savedUserName: StateFlow<String?> = _savedUserName.asStateFlow()
     
     init {
         setupConnectionCallback()
@@ -64,12 +68,16 @@ class LobbyViewModel @Inject constructor(
     
     fun startHostMode(roomName: String, userName: String) {
         if (roomName.isBlank()) {
-            _events.value = LobbyEvent.ShowError("Please enter a room name")
+            viewModelScope.launch {
+                _events.emit(LobbyEvent.ShowError("Please enter a room name"))
+            }
             return
         }
         
         if (userName.isBlank()) {
-            _events.value = LobbyEvent.ShowError("Please enter your name")
+            viewModelScope.launch {
+                _events.emit(LobbyEvent.ShowError("Please enter your name"))
+            }
             return
         }
         
@@ -77,13 +85,15 @@ class LobbyViewModel @Inject constructor(
             _state.value = LobbyState.Connecting
             preferencesManager.saveUserName(userName)
             chatRepository.startHostMode(roomName)
-            _events.value = LobbyEvent.NavigateToChat(roomName, userName, true)
+            _events.emit(LobbyEvent.NavigateToChat(roomName, userName, true))
         }
     }
     
     fun startClientMode(userName: String) {
         if (userName.isBlank()) {
-            _events.value = LobbyEvent.ShowError("Please enter your name")
+            viewModelScope.launch {
+                _events.emit(LobbyEvent.ShowError("Please enter your name"))
+            }
             return
         }
         
@@ -96,11 +106,13 @@ class LobbyViewModel @Inject constructor(
     
     fun onConnectedToRoom(roomName: String, userName: String) {
         preferencesManager.saveUserName(userName)
-        _events.value = LobbyEvent.NavigateToChat(roomName, userName, false)
+        viewModelScope.launch {
+            _events.emit(LobbyEvent.NavigateToChat(roomName, userName, false))
+        }
     }
     
     fun clearEvent() {
-        _events.value = null
+        // SharedFlow doesn't need clearing, events are consumed once
     }
     
     fun saveUserName(name: String) {