Merge branch 'improve-device-list-error-handling'

15 files changed, 295 insertions, 117 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/Dialog.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/Dialog.kt
index fdadbd5b29..edfcebac67 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/Dialog.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/Dialog.kt
@@ -86,7 +86,7 @@ fun ShowDeviceRemovalDialog(viewModel: DeviceListViewModel, device: Device) {
                         contentColor = Color.White
                     ),
                     onClick = {
-                        viewModel.confirmRemoval()
+                        viewModel.confirmRemovalOfStagedDevice()
                     }
                 ) {
                     Text(
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/List.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/List.kt
index 9e17d5cf63..3babb29a21 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/List.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/component/List.kt
@@ -1,48 +1,32 @@
 package net.mullvad.mullvadvpn.compose.component
 
+import androidx.annotation.DrawableRes
 import androidx.compose.foundation.Image
-import androidx.compose.foundation.ScrollState
 import androidx.compose.foundation.background
 import androidx.compose.foundation.clickable
 import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
-import androidx.compose.foundation.verticalScroll
+import androidx.compose.foundation.layout.width
+import androidx.compose.material.CircularProgressIndicator
 import androidx.compose.material.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
-import androidx.compose.ui.graphics.painter.Painter
 import androidx.compose.ui.res.colorResource
+import androidx.compose.ui.res.painterResource
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.unit.sp
 import net.mullvad.mullvadvpn.R
-import net.mullvad.mullvadvpn.model.Device
 
 @Composable
-fun DeviceList(
-    devices: List<Device>,
-    onItemClicked: (Device) -> Unit
-) {
-    Column(
-        modifier = Modifier.verticalScroll(ScrollState(0))
-    ) {
-        devices.forEach { device ->
-            DeviceRow(device.name) {
-                onItemClicked(device)
-            }
-        }
-    }
-}
-
-@Composable
-fun DeviceRow(
-    name: String,
-    painter: Painter? = null,
-    onItemClicked: () -> Unit
+fun ListItem(
+    text: String,
+    isLoading: Boolean,
+    @DrawableRes iconResourceId: Int? = null,
+    onClick: () -> Unit
 ) {
     val itemColor = colorResource(id = R.color.blue)
 
@@ -51,13 +35,10 @@ fun DeviceRow(
             .fillMaxWidth()
             .padding(vertical = 1.dp)
             .height(50.dp)
-            .background(itemColor)
-            .clickable {
-                onItemClicked()
-            },
+            .background(itemColor),
     ) {
         Text(
-            text = name,
+            text = text,
             fontSize = 18.sp,
             color = Color.White,
             modifier = Modifier
@@ -67,32 +48,27 @@ fun DeviceRow(
                 .align(Alignment.CenterStart)
         )
 
-        if (painter != null) {
-            Image(
-                painter = painter,
-                contentDescription = "Remove",
-                modifier = Modifier
-                    .align(Alignment.CenterEnd)
-                    .padding(horizontal = 12.dp)
-            )
-        }
-    }
-}
-
-@Composable
-fun <T> ItemList(
-    items: List<T>,
-    itemText: (T) -> String,
-    onItemClicked: (T) -> Unit,
-    itemPainter: Painter? = null,
-    modifier: Modifier = Modifier,
-) {
-    Column(
-        modifier = modifier
-    ) {
-        items.forEach { item ->
-            DeviceRow(itemText.invoke(item), itemPainter) {
-                onItemClicked(item)
+        Box(
+            modifier = Modifier
+                .align(Alignment.CenterEnd)
+                .padding(horizontal = 12.dp)
+        ) {
+            if (isLoading) {
+                CircularProgressIndicator(
+                    strokeWidth = 3.dp,
+                    color = Color.White,
+                    modifier = Modifier
+                        .height(24.dp)
+                        .width(24.dp)
+                )
+            } else if (iconResourceId != null) {
+                Image(
+                    painter = painterResource(id = iconResourceId),
+                    contentDescription = "Remove",
+                    modifier = Modifier
+                        .align(Alignment.CenterEnd)
+                        .clickable { onClick() }
+                )
             }
         }
     }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceListScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceListScreen.kt
index 7794a25603..99218a8837 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceListScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/DeviceListScreen.kt
@@ -29,7 +29,7 @@ import androidx.constraintlayout.compose.ConstraintLayout
 import androidx.constraintlayout.compose.Dimension
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.component.ActionButton
-import net.mullvad.mullvadvpn.compose.component.ItemList
+import net.mullvad.mullvadvpn.compose.component.ListItem
 import net.mullvad.mullvadvpn.compose.component.ShowDeviceRemovalDialog
 import net.mullvad.mullvadvpn.util.capitalizeFirstCharOfEachWord
 import net.mullvad.mullvadvpn.viewmodel.DeviceListViewModel
@@ -42,10 +42,10 @@ fun DeviceListScreen(
 ) {
     val state = viewModel.uiState.collectAsState().value
 
-    if (state.deviceStagedForRemoval != null) {
+    if (state.stagedDevice != null) {
         ShowDeviceRemovalDialog(
             viewModel = viewModel,
-            device = state.deviceStagedForRemoval
+            device = state.stagedDevice
         )
     }
 
@@ -140,14 +140,17 @@ fun DeviceListScreen(
                             width = Dimension.matchParent
                         }
                 ) {
-                    ItemList(
-                        state.devices,
-                        itemText = { it.name.capitalizeFirstCharOfEachWord() },
-                        onItemClicked = {
-                            viewModel.stageDeviceForRemoval(it)
-                        },
-                        itemPainter = painterResource(id = R.drawable.icon_close)
-                    )
+                    Column {
+                        state.deviceUiItems.forEach { deviceUiState ->
+                            ListItem(
+                                text = deviceUiState.device.name.capitalizeFirstCharOfEachWord(),
+                                isLoading = deviceUiState.isLoading,
+                                iconResourceId = R.drawable.icon_close
+                            ) {
+                                viewModel.stageDeviceForRemoval(deviceUiState.device.id)
+                            }
+                        }
+                    }
                 }
             }
         }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/DeviceListUiState.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/DeviceListUiState.kt
index 9e048c2926..e989960301 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/DeviceListUiState.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/state/DeviceListUiState.kt
@@ -3,17 +3,22 @@ package net.mullvad.mullvadvpn.compose.state
 import net.mullvad.mullvadvpn.model.Device
 
 data class DeviceListUiState(
-    val devices: List<Device>,
+    val deviceUiItems: List<DeviceListItemUiState>,
     val isLoading: Boolean,
-    val deviceStagedForRemoval: Device?
+    val stagedDevice: Device?
 ) {
-    val hasTooManyDevices = devices.count() >= 5
+    val hasTooManyDevices = deviceUiItems.count() >= 5
 
     companion object {
         val INITIAL = DeviceListUiState(
-            devices = listOf(),
+            deviceUiItems = emptyList(),
             isLoading = true,
-            deviceStagedForRemoval = null
+            stagedDevice = null
         )
     }
 }
+
+data class DeviceListItemUiState(
+    val device: Device,
+    val isLoading: Boolean
+)
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
index 38fa9117f7..e71dcb8416 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/di/UiModule.kt
@@ -41,6 +41,7 @@ val uiModule = module {
     }
 
     single { ServiceConnectionManager(androidContext()) }
+    single { androidContext().resources }
 
     single { AccountExpiryNotification(get()) }
     single { TunnelStateNotification(get()) }
@@ -50,7 +51,7 @@ val uiModule = module {
     single { DeviceRepository(get()) }
     viewModel { LoginViewModel(get(), get()) }
     viewModel { DeviceRevokedViewModel(get(), get()) }
-    viewModel { DeviceListViewModel(get()) }
+    viewModel { DeviceListViewModel(get(), get()) }
 }
 
 const val APPS_SCOPE = "APPS_SCOPE"
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt
index c351cc8130..ea7ecdcd7b 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ipc/Event.kt
@@ -12,6 +12,7 @@ import net.mullvad.mullvadvpn.model.DeviceState
 import net.mullvad.mullvadvpn.model.GeoIpLocation
 import net.mullvad.mullvadvpn.model.LoginResult
 import net.mullvad.mullvadvpn.model.RelayList
+import net.mullvad.mullvadvpn.model.RemoveDeviceResult
 import net.mullvad.mullvadvpn.model.Settings
 import net.mullvad.mullvadvpn.model.TunnelState
 import net.mullvad.mullvadvpn.model.VoucherSubmissionResult as VoucherSubmissionResultData
@@ -45,6 +46,9 @@ sealed class Event : Message.EventMessage() {
     data class DeviceListUpdate(val event: DeviceListEvent) : Event()
 
     @Parcelize
+    data class DeviceRemovalEvent(val deviceId: String, val result: RemoveDeviceResult) : Event()
+
+    @Parcelize
     data class ListenerReady(val connection: Messenger, val listenerId: Int) : Event()
 
     @Parcelize
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/DeviceList.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/DeviceList.kt
new file mode 100644
index 0000000000..de1acb0e23
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/model/DeviceList.kt
@@ -0,0 +1,7 @@
+package net.mullvad.mullvadvpn.model
+
+sealed class DeviceList {
+    object Unavailable : DeviceList()
+    data class Available(val devices: List<Device>) : DeviceList()
+    object Error : DeviceList()
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt
index cb290a6d27..9bdf30de21 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/DaemonDeviceDataSource.kt
@@ -48,7 +48,9 @@ class DaemonDeviceDataSource(
 
         endpoint.dispatcher.registerHandler(Request.RemoveDevice::class) { request ->
             tracker.newBackgroundJob("removeDeviceJob") {
-                daemon.removeDevice(request.accountToken, request.deviceId)
+                daemon.removeDevice(request.accountToken, request.deviceId).also { result ->
+                    endpoint.sendEvent(Event.DeviceRemovalEvent(request.deviceId, result))
+                }
             }
         }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragments/DeviceListFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragments/DeviceListFragment.kt
index 6d5e47f8dd..7c7e3c4658 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragments/DeviceListFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragments/DeviceListFragment.kt
@@ -4,9 +4,16 @@ import android.os.Bundle
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.widget.Toast
 import androidx.compose.ui.platform.ComposeView
 import androidx.compose.ui.res.colorResource
 import androidx.fragment.app.Fragment
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.flowWithLifecycle
+import androidx.lifecycle.lifecycleScope
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.component.ScaffoldWithTopBar
 import net.mullvad.mullvadvpn.compose.screen.DeviceListScreen
@@ -19,6 +26,11 @@ class DeviceListFragment : Fragment() {
 
     private val deviceListViewModel by viewModel<DeviceListViewModel>()
 
+    override fun onCreate(savedInstanceState: Bundle?) {
+        super.onCreate(savedInstanceState)
+        lifecycleScope.launchUiSubscriptionsOnResume()
+    }
+
     override fun onCreateView(
         inflater: LayoutInflater,
         container: ViewGroup?,
@@ -37,8 +49,8 @@ class DeviceListFragment : Fragment() {
                     content = {
                         DeviceListScreen(
                             viewModel = deviceListViewModel,
-                            onBackClick = this@DeviceListFragment::goBack,
-                            onContinueWithLogin = this@DeviceListFragment::openLoginView
+                            onBackClick = { openLoginView(doTriggerAutoLogin = false) },
+                            onContinueWithLogin = { openLoginView(doTriggerAutoLogin = true) }
                         )
                     }
                 )
@@ -46,10 +58,23 @@ class DeviceListFragment : Fragment() {
         }
     }
 
-    private fun openLoginView() {
+    override fun onResume() {
+        super.onResume()
+        deviceListViewModel.clearStagedDevice()
+    }
+
+    private fun CoroutineScope.launchUiSubscriptionsOnResume() = launch {
+        deviceListViewModel.toastMessages
+            .flowWithLifecycle(lifecycle, Lifecycle.State.RESUMED)
+            .collect {
+                Toast.makeText(context, it, Toast.LENGTH_SHORT).show()
+            }
+    }
+
+    private fun openLoginView(doTriggerAutoLogin: Boolean) {
         parentActivity()?.clearBackStack()
         val loginFragment = LoginFragment().apply {
-            if (deviceListViewModel.accountToken != null) {
+            if (doTriggerAutoLogin && deviceListViewModel.accountToken != null) {
                 arguments = Bundle().apply {
                     putString(
                         ACCOUNT_TOKEN_ARGUMENT_KEY,
@@ -64,10 +89,6 @@ class DeviceListFragment : Fragment() {
         }
     }
 
-    private fun goBack() {
-        parentActivity()?.onBackPressed()
-    }
-
     private fun parentActivity(): MainActivity? {
         return (context as? MainActivity)
     }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
index 790c2404f2..044888e9af 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/DeviceRepository.kt
@@ -4,26 +4,27 @@ import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.withTimeoutOrNull
-import net.mullvad.mullvadvpn.model.Device
+import net.mullvad.mullvadvpn.ipc.Event
+import net.mullvad.mullvadvpn.model.DeviceList
 import net.mullvad.mullvadvpn.model.DeviceListEvent
 import net.mullvad.mullvadvpn.model.DeviceState
 
 class DeviceRepository(
     private val serviceConnectionManager: ServiceConnectionManager,
-    private val deviceListTimeoutMillis: Long = 5000L,
     dispatcher: CoroutineDispatcher = Dispatchers.IO
 ) {
-    private val cachedDeviceList = MutableStateFlow<List<Device>>(emptyList())
+    private val cachedDeviceList = MutableStateFlow<DeviceList>(DeviceList.Unavailable)
 
     val deviceState = serviceConnectionManager.connectionState
         .flatMapLatest { state ->
@@ -52,20 +53,42 @@ class DeviceRepository(
         }
 
     val deviceList = deviceListEvents
-        .map { (it as? DeviceListEvent.Available)?.devices ?: emptyList() }
+        .map {
+            if (it is DeviceListEvent.Available) {
+                DeviceList.Available(it.devices)
+            } else {
+                DeviceList.Error
+            }
+        }
         .onStart {
-            if (cachedDeviceList.value.isNotEmpty()) {
+            if (cachedDeviceList.value is DeviceList.Available) {
                 emit(cachedDeviceList.value)
             }
         }
-        .stateIn(CoroutineScope(Dispatchers.IO), SharingStarted.WhileSubscribed(), emptyList())
+        .shareIn(
+            CoroutineScope(Dispatchers.IO),
+            SharingStarted.WhileSubscribed()
+        )
+
+    val deviceRemovalEvent: SharedFlow<Event.DeviceRemovalEvent> =
+        serviceConnectionManager.connectionState
+            .flatMapLatest { state ->
+                if (state is ServiceConnectionState.ConnectedReady) {
+                    state.container.deviceDataSource.deviceRemovalResult
+                } else {
+                    emptyFlow()
+                }
+            }
+            .shareIn(
+                CoroutineScope(dispatcher),
+                SharingStarted.WhileSubscribed()
+            )
 
     fun refreshDeviceState() {
         serviceConnectionManager.deviceDataSource()?.refreshDevice()
     }
 
     fun removeDevice(accountToken: String, deviceId: String) {
-        cachedDeviceList.value = emptyList()
         serviceConnectionManager.deviceDataSource()?.removeDevice(accountToken, deviceId)
     }
 
@@ -73,17 +96,43 @@ class DeviceRepository(
         serviceConnectionManager.deviceDataSource()?.refreshDeviceList(accountToken)
     }
 
-    suspend fun getDeviceList(accountToken: String): DeviceListEvent {
-        return withTimeoutOrNull(deviceListTimeoutMillis) {
+    fun clearCache() {
+        cachedDeviceList.value = DeviceList.Unavailable
+    }
+
+    private fun updateCache(event: DeviceListEvent, accountToken: String) {
+        cachedDeviceList.value =
+            if (event is DeviceListEvent.Available && event.accountToken == accountToken) {
+                DeviceList.Available(event.devices)
+            } else if (event is DeviceListEvent.Error) {
+                DeviceList.Error
+            } else {
+                DeviceList.Unavailable
+            }
+    }
+
+    suspend fun refreshAndAwaitDeviceListWithTimeout(
+        accountToken: String,
+        shouldClearCache: Boolean,
+        shouldOverrideCache: Boolean,
+        timeoutMillis: Long,
+    ): DeviceListEvent {
+        if (shouldClearCache) {
+            clearCache()
+        }
+
+        val result = withTimeoutOrNull(timeoutMillis) {
             deviceListEvents
                 .onStart {
                     refreshDeviceList(accountToken)
                 }
-                .onEach {
-                    cachedDeviceList.value =
-                        (it as? DeviceListEvent.Available)?.devices ?: emptyList()
-                }
                 .firstOrNull() ?: DeviceListEvent.Error
         } ?: DeviceListEvent.Error
+
+        if (shouldOverrideCache) {
+            updateCache(result, accountToken)
+        }
+
+        return result
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt
index 6f018a27b1..6e445a6d34 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/serviceconnection/ServiceConnectionDeviceDataSource.kt
@@ -31,6 +31,16 @@ class ServiceConnectionDeviceDataSource(
         }
     }
 
+    val deviceRemovalResult = callbackFlow {
+        val handler: (Event.DeviceRemovalEvent) -> Unit = { event ->
+            trySend(event)
+        }
+        dispatcher.registerHandler(Event.DeviceRemovalEvent::class, handler)
+        awaitClose {
+            // The current dispatcher doesn't support unregistration of handlers.
+        }
+    }
+
     // Async result: Event.DeviceChanged
     fun refreshDevice() {
         connection.send(Request.RefreshDeviceState.message)
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
index 027520d293..4a60508a38 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/DeviceListViewModel.kt
@@ -1,44 +1,113 @@
 package net.mullvad.mullvadvpn.viewmodel
 
+import android.content.res.Resources
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.onSubscription
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import kotlinx.coroutines.withTimeoutOrNull
+import net.mullvad.mullvadvpn.R
+import net.mullvad.mullvadvpn.compose.state.DeviceListItemUiState
 import net.mullvad.mullvadvpn.compose.state.DeviceListUiState
 import net.mullvad.mullvadvpn.model.Device
+import net.mullvad.mullvadvpn.model.DeviceList
+import net.mullvad.mullvadvpn.model.RemoveDeviceResult
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
-import net.mullvad.mullvadvpn.util.safeLet
+
+typealias DeviceId = String
 
 class DeviceListViewModel(
-    private val deviceRepository: DeviceRepository
+    private val deviceRepository: DeviceRepository,
+    private val resources: Resources,
+    private val dispatcher: CoroutineDispatcher = Dispatchers.Default
 ) : ViewModel() {
-    private val _stagedForRemoval = MutableStateFlow<Device?>(null)
+    private val _stagedDeviceId = MutableStateFlow<DeviceId?>(null)
+    private val _loadingDevices = MutableStateFlow<List<DeviceId>>(emptyList())
+
+    private val _toastMessages = MutableSharedFlow<String>(extraBufferCapacity = 1)
+    val toastMessages = _toastMessages.asSharedFlow()
+
     var accountToken: String? = null
+    private var cachedDeviceList: List<Device>? = null
 
-    val uiState = deviceRepository.deviceList
-        .combine(_stagedForRemoval) { deviceList, deviceStagedForRemoval ->
-            DeviceListUiState(
-                devices = deviceList,
-                isLoading = false,
-                deviceStagedForRemoval = deviceStagedForRemoval
-            )
+    val uiState = combine(
+        deviceRepository.deviceList,
+        _stagedDeviceId,
+        _loadingDevices
+    ) { deviceList, stagedDeviceId, loadingDevices ->
+        val devices = if (deviceList is DeviceList.Available) {
+            deviceList.devices.also { cachedDeviceList = it }
+        } else {
+            cachedDeviceList
+        }
+        val deviceUiItems = devices?.map { device ->
+            DeviceListItemUiState(device, loadingDevices.any { device.id == it })
+        } ?: emptyList()
+        val isLoading = devices == null
+        val stagedDevice = devices?.firstOrNull { device ->
+            device.id == stagedDeviceId
         }
-        .stateIn(viewModelScope, SharingStarted.WhileSubscribed(), DeviceListUiState.INITIAL)
+        DeviceListUiState(
+            deviceUiItems = deviceUiItems,
+            isLoading = isLoading,
+            stagedDevice = stagedDevice
+        )
+    }.stateIn(viewModelScope, SharingStarted.WhileSubscribed(), DeviceListUiState.INITIAL)
 
-    fun stageDeviceForRemoval(device: Device) {
-        _stagedForRemoval.value = device
+    fun stageDeviceForRemoval(deviceId: DeviceId) {
+        _stagedDeviceId.value = deviceId
     }
 
     fun clearStagedDevice() {
-        _stagedForRemoval.value = null
+        _stagedDeviceId.value = null
     }
 
-    fun confirmRemoval() {
-        safeLet(accountToken, _stagedForRemoval.value) { token, device ->
-            deviceRepository.removeDevice(token, device.id)
-            _stagedForRemoval.value = null
+    fun confirmRemovalOfStagedDevice() {
+        val token = accountToken
+        val stagedDeviceId = _stagedDeviceId.value
+
+        if (token != null && stagedDeviceId != null) {
+            viewModelScope.launch {
+                withContext(dispatcher) {
+                    val result = withTimeoutOrNull(DEVICE_REMOVAL_TIMEOUT_MILLIS) {
+                        deviceRepository.deviceRemovalEvent
+                            .onSubscription {
+                                clearStagedDevice()
+                                setLoadingDevice(stagedDeviceId)
+                                deviceRepository.removeDevice(token, stagedDeviceId)
+                            }
+                            .filter { (deviceId, result) ->
+                                deviceId == stagedDeviceId && result == RemoveDeviceResult.Ok
+                            }
+                            .first()
+                    }
+
+                    clearLoadingDevice(stagedDeviceId)
+
+                    if (result == null) {
+                        _toastMessages.tryEmit(
+                            resources.getString(R.string.failed_to_remove_device)
+                        )
+                        refreshDeviceList()
+                    }
+                }
+            }
+        } else {
+            _toastMessages.tryEmit(resources.getString(R.string.error_occurred))
+            clearLoadingDevices()
+            clearStagedDevice()
+            refreshDeviceList()
         }
     }
 
@@ -47,4 +116,20 @@ class DeviceListViewModel(
     fun refreshDeviceList() = accountToken?.let { token ->
         deviceRepository.refreshDeviceList(token)
     }
+
+    private fun setLoadingDevice(deviceId: DeviceId) {
+        _loadingDevices.value = _loadingDevices.value.toMutableList().apply { add(deviceId) }
+    }
+
+    private fun clearLoadingDevice(deviceId: DeviceId) {
+        _loadingDevices.value = _loadingDevices.value.toMutableList().apply { remove(deviceId) }
+    }
+
+    private fun clearLoadingDevices() {
+        _loadingDevices.value = emptyList()
+    }
+
+    companion object {
+        private const val DEVICE_REMOVAL_TIMEOUT_MILLIS = 5000L
+    }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
index 8751089cf8..e73d42438a 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModel.kt
@@ -80,7 +80,14 @@ class LoginViewModel(
             LoginResult.Ok -> LoginUiState.Success(false)
             LoginResult.InvalidAccount -> LoginUiState.InvalidAccountError
             LoginResult.MaxDevicesReached -> {
-                if (deviceRepository.getDeviceList(accountToken).isAvailable()) {
+                val refreshResult = deviceRepository.refreshAndAwaitDeviceListWithTimeout(
+                    accountToken = accountToken,
+                    shouldClearCache = true,
+                    shouldOverrideCache = true,
+                    timeoutMillis = 5000L
+                )
+
+                if (refreshResult.isAvailable()) {
                     LoginUiState.TooManyDevicesError(accountToken)
                 } else {
                     LoginUiState.TooManyDevicesMissingListError
diff --git a/android/app/src/main/res/values/strings.xml b/android/app/src/main/res/values/strings.xml
index 70cccd341e..4a62b94310 100644
--- a/android/app/src/main/res/values/strings.xml
+++ b/android/app/src/main/res/values/strings.xml
@@ -199,4 +199,5 @@
     <string name="copy_account_number">Copy account number</string>
     <string name="hide_account_number">Hide account number</string>
     <string name="show_account_number">Show account number</string>
+    <string name="failed_to_remove_device">Failed to remove device</string>
 </resources>
diff --git a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
index 4af15c0cee..c6ae770969 100644
--- a/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
+++ b/android/app/src/test/kotlin/net/mullvad/mullvadvpn/viewmodel/LoginViewModelTest.kt
@@ -108,7 +108,14 @@ class LoginViewModelTest {
 
     @Test
     fun testLoginWithTooManyDevicesError() = runBlockingTest {
-        coEvery { mockedDeviceRepository.getDeviceList(any()) } returns DeviceListEvent.Available(
+        coEvery {
+            mockedDeviceRepository.refreshAndAwaitDeviceListWithTimeout(
+                any(),
+                any(),
+                any(),
+                any()
+            )
+        } returns DeviceListEvent.Available(
             DUMMY_ACCOUNT_TOKEN, listOf()
         )