diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeScreen.kt
index 122b8cb69b3..fd53f573e78 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeScreen.kt
@@ -25,13 +25,13 @@ import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.fillMaxHeight
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Alignment
-import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.focus.FocusRequester
 import androidx.compose.ui.focus.focusRequester
@@ -58,7 +58,6 @@ import com.wire.android.ui.common.WireDialogButtonProperties
 import com.wire.android.ui.common.WireDialogButtonType
 import com.wire.android.ui.common.progress.WireCircularProgressIndicator
 import com.wire.android.ui.common.scaffold.WireScaffold
-import com.wire.android.ui.common.textfield.CodeFieldValue
 import com.wire.android.ui.common.textfield.CodeTextField
 import com.wire.android.ui.common.textfield.WireTextFieldState
 import com.wire.android.ui.common.topappbar.WireCenterAlignedTopAppBar
@@ -90,29 +89,44 @@ fun CreateAccountCodeScreen(
 
         CodeContent(
             state = codeState,
-            onCodeChange = { onCodeChange(it, ::navigateToSummaryScreen) },
+            textState = codeTextState,
             onResendCodePressed = ::resendCode,
             onBackPressed = navigator::navigateBack,
-            onErrorDismiss = ::clearCodeError,
-            onRemoveDeviceOpen = {
+            serverConfig = serverConfig
+        )
+
+        (codeState.result as? CreateAccountCodeViewState.Result.Error.DialogError)?.let {
+            val (title, message) = it.getResources(type = codeState.type)
+            WireDialog(
+                title = title,
+                text = message,
+                onDismiss = ::clearCodeError,
+                optionButton1Properties = WireDialogButtonProperties(
+                    onClick = ::clearCodeError,
+                    text = stringResource(id = R.string.label_ok),
+                    type = WireDialogButtonType.Primary,
+                )
+            )
+        }
+        LaunchedEffect(codeState.result) {
+            if (codeState.result is CreateAccountCodeViewState.Result.Success) {
+                navigateToSummaryScreen()
+            }
+            if (codeState.result is CreateAccountCodeViewState.Result.Error.TooManyDevicesError) {
                 clearCodeError()
                 clearCodeField()
                 navigator.navigate(NavigationCommand(RemoveDeviceScreenDestination, BackStackMode.CLEAR_WHOLE))
-            },
-            serverConfig = serverConfig
-        )
+            }
+        }
     }
 }
 
-@OptIn(ExperimentalComposeUiApi::class)
 @Composable
 private fun CodeContent(
     state: CreateAccountCodeViewState,
-    onCodeChange: (CodeFieldValue) -> Unit,
+    textState: TextFieldState,
     onResendCodePressed: () -> Unit,
     onBackPressed: () -> Unit,
-    onErrorDismiss: () -> Unit,
-    onRemoveDeviceOpen: () -> Unit,
     serverConfig: ServerConfig.Links
 ) {
     val focusRequester = remember { FocusRequester() }
@@ -152,10 +166,10 @@ private fun CodeContent(
             Spacer(modifier = Modifier.weight(1f))
             Column(horizontalAlignment = Alignment.CenterHorizontally) {
                 CodeTextField(
-                    value = state.code.text,
-                    onValueChange = onCodeChange,
-                    state = when (state.error) {
-                        is CreateAccountCodeViewState.CodeError.TextFieldError.InvalidActivationCodeError ->
+                    codeLength = state.codeLength,
+                    textState = textState,
+                    state = when (state.result) {
+                        is CreateAccountCodeViewState.Result.Error.TextFieldError.InvalidActivationCodeError ->
                             WireTextFieldState.Error(stringResource(id = R.string.create_account_code_error))
 
                         else -> WireTextFieldState.Default
@@ -180,51 +194,36 @@ private fun CodeContent(
             keyboardController?.show()
         }
     }
-    if (state.error is CreateAccountCodeViewState.CodeError.DialogError) {
-        val (title, message) = state.error.getResources(type = state.type)
-        WireDialog(
-            title = title,
-            text = message,
-            onDismiss = onErrorDismiss,
-            optionButton1Properties = WireDialogButtonProperties(
-                onClick = onErrorDismiss,
-                text = stringResource(id = R.string.label_ok),
-                type = WireDialogButtonType.Primary,
-            )
-        )
-    } else if (state.error is CreateAccountCodeViewState.CodeError.TooManyDevicesError) {
-        onRemoveDeviceOpen()
-    }
 }
 
 @Composable
-private fun CreateAccountCodeViewState.CodeError.DialogError.getResources(type: CreateAccountFlowType) = when (this) {
-    CreateAccountCodeViewState.CodeError.DialogError.AccountAlreadyExistsError -> DialogErrorStrings(
+private fun CreateAccountCodeViewState.Result.Error.DialogError.getResources(type: CreateAccountFlowType) = when (this) {
+    CreateAccountCodeViewState.Result.Error.DialogError.AccountAlreadyExistsError -> DialogErrorStrings(
         stringResource(id = R.string.create_account_code_error_title),
         stringResource(id = R.string.create_account_email_already_in_use_error)
     )
 
-    CreateAccountCodeViewState.CodeError.DialogError.BlackListedError -> DialogErrorStrings(
+    CreateAccountCodeViewState.Result.Error.DialogError.BlackListedError -> DialogErrorStrings(
         stringResource(id = R.string.create_account_code_error_title),
         stringResource(id = R.string.create_account_email_blacklisted_error)
     )
 
-    CreateAccountCodeViewState.CodeError.DialogError.EmailDomainBlockedError -> DialogErrorStrings(
+    CreateAccountCodeViewState.Result.Error.DialogError.EmailDomainBlockedError -> DialogErrorStrings(
         stringResource(id = R.string.create_account_code_error_title),
         stringResource(id = R.string.create_account_email_domain_blocked_error)
     )
 
-    CreateAccountCodeViewState.CodeError.DialogError.InvalidEmailError -> DialogErrorStrings(
+    CreateAccountCodeViewState.Result.Error.DialogError.InvalidEmailError -> DialogErrorStrings(
         stringResource(id = R.string.create_account_code_error_title),
         stringResource(id = R.string.create_account_email_invalid_error)
     )
 
-    CreateAccountCodeViewState.CodeError.DialogError.TeamMembersLimitError -> DialogErrorStrings(
+    CreateAccountCodeViewState.Result.Error.DialogError.TeamMembersLimitError -> DialogErrorStrings(
         stringResource(id = R.string.create_account_code_error_title),
         stringResource(id = R.string.create_account_code_error_team_members_limit_reached)
     )
 
-    CreateAccountCodeViewState.CodeError.DialogError.CreationRestrictedError -> DialogErrorStrings(
+    CreateAccountCodeViewState.Result.Error.DialogError.CreationRestrictedError -> DialogErrorStrings(
         stringResource(id = R.string.create_account_code_error_title),
         stringResource(
             id = when (type) {
@@ -234,15 +233,21 @@ private fun CreateAccountCodeViewState.CodeError.DialogError.getResources(type:
         )
     )
     // TODO: sync with design about the error message
-    CreateAccountCodeViewState.CodeError.DialogError.UserAlreadyExists ->
+    CreateAccountCodeViewState.Result.Error.DialogError.UserAlreadyExistsError ->
         DialogErrorStrings("User Already LoggedIn", "UserAlreadyLoggedIn")
 
-    is CreateAccountCodeViewState.CodeError.DialogError.GenericError ->
+    is CreateAccountCodeViewState.Result.Error.DialogError.GenericError ->
         this.coreFailure.dialogErrorStrings(LocalContext.current.resources)
 }
 
 @Composable
 @Preview
 fun PreviewCreateAccountCodeScreen() {
-    CodeContent(CreateAccountCodeViewState(CreateAccountFlowType.CreatePersonalAccount), {}, {}, {}, {}, {}, ServerConfig.DEFAULT)
+    CodeContent(
+        textState = TextFieldState(),
+        state = CreateAccountCodeViewState(CreateAccountFlowType.CreatePersonalAccount),
+        onResendCodePressed = {},
+        onBackPressed = {},
+        serverConfig = ServerConfig.DEFAULT
+    )
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewModel.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewModel.kt
index b2e1adf74e6..1d6d7b18dce 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewModel.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewModel.kt
@@ -17,10 +17,11 @@
  */
 package com.wire.android.ui.authentication.create.code
 
+import androidx.compose.foundation.text.input.TextFieldState
+import androidx.compose.foundation.text.input.clearText
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
@@ -30,7 +31,7 @@ import com.wire.android.di.ClientScopeProvider
 import com.wire.android.di.KaliumCoreLogic
 import com.wire.android.ui.authentication.create.common.CreateAccountFlowType
 import com.wire.android.ui.authentication.create.common.CreateAccountNavArgs
-import com.wire.android.ui.common.textfield.CodeFieldValue
+import com.wire.android.ui.common.textfield.textAsFlow
 import com.wire.android.ui.navArgs
 import com.wire.android.util.WillNeverOccurError
 import com.wire.kalium.logic.CoreLogic
@@ -44,6 +45,7 @@ import com.wire.kalium.logic.feature.register.RegisterParam
 import com.wire.kalium.logic.feature.register.RegisterResult
 import com.wire.kalium.logic.feature.register.RequestActivationCodeResult
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -61,11 +63,15 @@ class CreateAccountCodeViewModel @Inject constructor(
 
     val serverConfig: ServerConfig.Links = authServerConfigProvider.authServer.value
 
+    val codeTextState: TextFieldState = TextFieldState()
     var codeState: CreateAccountCodeViewState by mutableStateOf(CreateAccountCodeViewState(createAccountNavArgs.flowType))
 
-    fun onCodeChange(newValue: CodeFieldValue, onSuccess: () -> Unit) {
-        codeState = codeState.copy(code = newValue, error = CreateAccountCodeViewState.CodeError.None)
-        if (newValue.isFullyFilled) onCodeContinue(onSuccess)
+    init {
+        viewModelScope.launch {
+            codeTextState.textAsFlow().collectLatest {
+                if (it.length == codeState.codeLength) onCodeContinue()
+            }
+        }
     }
 
     fun resendCode() {
@@ -92,17 +98,17 @@ class CreateAccountCodeViewModel @Inject constructor(
                 }
             }
 
-            val codeError = authScope.registerScope.requestActivationCode(createAccountNavArgs.userRegistrationInfo.email).toCodeError()
-            codeState = codeState.copy(loading = false, error = codeError)
+            val result = authScope.registerScope.requestActivationCode(createAccountNavArgs.userRegistrationInfo.email).toCodeError()
+            codeState = codeState.copy(loading = false, result = result)
         }
     }
 
     fun clearCodeError() {
-        codeState = codeState.copy(error = CreateAccountCodeViewState.CodeError.None)
+        codeState = codeState.copy(result = CreateAccountCodeViewState.Result.None)
     }
 
     fun clearCodeField() {
-        codeState = codeState.copy(code = CodeFieldValue(text = TextFieldValue(""), isFullyFilled = false))
+        codeTextState.clearText()
     }
 
     private fun registerParamFromType() = when (createAccountNavArgs.flowType) {
@@ -112,7 +118,7 @@ class CreateAccountCodeViewModel @Inject constructor(
                 lastName = createAccountNavArgs.userRegistrationInfo.lastName,
                 password = createAccountNavArgs.userRegistrationInfo.password,
                 email = createAccountNavArgs.userRegistrationInfo.email,
-                emailActivationCode = codeState.code.text.text
+                emailActivationCode = codeTextState.text.toString()
             )
 
         CreateAccountFlowType.CreateTeam ->
@@ -121,14 +127,14 @@ class CreateAccountCodeViewModel @Inject constructor(
                 lastName = createAccountNavArgs.userRegistrationInfo.lastName,
                 password = createAccountNavArgs.userRegistrationInfo.password,
                 email = createAccountNavArgs.userRegistrationInfo.email,
-                emailActivationCode = codeState.code.text.text,
+                emailActivationCode = codeTextState.text.toString(),
                 teamName = createAccountNavArgs.userRegistrationInfo.teamName,
                 teamIcon = "default"
             )
     }
 
     @Suppress("ComplexMethod")
-    private fun onCodeContinue(onSuccess: () -> Unit) {
+    private fun onCodeContinue() {
         codeState = codeState.copy(loading = true)
         viewModelScope.launch {
             // create account does not support proxy yet
@@ -188,24 +194,20 @@ class CreateAccountCodeViewModel @Inject constructor(
                     }
 
                     is RegisterClientResult.Success -> {
-                        onSuccess()
+                        codeState = codeState.copy(result = CreateAccountCodeViewState.Result.Success)
                     }
 
                     is RegisterClientResult.E2EICertificateRequired -> {
                         // TODO
-                        onSuccess()
+                        codeState = codeState.copy(result = CreateAccountCodeViewState.Result.Success)
                     }
                 }
             }
         }
     }
 
-    private fun updateCodeErrorState(codeError: CreateAccountCodeViewState.CodeError) {
-        codeState = if (codeError is CreateAccountCodeViewState.CodeError.None) {
-            codeState.copy(error = codeError)
-        } else {
-            codeState.copy(loading = false, error = codeError)
-        }
+    private fun updateCodeErrorState(codeError: CreateAccountCodeViewState.Result.Error) {
+        codeState = codeState.copy(loading = false, result = codeError)
     }
 
     private suspend fun registerClient(userId: UserId, password: String) =
@@ -218,8 +220,8 @@ class CreateAccountCodeViewModel @Inject constructor(
         )
 
     private fun RegisterClientResult.Failure.toCodeError() = when (this) {
-        is RegisterClientResult.Failure.TooManyClients -> CreateAccountCodeViewState.CodeError.TooManyDevicesError
-        is RegisterClientResult.Failure.Generic -> CreateAccountCodeViewState.CodeError.DialogError.GenericError(this.genericFailure)
+        is RegisterClientResult.Failure.TooManyClients -> CreateAccountCodeViewState.Result.Error.TooManyDevicesError
+        is RegisterClientResult.Failure.Generic -> CreateAccountCodeViewState.Result.Error.DialogError.GenericError(this.genericFailure)
         is RegisterClientResult.Failure.InvalidCredentials ->
             throw WillNeverOccurError("RegisterClient: wrong password when register client after creating a new account")
 
@@ -228,29 +230,30 @@ class CreateAccountCodeViewModel @Inject constructor(
     }
 
     private fun RegisterResult.Failure.toCodeError() = when (this) {
-        RegisterResult.Failure.InvalidActivationCode -> CreateAccountCodeViewState.CodeError.TextFieldError.InvalidActivationCodeError
-        RegisterResult.Failure.AccountAlreadyExists -> CreateAccountCodeViewState.CodeError.DialogError.AccountAlreadyExistsError
-        RegisterResult.Failure.BlackListed -> CreateAccountCodeViewState.CodeError.DialogError.BlackListedError
-        RegisterResult.Failure.EmailDomainBlocked -> CreateAccountCodeViewState.CodeError.DialogError.EmailDomainBlockedError
-        RegisterResult.Failure.InvalidEmail -> CreateAccountCodeViewState.CodeError.DialogError.InvalidEmailError
-        RegisterResult.Failure.TeamMembersLimitReached -> CreateAccountCodeViewState.CodeError.DialogError.TeamMembersLimitError
-        RegisterResult.Failure.UserCreationRestricted -> CreateAccountCodeViewState.CodeError.DialogError.CreationRestrictedError
-        is RegisterResult.Failure.Generic -> CreateAccountCodeViewState.CodeError.DialogError.GenericError(this.failure)
+        RegisterResult.Failure.InvalidActivationCode -> CreateAccountCodeViewState.Result.Error.TextFieldError.InvalidActivationCodeError
+        RegisterResult.Failure.AccountAlreadyExists -> CreateAccountCodeViewState.Result.Error.DialogError.AccountAlreadyExistsError
+        RegisterResult.Failure.BlackListed -> CreateAccountCodeViewState.Result.Error.DialogError.BlackListedError
+        RegisterResult.Failure.EmailDomainBlocked -> CreateAccountCodeViewState.Result.Error.DialogError.EmailDomainBlockedError
+        RegisterResult.Failure.InvalidEmail -> CreateAccountCodeViewState.Result.Error.DialogError.InvalidEmailError
+        RegisterResult.Failure.TeamMembersLimitReached -> CreateAccountCodeViewState.Result.Error.DialogError.TeamMembersLimitError
+        RegisterResult.Failure.UserCreationRestricted -> CreateAccountCodeViewState.Result.Error.DialogError.CreationRestrictedError
+        is RegisterResult.Failure.Generic -> CreateAccountCodeViewState.Result.Error.DialogError.GenericError(this.failure)
     }
 
     private fun AddAuthenticatedUserUseCase.Result.Failure.toCodeError() = when (this) {
         is AddAuthenticatedUserUseCase.Result.Failure.Generic ->
-            CreateAccountCodeViewState.CodeError.DialogError.GenericError(this.genericFailure)
+            CreateAccountCodeViewState.Result.Error.DialogError.GenericError(this.genericFailure)
 
-        AddAuthenticatedUserUseCase.Result.Failure.UserAlreadyExists -> CreateAccountCodeViewState.CodeError.DialogError.UserAlreadyExists
+        AddAuthenticatedUserUseCase.Result.Failure.UserAlreadyExists ->
+            CreateAccountCodeViewState.Result.Error.DialogError.UserAlreadyExistsError
     }
 
     private fun RequestActivationCodeResult.toCodeError() = when (this) {
-        RequestActivationCodeResult.Failure.AlreadyInUse -> CreateAccountCodeViewState.CodeError.DialogError.AccountAlreadyExistsError
-        RequestActivationCodeResult.Failure.BlacklistedEmail -> CreateAccountCodeViewState.CodeError.DialogError.BlackListedError
-        RequestActivationCodeResult.Failure.DomainBlocked -> CreateAccountCodeViewState.CodeError.DialogError.EmailDomainBlockedError
-        RequestActivationCodeResult.Failure.InvalidEmail -> CreateAccountCodeViewState.CodeError.DialogError.InvalidEmailError
-        is RequestActivationCodeResult.Failure.Generic -> CreateAccountCodeViewState.CodeError.DialogError.GenericError(this.failure)
-        RequestActivationCodeResult.Success -> CreateAccountCodeViewState.CodeError.None
+        RequestActivationCodeResult.Failure.AlreadyInUse -> CreateAccountCodeViewState.Result.Error.DialogError.AccountAlreadyExistsError
+        RequestActivationCodeResult.Failure.BlacklistedEmail -> CreateAccountCodeViewState.Result.Error.DialogError.BlackListedError
+        RequestActivationCodeResult.Failure.DomainBlocked -> CreateAccountCodeViewState.Result.Error.DialogError.EmailDomainBlockedError
+        RequestActivationCodeResult.Failure.InvalidEmail -> CreateAccountCodeViewState.Result.Error.DialogError.InvalidEmailError
+        is RequestActivationCodeResult.Failure.Generic -> CreateAccountCodeViewState.Result.Error.DialogError.GenericError(this.failure)
+        RequestActivationCodeResult.Success -> CreateAccountCodeViewState.Result.None
     }
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewState.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewState.kt
index 3ab9b4e59b6..e11175f8a22 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewState.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/code/CreateAccountCodeViewState.kt
@@ -18,35 +18,38 @@
 
 package com.wire.android.ui.authentication.create.code
 
-import androidx.compose.ui.text.input.TextFieldValue
 import com.wire.android.ui.authentication.create.common.CreateAccountFlowType
-import com.wire.android.ui.common.textfield.CodeFieldValue
 import com.wire.kalium.logic.CoreFailure
 
 data class CreateAccountCodeViewState(
     val type: CreateAccountFlowType,
-    val code: CodeFieldValue = CodeFieldValue(TextFieldValue(""), false),
+    val codeLength: Int = DEFAULT_VERIFICATION_CODE_LENGTH,
     val email: String = "",
     val loading: Boolean = false,
-    val error: CodeError = CodeError.None
+    val result: Result = Result.None,
 ) {
-    sealed class CodeError {
-        object None : CodeError()
-        sealed class TextFieldError : CodeError() {
-            object InvalidActivationCodeError : TextFieldError()
-        }
+    sealed interface Result {
+        data object None : Result
+        data object Success : Result
+        sealed class Error : Result {
+            sealed class TextFieldError : Error() {
+                data object InvalidActivationCodeError : TextFieldError()
+            }
 
-        sealed class DialogError : CodeError() {
-            object InvalidEmailError : DialogError()
-            object AccountAlreadyExistsError : DialogError()
-            object BlackListedError : DialogError()
-            object EmailDomainBlockedError : DialogError()
-            object TeamMembersLimitError : DialogError()
-            object CreationRestrictedError : DialogError()
-            object UserAlreadyExists: DialogError()
-            data class GenericError(val coreFailure: CoreFailure) : DialogError()
+            sealed class DialogError : Error() {
+                data object InvalidEmailError : DialogError()
+                data object AccountAlreadyExistsError : DialogError()
+                data object BlackListedError : DialogError()
+                data object EmailDomainBlockedError : DialogError()
+                data object TeamMembersLimitError : DialogError()
+                data object CreationRestrictedError : DialogError()
+                data object UserAlreadyExistsError : DialogError()
+                data class GenericError(val coreFailure: CoreFailure) : DialogError()
+            }
+            data object TooManyDevicesError : Error()
         }
-
-        object TooManyDevicesError : CodeError()
+    }
+    companion object {
+        const val DEFAULT_VERIFICATION_CODE_LENGTH = 6
     }
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsScreen.kt
index 6fbbce6e6e6..7cf62ffbc67 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsScreen.kt
@@ -25,6 +25,7 @@ import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.rememberScrollState
 import androidx.compose.foundation.text.KeyboardOptions
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Surface
@@ -42,8 +43,6 @@ import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.input.ImeAction
 import androidx.compose.ui.text.input.KeyboardCapitalization
 import androidx.compose.ui.text.input.KeyboardType
-import androidx.compose.ui.text.input.TextFieldValue
-import androidx.compose.ui.tooling.preview.Preview
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.ramcosta.composedestinations.annotation.Destination
 import com.wire.android.R
@@ -60,14 +59,17 @@ import com.wire.android.ui.common.error.CoreFailureErrorDialog
 import com.wire.android.ui.common.rememberBottomBarElevationState
 import com.wire.android.ui.common.rememberTopBarElevationState
 import com.wire.android.ui.common.scaffold.WireScaffold
+import com.wire.android.ui.common.textfield.DefaultPassword
 import com.wire.android.ui.common.textfield.WirePasswordTextField
 import com.wire.android.ui.common.textfield.WireTextField
 import com.wire.android.ui.common.textfield.WireTextFieldState
 import com.wire.android.ui.common.topappbar.WireCenterAlignedTopAppBar
 import com.wire.android.ui.destinations.CreateAccountCodeScreenDestination
+import com.wire.android.ui.theme.WireTheme
 import com.wire.android.ui.theme.wireColorScheme
 import com.wire.android.ui.theme.wireDimensions
 import com.wire.android.ui.theme.wireTypography
+import com.wire.android.util.ui.PreviewMultipleThemes
 import com.wire.kalium.logic.configuration.server.ServerConfig
 
 @CreatePersonalAccountNavGraph
@@ -84,10 +86,10 @@ fun CreateAccountDetailsScreen(
                 CreateAccountCodeScreenDestination(
                     createAccountNavArgs.copy(
                         userRegistrationInfo = createAccountNavArgs.userRegistrationInfo.copy(
-                            firstName = detailsState.firstName.text.trim(),
-                            lastName = detailsState.lastName.text.trim(),
-                            password = detailsState.password.text,
-                            teamName = detailsState.teamName.text.trim()
+                            firstName = firstNameTextState.text.toString().trim(),
+                            lastName = lastNameTextState.text.toString().trim(),
+                            password = passwordTextState.text.toString(),
+                            teamName = teamNameTextState.text.toString().trim()
                         )
                     )
                 )
@@ -96,21 +98,11 @@ fun CreateAccountDetailsScreen(
 
         DetailsContent(
             state = detailsState,
-            onFirstNameChange = {
-                onDetailsChange(
-                    it,
-                    CreateAccountDetailsViewModel.DetailsFieldType.FirstName
-                )
-            },
-            onLastNameChange = { onDetailsChange(it, CreateAccountDetailsViewModel.DetailsFieldType.LastName) },
-            onPasswordChange = { onDetailsChange(it, CreateAccountDetailsViewModel.DetailsFieldType.Password) },
-            onConfirmPasswordChange = {
-                onDetailsChange(
-                    it,
-                    CreateAccountDetailsViewModel.DetailsFieldType.ConfirmPassword
-                )
-            },
-            onTeamNameChange = { onDetailsChange(it, CreateAccountDetailsViewModel.DetailsFieldType.TeamName) },
+            firstNameTextState = firstNameTextState,
+            lastNameTextState = lastNameTextState,
+            passwordTextState = passwordTextState,
+            confirmPasswordTextState = confirmPasswordTextState,
+            teamNameTextState = teamNameTextState,
             onBackPressed = navigator::navigateBack,
             onContinuePressed = { onDetailsContinue(::navigateToCodeScreen) },
             onErrorDismiss = ::onDetailsErrorDismiss,
@@ -122,11 +114,11 @@ fun CreateAccountDetailsScreen(
 @Composable
 private fun DetailsContent(
     state: CreateAccountDetailsViewState,
-    onFirstNameChange: (TextFieldValue) -> Unit,
-    onLastNameChange: (TextFieldValue) -> Unit,
-    onPasswordChange: (TextFieldValue) -> Unit,
-    onConfirmPasswordChange: (TextFieldValue) -> Unit,
-    onTeamNameChange: (TextFieldValue) -> Unit,
+    firstNameTextState: TextFieldState,
+    lastNameTextState: TextFieldState,
+    passwordTextState: TextFieldState,
+    confirmPasswordTextState: TextFieldState,
+    teamNameTextState: TextFieldState,
     onBackPressed: () -> Unit,
     onContinuePressed: () -> Unit,
     onErrorDismiss: () -> Unit,
@@ -179,8 +171,7 @@ private fun DetailsContent(
                 )
 
                 WireTextField(
-                    value = state.firstName,
-                    onValueChange = onFirstNameChange,
+                    textState = firstNameTextState,
                     placeholderText = stringResource(R.string.create_account_details_first_name_placeholder),
                     labelText = stringResource(R.string.create_account_details_first_name_label),
                     labelMandatoryIcon = true,
@@ -197,8 +188,7 @@ private fun DetailsContent(
                 )
 
                 WireTextField(
-                    value = state.lastName,
-                    onValueChange = onLastNameChange,
+                    textState = lastNameTextState,
                     placeholderText = stringResource(R.string.create_account_details_last_name_placeholder),
                     labelText = stringResource(R.string.create_account_details_last_name_label),
                     labelMandatoryIcon = true,
@@ -215,8 +205,7 @@ private fun DetailsContent(
 
                 if (state.type == CreateAccountFlowType.CreateTeam) {
                     WireTextField(
-                        value = state.teamName,
-                        onValueChange = onTeamNameChange,
+                        textState = teamNameTextState,
                         placeholderText = stringResource(R.string.create_account_details_team_name_placeholder),
                         labelText = stringResource(R.string.create_account_details_team_name_label),
                         labelMandatoryIcon = true,
@@ -233,11 +222,10 @@ private fun DetailsContent(
                 }
 
                 WirePasswordTextField(
-                    value = state.password,
-                    onValueChange = onPasswordChange,
+                    textState = passwordTextState,
                     labelMandatoryIcon = true,
                     descriptionText = stringResource(R.string.create_account_details_password_description),
-                    imeAction = ImeAction.Next,
+                    keyboardOptions = KeyboardOptions.DefaultPassword.copy(imeAction = ImeAction.Next),
                     modifier = Modifier
                         .padding(horizontal = MaterialTheme.wireDimensions.spacing16x)
                         .testTag("password"),
@@ -246,16 +234,15 @@ private fun DetailsContent(
                     } else {
                         WireTextFieldState.Default
                     },
-                    autofill = false,
+                    autoFill = false,
                 )
 
                 WirePasswordTextField(
-                    value = state.confirmPassword,
-                    onValueChange = onConfirmPasswordChange,
+                    textState = confirmPasswordTextState,
                     labelText = stringResource(R.string.create_account_details_confirm_password_label),
                     labelMandatoryIcon = true,
-                    imeAction = ImeAction.Done,
-                    onImeAction = { keyboardController?.hide() },
+                    keyboardOptions = KeyboardOptions.DefaultPassword.copy(imeAction = ImeAction.Done),
+                    onKeyboardAction = { keyboardController?.hide() },
                     modifier = Modifier
                         .padding(
                             horizontal = MaterialTheme.wireDimensions.spacing16x,
@@ -269,7 +256,7 @@ private fun DetailsContent(
                         CreateAccountDetailsViewState.DetailsError.TextFieldError.InvalidPasswordError ->
                             WireTextFieldState.Error(stringResource(id = R.string.create_account_details_password_error))
                     } else WireTextFieldState.Default,
-                    autofill = false,
+                    autoFill = false,
                 )
             }
 
@@ -301,7 +288,18 @@ private fun DetailsContent(
 }
 
 @Composable
-@Preview
-fun PreviewCreateAccountDetailsScreen() {
-    DetailsContent(CreateAccountDetailsViewState(CreateAccountFlowType.CreateTeam), {}, {}, {}, {}, {}, {}, {}, {}, ServerConfig.DEFAULT)
+@PreviewMultipleThemes
+fun PreviewCreateAccountDetailsScreen() = WireTheme {
+    DetailsContent(
+        state = CreateAccountDetailsViewState(CreateAccountFlowType.CreateTeam),
+        firstNameTextState = TextFieldState(),
+        lastNameTextState = TextFieldState(),
+        passwordTextState = TextFieldState(),
+        confirmPasswordTextState = TextFieldState(),
+        teamNameTextState = TextFieldState(),
+        onBackPressed = {},
+        onContinuePressed = {},
+        onErrorDismiss = {},
+        serverConfig = ServerConfig.DEFAULT
+    )
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewModel.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewModel.kt
index f71384d8d35..98c3352a3e1 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewModel.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewModel.kt
@@ -17,19 +17,22 @@
  */
 package com.wire.android.ui.authentication.create.details
 
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.wire.android.di.AuthServerConfigProvider
+import com.wire.android.ui.authentication.create.common.CreateAccountFlowType
 import com.wire.android.ui.authentication.create.common.CreateAccountNavArgs
+import com.wire.android.ui.common.textfield.textAsFlow
 import com.wire.android.ui.navArgs
 import com.wire.kalium.logic.configuration.server.ServerConfig
 import com.wire.kalium.logic.feature.auth.ValidatePasswordUseCase
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -43,22 +46,32 @@ class CreateAccountDetailsViewModel @Inject constructor(
 
     val createAccountNavArgs: CreateAccountNavArgs = savedStateHandle.navArgs()
 
+    val firstNameTextState: TextFieldState = TextFieldState()
+    val lastNameTextState: TextFieldState = TextFieldState()
+    val passwordTextState: TextFieldState = TextFieldState()
+    val confirmPasswordTextState: TextFieldState = TextFieldState()
+    val teamNameTextState: TextFieldState = TextFieldState()
     var detailsState: CreateAccountDetailsViewState by mutableStateOf(CreateAccountDetailsViewState(createAccountNavArgs.flowType))
 
     val serverConfig: ServerConfig.Links = authServerConfigProvider.authServer.value
 
-    fun onDetailsChange(newText: TextFieldValue, fieldType: DetailsFieldType) {
-        detailsState = when (fieldType) {
-            DetailsFieldType.FirstName -> detailsState.copy(firstName = newText)
-            DetailsFieldType.LastName -> detailsState.copy(lastName = newText)
-            DetailsFieldType.Password -> detailsState.copy(password = newText)
-            DetailsFieldType.ConfirmPassword -> detailsState.copy(confirmPassword = newText)
-            DetailsFieldType.TeamName -> detailsState.copy(teamName = newText)
-        }.let {
-            it.copy(
-                error = CreateAccountDetailsViewState.DetailsError.None,
-                continueEnabled = it.fieldsNotEmpty() && !it.loading
-            )
+    init {
+        viewModelScope.launch {
+            combine(
+                firstNameTextState.textAsFlow(),
+                lastNameTextState.textAsFlow(),
+                passwordTextState.textAsFlow(),
+                confirmPasswordTextState.textAsFlow(),
+                teamNameTextState.textAsFlow(),
+            ) { firstName, lastName, password, confirmPassword, teamName ->
+                firstName.isNotBlank() && lastName.isNotBlank() && password.isNotBlank() && confirmPassword.isNotBlank()
+                        && (detailsState.type == CreateAccountFlowType.CreatePersonalAccount || teamName.isNotBlank())
+            }.collect { fieldsNotEmpty ->
+                detailsState = detailsState.copy(
+                    error = CreateAccountDetailsViewState.DetailsError.None,
+                    continueEnabled = fieldsNotEmpty && !detailsState.loading
+                )
+            }
         }
     }
 
@@ -66,10 +79,10 @@ class CreateAccountDetailsViewModel @Inject constructor(
         detailsState = detailsState.copy(loading = true, continueEnabled = false)
         viewModelScope.launch {
             val detailsError = when {
-                !validatePasswordUseCase(detailsState.password.text).isValid ->
+                !validatePasswordUseCase(passwordTextState.text.toString()).isValid ->
                     CreateAccountDetailsViewState.DetailsError.TextFieldError.InvalidPasswordError
 
-                detailsState.password.text != detailsState.confirmPassword.text ->
+                passwordTextState.text.toString() != confirmPasswordTextState.text.toString() ->
                     CreateAccountDetailsViewState.DetailsError.TextFieldError.PasswordsNotMatchingError
 
                 else -> CreateAccountDetailsViewState.DetailsError.None
@@ -86,8 +99,4 @@ class CreateAccountDetailsViewModel @Inject constructor(
     fun onDetailsErrorDismiss() {
         detailsState = detailsState.copy(error = CreateAccountDetailsViewState.DetailsError.None)
     }
-
-    enum class DetailsFieldType {
-        FirstName, LastName, Password, ConfirmPassword, TeamName
-    }
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewState.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewState.kt
index 83241ac4669..d5ca77a8021 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewState.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/details/CreateAccountDetailsViewState.kt
@@ -18,30 +18,20 @@
 
 package com.wire.android.ui.authentication.create.details
 
-import androidx.compose.ui.text.input.TextFieldValue
 import com.wire.android.ui.authentication.create.common.CreateAccountFlowType
 import com.wire.kalium.logic.NetworkFailure
 
 data class CreateAccountDetailsViewState(
     val type: CreateAccountFlowType,
-    val firstName: TextFieldValue = TextFieldValue(""),
-    val lastName: TextFieldValue = TextFieldValue(""),
-    val password: TextFieldValue = TextFieldValue(""),
-    val confirmPassword: TextFieldValue = TextFieldValue(""),
-    val teamName: TextFieldValue = TextFieldValue(""),
     val continueEnabled: Boolean = false,
     val loading: Boolean = false,
     val error: DetailsError = DetailsError.None
 ) {
-    fun fieldsNotEmpty(): Boolean =
-        firstName.text.isNotBlank() && lastName.text.isNotBlank() && password.text.isNotBlank() && confirmPassword.text.isNotBlank()
-                && (type == CreateAccountFlowType.CreatePersonalAccount || teamName.text.isNotBlank())
-
     sealed class DetailsError {
-        object None : DetailsError()
+        data object None : DetailsError()
         sealed class TextFieldError : DetailsError() {
-            object InvalidPasswordError : TextFieldError()
-            object PasswordsNotMatchingError : TextFieldError()
+            data object InvalidPasswordError : TextFieldError()
+            data object PasswordsNotMatchingError : TextFieldError()
         }
 
         sealed class DialogError : DetailsError() {
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailScreen.kt
index 5b007c6e4b5..e413f6bc9bc 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailScreen.kt
@@ -29,14 +29,13 @@ import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.text.ClickableText
 import androidx.compose.foundation.text.KeyboardOptions
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.material3.MaterialTheme
-import com.wire.android.ui.common.scaffold.WireScaffold
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Alignment
-import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.focus.FocusRequester
 import androidx.compose.ui.focus.focusRequester
@@ -48,11 +47,9 @@ import androidx.compose.ui.text.SpanStyle
 import androidx.compose.ui.text.buildAnnotatedString
 import androidx.compose.ui.text.input.ImeAction
 import androidx.compose.ui.text.input.KeyboardType
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.text.style.TextDecoration
 import androidx.compose.ui.text.withStyle
-import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.unit.dp
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.ramcosta.composedestinations.annotation.Destination
@@ -73,15 +70,18 @@ import com.wire.android.ui.common.button.WireButtonState
 import com.wire.android.ui.common.button.WirePrimaryButton
 import com.wire.android.ui.common.button.WireSecondaryButton
 import com.wire.android.ui.common.error.CoreFailureErrorDialog
+import com.wire.android.ui.common.scaffold.WireScaffold
 import com.wire.android.ui.common.textfield.WireTextField
 import com.wire.android.ui.common.textfield.WireTextFieldState
 import com.wire.android.ui.common.topappbar.WireCenterAlignedTopAppBar
 import com.wire.android.ui.destinations.CreateAccountDetailsScreenDestination
 import com.wire.android.ui.destinations.LoginScreenDestination
+import com.wire.android.ui.theme.WireTheme
 import com.wire.android.ui.theme.wireColorScheme
 import com.wire.android.ui.theme.wireDimensions
 import com.wire.android.ui.theme.wireTypography
 import com.wire.android.util.CustomTabsHelper
+import com.wire.android.util.ui.PreviewMultipleThemes
 import com.wire.kalium.logic.configuration.server.ServerConfig
 
 @CreatePersonalAccountNavGraph
@@ -98,7 +98,7 @@ fun CreateAccountEmailScreen(
                 CreateAccountDetailsScreenDestination(
                     navArgs = createAccountNavArgs.copy(
                         userRegistrationInfo = UserRegistrationInfo(
-                            email = emailState.email.text.trim().lowercase()
+                            email = emailTextState.text.trim().toString().lowercase()
                         )
                     )
                 )
@@ -107,7 +107,7 @@ fun CreateAccountEmailScreen(
 
         EmailContent(
             state = emailState,
-            onEmailChange = ::onEmailChange,
+            emailTextState = emailTextState,
             onBackPressed = navigator::navigateBack,
             onContinuePressed = { onEmailContinue(::navigateToDetailsScreen) },
             onLoginPressed = { navigator.navigate(NavigationCommand(LoginScreenDestination(), BackStackMode.CLEAR_TILL_START)) },
@@ -120,11 +120,10 @@ fun CreateAccountEmailScreen(
     }
 }
 
-@OptIn(ExperimentalComposeUiApi::class)
 @Composable
 private fun EmailContent(
     state: CreateAccountEmailViewState,
-    onEmailChange: (TextFieldValue) -> Unit,
+    emailTextState: TextFieldState,
     onBackPressed: () -> Unit,
     onContinuePressed: () -> Unit,
     onLoginPressed: () -> Unit,
@@ -169,8 +168,7 @@ private fun EmailContent(
                     .testTag("createTeamText")
             )
             WireTextField(
-                value = state.email,
-                onValueChange = onEmailChange,
+                textState = emailTextState,
                 placeholderText = stringResource(R.string.create_account_email_placeholder),
                 labelText = stringResource(R.string.create_account_email_label),
                 state = if (state.error is CreateAccountEmailViewState.EmailError.None) WireTextFieldState.Default
@@ -328,9 +326,18 @@ private fun TermsConditionsDialog(onDialogDismiss: () -> Unit, onContinuePressed
 }
 
 @Composable
-@Preview
-fun PreviewCreateAccountEmailScreen() {
-    EmailContent(CreateAccountEmailViewState(CreateAccountFlowType.CreatePersonalAccount), {}, {}, {}, {}, {}, {}, {}, "",
-        ServerConfig.DEFAULT
+@PreviewMultipleThemes
+fun PreviewCreateAccountEmailScreen() = WireTheme {
+    EmailContent(
+        emailTextState = TextFieldState(),
+        state = CreateAccountEmailViewState(CreateAccountFlowType.CreatePersonalAccount),
+        onBackPressed = {},
+        onContinuePressed = {},
+        onLoginPressed = {},
+        onTermsDialogDismiss = {},
+        onTermsAccept = {},
+        onErrorDismiss = {},
+        tosUrl = "",
+        serverConfig = ServerConfig.DEFAULT
     )
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewModel.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewModel.kt
index 5656091bdc5..406faf4ebf9 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewModel.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewModel.kt
@@ -17,16 +17,17 @@
  */
 package com.wire.android.ui.authentication.create.email
 
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.wire.android.di.AuthServerConfigProvider
 import com.wire.android.di.KaliumCoreLogic
 import com.wire.android.ui.authentication.create.common.CreateAccountNavArgs
+import com.wire.android.ui.common.textfield.textAsFlow
 import com.wire.android.ui.navArgs
 import com.wire.kalium.logic.CoreLogic
 import com.wire.kalium.logic.configuration.server.ServerConfig
@@ -34,6 +35,7 @@ import com.wire.kalium.logic.feature.auth.ValidateEmailUseCase
 import com.wire.kalium.logic.feature.auth.autoVersioningAuth.AutoVersionAuthScopeUseCase
 import com.wire.kalium.logic.feature.register.RequestActivationCodeResult
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -48,6 +50,7 @@ class CreateAccountEmailViewModel @Inject constructor(
 
     val createAccountNavArgs: CreateAccountNavArgs = savedStateHandle.navArgs()
 
+    val emailTextState: TextFieldState = TextFieldState()
     var emailState: CreateAccountEmailViewState by mutableStateOf(CreateAccountEmailViewState(createAccountNavArgs.flowType))
         private set
 
@@ -55,20 +58,25 @@ class CreateAccountEmailViewModel @Inject constructor(
 
     fun tosUrl(): String = authServerConfigProvider.authServer.value.tos
 
-    fun onEmailChange(newText: TextFieldValue) {
-        emailState = emailState.copy(
-            email = newText,
-            error = CreateAccountEmailViewState.EmailError.None,
-            continueEnabled = newText.text.isNotEmpty() && !emailState.loading
-        )
+    init {
+        viewModelScope.launch {
+            emailTextState.textAsFlow().collectLatest {
+                emailState = emailState.copy(
+                    error = CreateAccountEmailViewState.EmailError.None,
+                    continueEnabled = it.isNotEmpty() && !emailState.loading
+                )
+            }
+        }
     }
 
     fun onEmailContinue(onSuccess: () -> Unit) {
         emailState = emailState.copy(loading = true, continueEnabled = false)
         viewModelScope.launch {
-            val emailError =
-                if (validateEmail(emailState.email.text.trim().lowercase())) CreateAccountEmailViewState.EmailError.None
-                else CreateAccountEmailViewState.EmailError.TextFieldError.InvalidEmailError
+            val email = emailTextState.text.toString().trim().lowercase()
+            val emailError = when (validateEmail(email)) {
+                true -> CreateAccountEmailViewState.EmailError.None
+                false -> CreateAccountEmailViewState.EmailError.TextFieldError.InvalidEmailError
+            }
             emailState = emailState.copy(
                 loading = false,
                 continueEnabled = true,
@@ -104,7 +112,8 @@ class CreateAccountEmailViewModel @Inject constructor(
                 }
             }
 
-            val emailError = authScope.registerScope.requestActivationCode(emailState.email.text.trim().lowercase()).toEmailError()
+            val email = emailTextState.text.toString().trim().lowercase()
+            val emailError = authScope.registerScope.requestActivationCode(email).toEmailError()
             emailState = emailState.copy(loading = false, continueEnabled = true, error = emailError)
             if (emailError is CreateAccountEmailViewState.EmailError.None) onSuccess()
         }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewState.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewState.kt
index 91a700ad09f..280cb56ded7 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewState.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/create/email/CreateAccountEmailViewState.kt
@@ -18,13 +18,11 @@
 
 package com.wire.android.ui.authentication.create.email
 
-import androidx.compose.ui.text.input.TextFieldValue
 import com.wire.android.ui.authentication.create.common.CreateAccountFlowType
 import com.wire.kalium.logic.CoreFailure
 
 data class CreateAccountEmailViewState(
     val type: CreateAccountFlowType,
-    val email: TextFieldValue = TextFieldValue(""),
     val termsDialogVisible: Boolean = false,
     val termsAccepted: Boolean = false,
     val continueEnabled: Boolean = false,
@@ -34,17 +32,16 @@ data class CreateAccountEmailViewState(
     val showServerVersionNotSupportedDialog: Boolean = false
 ) {
     sealed class EmailError {
-        object None : EmailError()
+        data object None : EmailError()
         sealed class TextFieldError : EmailError() {
-            object InvalidEmailError : TextFieldError()
-            object BlacklistedEmailError : TextFieldError()
-            object AlreadyInUseError : TextFieldError()
-            object DomainBlockedError : TextFieldError()
+            data object InvalidEmailError : TextFieldError()
+            data object BlacklistedEmailError : TextFieldError()
+            data object AlreadyInUseError : TextFieldError()
+            data object DomainBlockedError : TextFieldError()
         }
 
         sealed class DialogError : EmailError() {
             data class GenericError(val coreFailure: CoreFailure) : DialogError()
         }
     }
-
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginError.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginError.kt
deleted file mode 100644
index 1d44962b8ac..00000000000
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginError.kt
+++ /dev/null
@@ -1,45 +0,0 @@
-/*
- * Wire
- * Copyright (C) 2024 Wire Swiss GmbH
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program. If not, see http://www.gnu.org/licenses/.
- */
-
-package com.wire.android.ui.authentication.login
-
-import com.wire.android.util.deeplink.SSOFailureCodes
-import com.wire.kalium.logic.CoreFailure
-
-sealed class LoginError {
-    object None : LoginError()
-    sealed class TextFieldError : LoginError() {
-        object InvalidValue : TextFieldError()
-    }
-
-    sealed class DialogError : LoginError() {
-        data class GenericError(val coreFailure: CoreFailure) : DialogError()
-        object InvalidCredentialsError : DialogError()
-        object ProxyError : DialogError()
-        object InvalidSSOCookie : DialogError()
-        object InvalidSSOCodeError : DialogError()
-        object UserAlreadyExists : DialogError()
-        object PasswordNeededToRegisterClient : DialogError()
-        data class SSOResultError constructor(val result: SSOFailureCodes) :
-            DialogError()
-        object ServerVersionNotSupported: DialogError()
-        object ClientUpdateRequired: DialogError()
-    }
-
-    object TooManyDevicesError : LoginError()
-}
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginScreen.kt
index 96327f343ec..28063a2dfbc 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginScreen.kt
@@ -20,7 +20,6 @@ package com.wire.android.ui.authentication.login
 
 import androidx.annotation.StringRes
 import androidx.compose.animation.AnimatedContent
-import androidx.compose.animation.ExperimentalAnimationApi
 import androidx.compose.animation.togetherWith
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.LocalOverscrollConfiguration
@@ -33,7 +32,6 @@ import androidx.compose.foundation.pager.HorizontalPager
 import androidx.compose.foundation.pager.rememberPagerState
 import androidx.compose.foundation.rememberScrollState
 import androidx.compose.material3.MaterialTheme
-import com.wire.android.ui.common.scaffold.WireScaffold
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.LaunchedEffect
@@ -42,14 +40,12 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
-import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalFocusManager
 import androidx.compose.ui.platform.LocalSoftwareKeyboardController
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.AnnotatedString
-import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.window.DialogProperties
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.ramcosta.composedestinations.annotation.Destination
@@ -73,6 +69,7 @@ import com.wire.android.ui.common.calculateCurrentTab
 import com.wire.android.ui.common.dialogs.FeatureDisabledWithProxyDialogContent
 import com.wire.android.ui.common.dialogs.FeatureDisabledWithProxyDialogState
 import com.wire.android.ui.common.rememberTopBarElevationState
+import com.wire.android.ui.common.scaffold.WireScaffold
 import com.wire.android.ui.common.topappbar.WireCenterAlignedTopAppBar
 import com.wire.android.ui.common.visbility.rememberVisibilityState
 import com.wire.android.ui.destinations.E2EIEnrollmentScreenDestination
@@ -84,6 +81,7 @@ import com.wire.android.ui.theme.wireDimensions
 import com.wire.android.ui.theme.wireTypography
 import com.wire.android.util.deeplink.DeepLinkResult
 import com.wire.android.util.dialogErrorStrings
+import com.wire.android.util.ui.PreviewMultipleThemes
 import com.wire.android.util.ui.UIText
 import kotlinx.coroutines.launch
 
@@ -95,33 +93,31 @@ import kotlinx.coroutines.launch
 fun LoginScreen(
     navigator: Navigator,
     loginNavArgs: LoginNavArgs,
-    loginViewModel: LoginViewModel = hiltViewModel(),
     loginEmailViewModel: LoginEmailViewModel = hiltViewModel()
 ) {
 
     LoginContent(
-        navigator::navigateBack,
-        { initialSyncCompleted, isE2EIRequired ->
+        onBackPressed = navigator::navigateBack,
+        onSuccess = { initialSyncCompleted, isE2EIRequired ->
             val destination = if (isE2EIRequired) E2EIEnrollmentScreenDestination
             else if (initialSyncCompleted) HomeScreenDestination
             else InitialSyncScreenDestination
 
             navigator.navigate(NavigationCommand(destination, BackStackMode.CLEAR_WHOLE))
         },
-        { navigator.navigate(NavigationCommand(RemoveDeviceScreenDestination, BackStackMode.CLEAR_WHOLE)) },
-        loginViewModel,
-        loginEmailViewModel,
+        onRemoveDeviceNeeded = {
+            navigator.navigate(NavigationCommand(RemoveDeviceScreenDestination, BackStackMode.CLEAR_WHOLE))
+        },
+        loginEmailViewModel = loginEmailViewModel,
         ssoLoginResult = loginNavArgs.ssoLoginResult
     )
 }
 
-@OptIn(ExperimentalAnimationApi::class)
 @Composable
 private fun LoginContent(
     onBackPressed: () -> Unit,
     onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit,
     onRemoveDeviceNeeded: () -> Unit,
-    viewModel: LoginViewModel,
     loginEmailViewModel: LoginEmailViewModel,
     ssoLoginResult: DeepLinkResult.SSOLogin?
 ) {
@@ -135,22 +131,21 @@ private fun LoginContent(
             targetState = loginEmailViewModel.secondFactorVerificationCodeState.isCodeInputNecessary,
             transitionSpec = { TransitionAnimationType.SLIDE.enterTransition.togetherWith(TransitionAnimationType.SLIDE.exitTransition) }
         ) { isCodeInputNecessary ->
-            if (isCodeInputNecessary) LoginEmailVerificationCodeScreen(onSuccess, loginEmailViewModel)
-            else MainLoginContent(onBackPressed, onSuccess, onRemoveDeviceNeeded, viewModel, loginEmailViewModel, ssoLoginResult)
+            if (isCodeInputNecessary) {
+                LoginEmailVerificationCodeScreen(loginEmailViewModel)
+            } else {
+                MainLoginContent(onBackPressed, onSuccess, onRemoveDeviceNeeded, loginEmailViewModel, ssoLoginResult)
+            }
         }
     }
 }
 
-@OptIn(
-    ExperimentalComposeUiApi::class,
-    ExperimentalFoundationApi::class,
-)
+@OptIn(ExperimentalFoundationApi::class)
 @Composable
 private fun MainLoginContent(
     onBackPressed: () -> Unit,
     onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit,
     onRemoveDeviceNeeded: () -> Unit,
-    viewModel: LoginViewModel,
     loginEmailViewModel: LoginEmailViewModel,
     ssoLoginResult: DeepLinkResult.SSOLogin?
 ) {
@@ -172,9 +167,9 @@ private fun MainLoginContent(
                 elevation = scrollState.rememberTopBarElevationState().value,
                 title = stringResource(R.string.login_title),
                 subtitleContent = {
-                    if (viewModel.serverConfig.isOnPremises) {
+                    if (loginEmailViewModel.serverConfig.isOnPremises) {
                         ServerTitle(
-                            serverLinks = viewModel.serverConfig,
+                            serverLinks = loginEmailViewModel.serverConfig,
                             style = MaterialTheme.wireTypography.body01
                         )
                     }
@@ -186,7 +181,7 @@ private fun MainLoginContent(
                     selectedTabIndex = pagerState.calculateCurrentTab(),
                     onTabChange = {
 
-                        if (viewModel.loginState.isProxyEnabled) {
+                        if (loginEmailViewModel.serverConfig.isProxyEnabled) {
                             if (pagerState.currentPage != LoginTabItem.SSO.ordinal) {
                                 ssoDisabledWithProxyDialogState.show(
                                     ssoDisabledWithProxyDialogState.savedState ?: FeatureDisabledWithProxyDialogState(
@@ -237,25 +232,25 @@ private fun MainLoginContent(
 
 @Composable
 fun LoginErrorDialog(
-    error: LoginError,
+    error: LoginState.Error,
     onDialogDismiss: () -> Unit,
     updateTheApp: () -> Unit,
     ssoLoginResult: DeepLinkResult.SSOLogin? = null
 ) {
     val dialogErrorData: LoginDialogErrorData = when (error) {
-        is LoginError.DialogError.InvalidCredentialsError -> LoginDialogErrorData(
+        is LoginState.Error.DialogError.InvalidCredentialsError -> LoginDialogErrorData(
             title = stringResource(R.string.login_error_invalid_credentials_title),
             body = AnnotatedString(stringResource(R.string.login_error_invalid_credentials_message)),
             onDismiss = onDialogDismiss
         )
 
-        is LoginError.DialogError.UserAlreadyExists -> LoginDialogErrorData(
+        is LoginState.Error.DialogError.UserAlreadyExists -> LoginDialogErrorData(
             title = stringResource(R.string.login_error_user_already_logged_in_title),
             body = AnnotatedString(stringResource(R.string.login_error_user_already_logged_in_message)),
             onDismiss = onDialogDismiss
         )
 
-        is LoginError.DialogError.ProxyError -> {
+        is LoginState.Error.DialogError.ProxyError -> {
             LoginDialogErrorData(
                 title = stringResource(R.string.error_socket_title),
                 body = AnnotatedString(stringResource(R.string.error_socket_message)),
@@ -263,7 +258,7 @@ fun LoginErrorDialog(
             )
         }
 
-        is LoginError.DialogError.GenericError -> {
+        is LoginState.Error.DialogError.GenericError -> {
             val strings = error.coreFailure.dialogErrorStrings(LocalContext.current.resources)
             LoginDialogErrorData(
                 strings.title,
@@ -272,19 +267,19 @@ fun LoginErrorDialog(
             )
         }
 
-        is LoginError.DialogError.InvalidSSOCodeError -> LoginDialogErrorData(
+        is LoginState.Error.DialogError.InvalidSSOCodeError -> LoginDialogErrorData(
             title = stringResource(R.string.login_error_invalid_credentials_title),
             body = AnnotatedString(stringResource(R.string.login_error_invalid_sso_code)),
             onDismiss = onDialogDismiss
         )
 
-        is LoginError.DialogError.InvalidSSOCookie -> LoginDialogErrorData(
+        is LoginState.Error.DialogError.InvalidSSOCookie -> LoginDialogErrorData(
             title = stringResource(R.string.login_sso_error_invalid_cookie_title),
             body = AnnotatedString(stringResource(R.string.login_sso_error_invalid_cookie_message)),
             onDismiss = onDialogDismiss
         )
 
-        is LoginError.DialogError.SSOResultError -> {
+        is LoginState.Error.DialogError.SSOResultError -> {
             with(ssoLoginResult as DeepLinkResult.SSOLogin.Failure) {
                 LoginDialogErrorData(
                     title = stringResource(R.string.sso_error_dialog_title),
@@ -294,7 +289,7 @@ fun LoginErrorDialog(
             }
         }
 
-        is LoginError.DialogError.ServerVersionNotSupported -> LoginDialogErrorData(
+        is LoginState.Error.DialogError.ServerVersionNotSupported -> LoginDialogErrorData(
             title = stringResource(R.string.api_versioning_server_version_not_supported_title),
             body = AnnotatedString(stringResource(R.string.api_versioning_server_version_not_supported_message)),
             onDismiss = onDialogDismiss,
@@ -302,7 +297,7 @@ fun LoginErrorDialog(
             dismissOnClickOutside = false
         )
 
-        is LoginError.DialogError.ClientUpdateRequired -> LoginDialogErrorData(
+        is LoginState.Error.DialogError.ClientUpdateRequired -> LoginDialogErrorData(
             title = stringResource(R.string.api_versioning_client_update_required_title),
             body = AnnotatedString(stringResource(R.string.api_versioning_client_update_required_message)),
             onDismiss = onDialogDismiss,
@@ -311,7 +306,7 @@ fun LoginErrorDialog(
             dismissOnClickOutside = false
         )
 
-        LoginError.DialogError.PasswordNeededToRegisterClient -> TODO()
+        LoginState.Error.DialogError.PasswordNeededToRegisterClient -> TODO()
 
         else -> LoginDialogErrorData(
             title = stringResource(R.string.error_unknown_title),
@@ -352,10 +347,16 @@ enum class LoginTabItem(@StringRes val titleResId: Int) : TabItem {
     override val title: UIText = UIText.StringResource(titleResId)
 }
 
-@Preview
+@PreviewMultipleThemes
 @Composable
-private fun PreviewLoginScreen() {
+private fun PreviewLoginScreen() = WireTheme {
     WireTheme {
-        MainLoginContent({}, { _, _ -> }, {}, hiltViewModel(), hiltViewModel(), ssoLoginResult = null)
+        MainLoginContent(
+            onBackPressed = {},
+            onSuccess = { _, _ -> },
+            onRemoveDeviceNeeded = {},
+            loginEmailViewModel = hiltViewModel(),
+            ssoLoginResult = null
+        )
     }
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginState.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginState.kt
index 84eb7d9f42f..a96243cccce 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginState.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginState.kt
@@ -18,39 +18,30 @@
 
 package com.wire.android.ui.authentication.login
 
-import androidx.compose.ui.text.input.TextFieldValue
-import com.wire.android.ui.common.dialogs.CustomServerDialogState
-import com.wire.kalium.logic.data.auth.login.ProxyCredentials
+import com.wire.android.util.deeplink.SSOFailureCodes
+import com.wire.kalium.logic.CoreFailure
 
-data class LoginState(
-    val userIdentifier: TextFieldValue = TextFieldValue(""),
-    val userIdentifierEnabled: Boolean = true,
-    val password: TextFieldValue = TextFieldValue(""),
-    val userInput: TextFieldValue = TextFieldValue(""),
-    val proxyIdentifier: TextFieldValue = TextFieldValue(""),
-    val proxyPassword: TextFieldValue = TextFieldValue(""),
-    val ssoLoginLoading: Boolean = false,
-    val emailLoginLoading: Boolean = false,
-    val ssoLoginEnabled: Boolean = false,
-    val emailLoginEnabled: Boolean = false,
-    val isProxyAuthRequired: Boolean = false,
-    val loginError: LoginError = LoginError.None,
-    val isProxyEnabled: Boolean = false,
-    val customServerDialogState: CustomServerDialogState? = null,
-) {
-    fun getProxyCredentials(): ProxyCredentials? =
-        if (proxyIdentifier.text.isNotBlank() && proxyPassword.text.isNotBlank()) {
-            ProxyCredentials(proxyIdentifier.text, proxyPassword.text)
-        } else {
-            null
+sealed class LoginState {
+    data object Default : LoginState()
+    data object Loading : LoginState()
+    data class Success(val initialSyncCompleted: Boolean, val isE2EIRequired: Boolean) : LoginState()
+    sealed class Error : LoginState() {
+        sealed class TextFieldError : Error() {
+            data object InvalidValue : TextFieldError()
         }
+        sealed class DialogError : Error() {
+            data class GenericError(val coreFailure: CoreFailure) : DialogError()
+            data object InvalidCredentialsError : DialogError()
+            data object ProxyError : DialogError()
+            data object InvalidSSOCookie : DialogError()
+            data object InvalidSSOCodeError : DialogError()
+            data object UserAlreadyExists : DialogError()
+            data object PasswordNeededToRegisterClient : DialogError()
+            data class SSOResultError(val result: SSOFailureCodes) :
+                DialogError()
+            data object ServerVersionNotSupported : DialogError()
+            data object ClientUpdateRequired : DialogError()
+        }
+        data object TooManyDevicesError : Error()
+    }
 }
-
-fun LoginState.updateEmailLoginEnabled() =
-    copy(
-        emailLoginEnabled = userIdentifier.text.isNotEmpty() && password.text.isNotEmpty() && !emailLoginLoading &&
-                (!isProxyAuthRequired || (isProxyAuthRequired && proxyIdentifier.text.isNotEmpty() && proxyPassword.text.isNotEmpty()))
-    )
-
-fun LoginState.updateSSOLoginEnabled() =
-    copy(ssoLoginEnabled = userInput.text.isNotEmpty() && !ssoLoginLoading)
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginViewModel.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginViewModel.kt
index 787b7e160c6..0750cd50fce 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginViewModel.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/LoginViewModel.kt
@@ -18,12 +18,9 @@
 
 package com.wire.android.ui.authentication.login
 
-import androidx.annotation.VisibleForTesting
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.compose.ui.text.input.TextFieldValue
-import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.wire.android.BuildConfig
@@ -31,8 +28,6 @@ import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.AuthServerConfigProvider
 import com.wire.android.di.ClientScopeProvider
 import com.wire.android.di.KaliumCoreLogic
-import com.wire.android.ui.navArgs
-import com.wire.android.util.EMPTY
 import com.wire.kalium.logic.CoreLogic
 import com.wire.kalium.logic.configuration.server.ServerConfig
 import com.wire.kalium.logic.data.client.ClientCapability
@@ -50,7 +45,6 @@ import javax.inject.Inject
 @HiltViewModel
 @Suppress("TooManyFunctions")
 open class LoginViewModel @Inject constructor(
-    savedStateHandle: SavedStateHandle,
     private val clientScopeProviderFactory: ClientScopeProvider.Factory,
     protected val authServerConfigProvider: AuthServerConfigProvider,
     private val userDataStoreProvider: UserDataStoreProvider,
@@ -67,62 +61,6 @@ open class LoginViewModel @Inject constructor(
         }
     }
 
-    private val loginNavArgs: LoginNavArgs = savedStateHandle.navArgs()
-    private val preFilledUserIdentifier: PreFilledUserIdentifierType = loginNavArgs.userHandle.let {
-        if (it.isNullOrEmpty()) PreFilledUserIdentifierType.None else PreFilledUserIdentifierType.PreFilled(it)
-    }
-
-    var loginState by mutableStateOf(
-        LoginState(
-            userInput = TextFieldValue(savedStateHandle[SSO_CODE_SAVED_STATE_KEY] ?: String.EMPTY),
-            userIdentifier = TextFieldValue(
-                if (preFilledUserIdentifier is PreFilledUserIdentifierType.PreFilled) preFilledUserIdentifier.userIdentifier
-                else savedStateHandle[USER_IDENTIFIER_SAVED_STATE_KEY] ?: String.EMPTY
-            ),
-            userIdentifierEnabled = preFilledUserIdentifier is PreFilledUserIdentifierType.None,
-            password = TextFieldValue(String.EMPTY),
-            isProxyAuthRequired =
-            if (serverConfig.apiProxy?.needsAuthentication != null) serverConfig.apiProxy?.needsAuthentication!!
-            else false,
-            isProxyEnabled = serverConfig.apiProxy != null
-        )
-    )
-        @VisibleForTesting
-        set
-
-    open fun updateSSOLoginError(error: LoginError) {
-        loginState = if (error is LoginError.None) {
-            loginState.copy(loginError = error)
-        } else {
-            loginState.copy(ssoLoginLoading = false, loginError = error).updateSSOLoginEnabled()
-        }
-    }
-
-    open fun updateEmailLoginError(error: LoginError) {
-        loginState = if (error is LoginError.None) {
-            loginState.copy(loginError = error)
-        } else {
-            loginState.copy(emailLoginLoading = false, loginError = error).updateEmailLoginEnabled()
-        }
-    }
-
-    fun onDialogDismiss() {
-        clearLoginErrors()
-    }
-
-    fun clearLoginErrors() {
-        clearSSOLoginError()
-        clearEmailLoginError()
-    }
-
-    fun clearSSOLoginError() {
-        updateSSOLoginError(LoginError.None)
-    }
-
-    fun clearEmailLoginError() {
-        updateEmailLoginError(LoginError.None)
-    }
-
     suspend fun registerClient(
         userId: UserId,
         password: String?,
@@ -146,34 +84,32 @@ open class LoginViewModel @Inject constructor(
     fun updateTheApp() {
         // todo : update the app after releasing on the store
     }
-
-    companion object {
-        const val SSO_CODE_SAVED_STATE_KEY = "sso_code"
-        const val USER_IDENTIFIER_SAVED_STATE_KEY = "user_identifier"
-    }
 }
 
 fun AuthenticationResult.Failure.toLoginError() = when (this) {
-    is AuthenticationResult.Failure.SocketError -> LoginError.DialogError.ProxyError
-    is AuthenticationResult.Failure.Generic -> LoginError.DialogError.GenericError(this.genericFailure)
-    is AuthenticationResult.Failure.InvalidCredentials -> LoginError.DialogError.InvalidCredentialsError
-    is AuthenticationResult.Failure.InvalidUserIdentifier -> LoginError.TextFieldError.InvalidValue
+    is AuthenticationResult.Failure.SocketError -> LoginState.Error.DialogError.ProxyError
+    is AuthenticationResult.Failure.Generic -> LoginState.Error.DialogError.GenericError(this.genericFailure)
+    is AuthenticationResult.Failure.InvalidCredentials -> LoginState.Error.DialogError.InvalidCredentialsError
+    is AuthenticationResult.Failure.InvalidUserIdentifier -> LoginState.Error.TextFieldError.InvalidValue
 }
 
 fun RegisterClientResult.Failure.toLoginError() = when (this) {
-    is RegisterClientResult.Failure.Generic -> LoginError.DialogError.GenericError(this.genericFailure)
-    is RegisterClientResult.Failure.InvalidCredentials -> LoginError.DialogError.InvalidCredentialsError
-    is RegisterClientResult.Failure.TooManyClients -> LoginError.TooManyDevicesError
-    is RegisterClientResult.Failure.PasswordAuthRequired -> LoginError.DialogError.PasswordNeededToRegisterClient
+    is RegisterClientResult.Failure.Generic -> LoginState.Error.DialogError.GenericError(this.genericFailure)
+    is RegisterClientResult.Failure.InvalidCredentials -> LoginState.Error.DialogError.InvalidCredentialsError
+    is RegisterClientResult.Failure.TooManyClients -> LoginState.Error.TooManyDevicesError
+    is RegisterClientResult.Failure.PasswordAuthRequired -> LoginState.Error.DialogError.PasswordNeededToRegisterClient
 }
 
-fun DomainLookupUseCase.Result.Failure.toLoginError() = LoginError.DialogError.GenericError(this.coreFailure)
-fun AddAuthenticatedUserUseCase.Result.Failure.toLoginError(): LoginError = when (this) {
-    is AddAuthenticatedUserUseCase.Result.Failure.Generic -> LoginError.DialogError.GenericError(this.genericFailure)
-    AddAuthenticatedUserUseCase.Result.Failure.UserAlreadyExists -> LoginError.DialogError.UserAlreadyExists
+fun DomainLookupUseCase.Result.Failure.toLoginError() = LoginState.Error.DialogError.GenericError(this.coreFailure)
+fun AddAuthenticatedUserUseCase.Result.Failure.toLoginError(): LoginState.Error = when (this) {
+    is AddAuthenticatedUserUseCase.Result.Failure.Generic -> LoginState.Error.DialogError.GenericError(this.genericFailure)
+    AddAuthenticatedUserUseCase.Result.Failure.UserAlreadyExists -> LoginState.Error.DialogError.UserAlreadyExists
 }
 
 sealed interface PreFilledUserIdentifierType {
-    object None : PreFilledUserIdentifierType
+    data object None : PreFilledUserIdentifierType
     data class PreFilled(val userIdentifier: String) : PreFilledUserIdentifierType
 }
+
+val ServerConfig.Links.isProxyEnabled get() = this.apiProxy != null
+val ServerConfig.Links.isProxyAuthRequired get() = apiProxy?.needsAuthentication ?: false
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailScreen.kt
index 3c7d851c1b9..d4fa8734ff1 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailScreen.kt
@@ -30,11 +30,13 @@ import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.rememberScrollState
 import androidx.compose.foundation.text.KeyboardOptions
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Surface
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.ui.Alignment
@@ -48,18 +50,18 @@ import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.semantics.testTagsAsResourceId
 import androidx.compose.ui.text.input.ImeAction
 import androidx.compose.ui.text.input.KeyboardType
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.text.style.TextDecoration
 import com.wire.android.R
 import com.wire.android.appLogger
-import com.wire.android.ui.authentication.login.LoginError
 import com.wire.android.ui.authentication.login.LoginErrorDialog
 import com.wire.android.ui.authentication.login.LoginState
+import com.wire.android.ui.authentication.login.isProxyAuthRequired
 import com.wire.android.ui.common.button.WireButtonState
 import com.wire.android.ui.common.button.WirePrimaryButton
 import com.wire.android.ui.common.colorsScheme
 import com.wire.android.ui.common.rememberBottomBarElevationState
+import com.wire.android.ui.common.textfield.DefaultPassword
 import com.wire.android.ui.common.textfield.WireAutoFillType
 import com.wire.android.ui.common.textfield.WirePasswordTextField
 import com.wire.android.ui.common.textfield.WireTextField
@@ -82,36 +84,47 @@ fun LoginEmailScreen(
     scrollState: ScrollState = rememberScrollState()
 ) {
     val scope = rememberCoroutineScope()
-    val loginEmailState: LoginState = loginEmailViewModel.loginState
 
     clearAutofillTree()
 
     LoginEmailContent(
         scrollState = scrollState,
-        loginState = loginEmailState,
-        onUserIdentifierChange = loginEmailViewModel::onUserIdentifierChange,
-        onPasswordChange = loginEmailViewModel::onPasswordChange,
-        onDialogDismiss = loginEmailViewModel::onDialogDismiss,
+        loginEmailState = loginEmailViewModel.loginState,
+        userIdentifierTextState = loginEmailViewModel.userIdentifierTextState,
+        proxyIdentifierState = loginEmailViewModel.proxyIdentifierTextState,
+        proxyPasswordState = loginEmailViewModel.proxyPasswordTextState,
+        passwordTextState = loginEmailViewModel.passwordTextState,
+        isProxyAuthRequired = loginEmailViewModel.serverConfig.isProxyAuthRequired,
+        apiProxyUrl = loginEmailViewModel.serverConfig.apiProxy?.host,
+        onDialogDismiss = loginEmailViewModel::clearLoginErrors,
         onRemoveDeviceOpen = {
             loginEmailViewModel.clearLoginErrors()
             onRemoveDeviceNeeded()
         },
-        onLoginButtonClick = {
-            loginEmailViewModel.login(onSuccess)
-        },
+        onLoginButtonClick = loginEmailViewModel::login,
         onUpdateApp = loginEmailViewModel::updateTheApp,
         forgotPasswordUrl = loginEmailViewModel.serverConfig.forgotPassword,
         scope = scope
     )
+
+    LaunchedEffect(loginEmailViewModel.loginState.flowState) {
+        (loginEmailViewModel.loginState.flowState as? LoginState.Success)?.let {
+            onSuccess(it.initialSyncCompleted, it.isE2EIRequired)
+        }
+    }
 }
 
 @OptIn(ExperimentalComposeUiApi::class)
 @Composable
 private fun LoginEmailContent(
     scrollState: ScrollState,
-    loginState: LoginState,
-    onUserIdentifierChange: (TextFieldValue) -> Unit,
-    onPasswordChange: (TextFieldValue) -> Unit,
+    userIdentifierTextState: TextFieldState,
+    passwordTextState: TextFieldState,
+    proxyIdentifierState: TextFieldState,
+    proxyPasswordState: TextFieldState,
+    loginEmailState: LoginEmailState,
+    isProxyAuthRequired: Boolean,
+    apiProxyUrl: String?,
     onDialogDismiss: () -> Unit,
     onRemoveDeviceOpen: () -> Unit,
     onLoginButtonClick: () -> Unit,
@@ -133,7 +146,7 @@ private fun LoginEmailContent(
                     testTagsAsResourceId = true
                 }
         ) {
-            if (loginState.isProxyAuthRequired) {
+            if (isProxyAuthRequired) {
                 Text(
                     text = stringResource(R.string.label_wire_credentials),
                     style = MaterialTheme.wireTypography.title03.copy(
@@ -150,20 +163,18 @@ private fun LoginEmailContent(
                 modifier = Modifier
                     .fillMaxWidth()
                     .padding(bottom = MaterialTheme.wireDimensions.spacing16x),
-                userIdentifier = loginState.userIdentifier,
-                onUserIdentifierChange = onUserIdentifierChange,
-                error = when (loginState.loginError) {
-                    LoginError.TextFieldError.InvalidValue -> stringResource(R.string.login_error_invalid_user_identifier)
+                userIdentifierState = userIdentifierTextState,
+                error = when (loginEmailState.flowState) {
+                    is LoginState.Error.TextFieldError.InvalidValue -> stringResource(R.string.login_error_invalid_user_identifier)
                     else -> null
                 },
-                isEnabled = loginState.userIdentifierEnabled
+                isEnabled = loginEmailState.userIdentifierEnabled,
             )
             PasswordInput(
                 modifier = Modifier
                     .fillMaxWidth()
                     .padding(bottom = MaterialTheme.wireDimensions.spacing16x),
-                password = loginState.password,
-                onPasswordChange = onPasswordChange
+                passwordState = passwordTextState,
             )
             ForgotPasswordLabel(
                 modifier = Modifier
@@ -171,8 +182,13 @@ private fun LoginEmailContent(
                     .padding(bottom = MaterialTheme.wireDimensions.spacing16x),
                 forgotPasswordUrl = forgotPasswordUrl
             )
-            if (loginState.isProxyAuthRequired) {
-                ProxyScreen()
+            if (isProxyAuthRequired) {
+                ProxyScreen(
+                    proxyIdentifierState = proxyIdentifierState,
+                    proxyPasswordState = proxyPasswordState,
+                    proxyState = loginEmailState,
+                    apiProxyUrl = apiProxyUrl,
+                )
             }
 
             Spacer(modifier = Modifier.weight(1f))
@@ -188,8 +204,8 @@ private fun LoginEmailContent(
             Box(modifier = Modifier.padding(MaterialTheme.wireDimensions.spacing16x)) {
                 LoginButton(
                     modifier = Modifier.fillMaxWidth(),
-                    loading = loginState.emailLoginLoading,
-                    enabled = loginState.emailLoginEnabled
+                    loading = loginEmailState.flowState is LoginState.Loading,
+                    enabled = loginEmailState.loginEnabled,
                 ) {
                     scope.launch {
                         onLoginButtonClick()
@@ -199,9 +215,9 @@ private fun LoginEmailContent(
         }
     }
 
-    if (loginState.loginError is LoginError.DialogError) {
-        LoginErrorDialog(loginState.loginError, onDialogDismiss, onUpdateApp)
-    } else if (loginState.loginError is LoginError.TooManyDevicesError) {
+    if (loginEmailState.flowState is LoginState.Error.DialogError) {
+        LoginErrorDialog(loginEmailState.flowState, onDialogDismiss, onUpdateApp)
+    } else if (loginEmailState.flowState is LoginState.Error.TooManyDevicesError) {
         onRemoveDeviceOpen()
     }
 }
@@ -209,15 +225,13 @@ private fun LoginEmailContent(
 @Composable
 private fun UserIdentifierInput(
     modifier: Modifier,
-    userIdentifier: TextFieldValue,
+    userIdentifierState: TextFieldState,
     error: String?,
-    onUserIdentifierChange: (TextFieldValue) -> Unit,
     isEnabled: Boolean,
 ) {
     WireTextField(
         autoFillType = WireAutoFillType.Login,
-        value = userIdentifier,
-        onValueChange = onUserIdentifierChange,
+        textState = userIdentifierState,
         placeholderText = stringResource(R.string.login_user_identifier_placeholder),
         labelText = stringResource(R.string.login_user_identifier_label),
         state = when {
@@ -232,15 +246,14 @@ private fun UserIdentifierInput(
 }
 
 @Composable
-private fun PasswordInput(modifier: Modifier, password: TextFieldValue, onPasswordChange: (TextFieldValue) -> Unit) {
+private fun PasswordInput(modifier: Modifier, passwordState: TextFieldState) {
     val keyboardController = LocalSoftwareKeyboardController.current
     WirePasswordTextField(
-        value = password,
-        onValueChange = onPasswordChange,
-        imeAction = ImeAction.Done,
-        onImeAction = { keyboardController?.hide() },
+        textState = passwordState,
+        keyboardOptions = KeyboardOptions.DefaultPassword.copy(imeAction = ImeAction.Done),
+        onKeyboardAction = { keyboardController?.hide() },
         modifier = modifier.testTag("passwordField"),
-        autofill = true,
+        autoFill = true,
         testTag = "PasswordInput"
     )
 }
@@ -293,20 +306,21 @@ private fun LoginButton(modifier: Modifier, loading: Boolean, enabled: Boolean,
 
 @PreviewMultipleThemes
 @Composable
-fun PreviewLoginEmailScreen() {
-    val scope = rememberCoroutineScope()
-    WireTheme {
-        LoginEmailContent(
-            scrollState = rememberScrollState(),
-            loginState = LoginState(),
-            onUserIdentifierChange = { },
-            onPasswordChange = { },
-            onDialogDismiss = { },
-            onRemoveDeviceOpen = { },
-            onLoginButtonClick = { },
-            onUpdateApp = {},
-            forgotPasswordUrl = "",
-            scope = scope
-        )
-    }
+fun PreviewLoginEmailScreen() = WireTheme {
+    LoginEmailContent(
+        scrollState = rememberScrollState(),
+        loginEmailState = LoginEmailState(),
+        userIdentifierTextState = TextFieldState(),
+        passwordTextState = TextFieldState(),
+        proxyIdentifierState = TextFieldState(),
+        proxyPasswordState = TextFieldState(),
+        isProxyAuthRequired = true,
+        apiProxyUrl = "",
+        onDialogDismiss = { },
+        onRemoveDeviceOpen = { },
+        onLoginButtonClick = { },
+        onUpdateApp = {},
+        forgotPasswordUrl = "",
+        scope = rememberCoroutineScope()
+    )
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailState.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailState.kt
new file mode 100644
index 00000000000..5bcede8d076
--- /dev/null
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailState.kt
@@ -0,0 +1,26 @@
+/*
+ * Wire
+ * Copyright (C) 2024 Wire Swiss GmbH
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see http://www.gnu.org/licenses/.
+ */
+package com.wire.android.ui.authentication.login.email
+
+import com.wire.android.ui.authentication.login.LoginState
+
+data class LoginEmailState(
+    val userIdentifierEnabled: Boolean = true,
+    val loginEnabled: Boolean = false,
+    val flowState: LoginState = LoginState.Default,
+)
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailVerificationCodeScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailVerificationCodeScreen.kt
index 7cd213cc179..80234d6be10 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailVerificationCodeScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailVerificationCodeScreen.kt
@@ -25,21 +25,21 @@ import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.res.stringResource
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.compose.ui.text.style.TextAlign
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.wire.android.R
+import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.authentication.verificationcode.VerificationCode
 import com.wire.android.ui.authentication.verificationcode.VerificationCodeState
 import com.wire.android.ui.common.colorsScheme
 import com.wire.android.ui.common.dimensions
 import com.wire.android.ui.common.scaffold.WireScaffold
-import com.wire.android.ui.common.textfield.CodeFieldValue
 import com.wire.android.ui.common.topappbar.WireCenterAlignedTopAppBar
 import com.wire.android.ui.common.typography
 import com.wire.android.ui.theme.WireTheme
@@ -48,21 +48,20 @@ import com.wire.android.util.ui.UIText
 
 @Composable
 fun LoginEmailVerificationCodeScreen(
-    onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit,
     viewModel: LoginEmailViewModel = hiltViewModel()
 ) = LoginEmailVerificationCodeContent(
+    viewModel.secondFactorVerificationCodeTextState,
     viewModel.secondFactorVerificationCodeState,
-    viewModel.loginState.emailLoginLoading,
-    { viewModel.onCodeChange(it, onSuccess) },
+    viewModel.loginState.flowState is LoginState.Loading,
     viewModel::onCodeResend,
     viewModel::onCodeVerificationBackPress
 )
 
 @Composable
 private fun LoginEmailVerificationCodeContent(
+    verificationCodeTextState: TextFieldState,
     verificationCodeState: VerificationCodeState,
     isLoading: Boolean,
-    onCodeChange: (CodeFieldValue) -> Unit,
     onCodeResend: () -> Unit,
     onBackPressed: () -> Unit,
 ) {
@@ -77,9 +76,9 @@ private fun LoginEmailVerificationCodeContent(
         }
     ) { internalPadding ->
         MainContent(
+            codeTextState = verificationCodeTextState,
             codeState = verificationCodeState,
             isLoading = isLoading,
-            onCodeChange = onCodeChange,
             onResendCode = onCodeResend,
             modifier = Modifier
                 .fillMaxSize()
@@ -91,9 +90,9 @@ private fun LoginEmailVerificationCodeContent(
 
 @Composable
 private fun MainContent(
+    codeTextState: TextFieldState,
     codeState: VerificationCodeState,
     isLoading: Boolean,
-    onCodeChange: (CodeFieldValue) -> Unit,
     onResendCode: () -> Unit,
     modifier: Modifier = Modifier
 ) = Column(
@@ -115,10 +114,9 @@ private fun MainContent(
         .weight(1f))
     VerificationCode(
         codeLength = codeState.codeLength,
-        currentCode = codeState.codeInput.text,
+        codeState = codeTextState,
         isLoading = isLoading,
         isCurrentCodeInvalid = codeState.isCurrentCodeInvalid,
-        onCodeChange = onCodeChange,
         onResendCode = onResendCode,
     )
     Spacer(modifier = Modifier
@@ -130,15 +128,14 @@ private fun MainContent(
 @Composable
 internal fun LoginEmailVerificationCodeScreenPreview() = WireTheme {
     LoginEmailVerificationCodeContent(
-        VerificationCodeState(
+        verificationCodeTextState = TextFieldState(),
+        verificationCodeState = VerificationCodeState(
             codeLength = 6,
-            codeInput = CodeFieldValue(TextFieldValue("12"), false),
             isCurrentCodeInvalid = false,
             emailUsed = ""
         ),
-        false,
-        {},
-        {},
-        {},
+        isLoading = false,
+        onCodeResend = {},
+        onBackPressed = {},
     )
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModel.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModel.kt
index 4fd15a6c2e3..eedd8a7b86b 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModel.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModel.kt
@@ -18,24 +18,31 @@
 
 package com.wire.android.ui.authentication.login.email
 
+import androidx.compose.foundation.text.input.TextFieldState
+import androidx.compose.foundation.text.input.clearText
+import androidx.compose.foundation.text.input.setTextAndPlaceCursorAtEnd
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.AuthServerConfigProvider
 import com.wire.android.di.ClientScopeProvider
 import com.wire.android.di.KaliumCoreLogic
-import com.wire.android.ui.authentication.login.LoginError
+import com.wire.android.ui.authentication.login.LoginNavArgs
+import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.authentication.login.LoginViewModel
+import com.wire.android.ui.authentication.login.PreFilledUserIdentifierType
+import com.wire.android.ui.authentication.login.isProxyAuthRequired
 import com.wire.android.ui.authentication.login.toLoginError
-import com.wire.android.ui.authentication.login.updateEmailLoginEnabled
 import com.wire.android.ui.authentication.verificationcode.VerificationCodeState
-import com.wire.android.ui.common.textfield.CodeFieldValue
+import com.wire.android.ui.common.textfield.textAsFlow
+import com.wire.android.ui.navArgs
+import com.wire.android.util.EMPTY
 import com.wire.android.util.dispatchers.DispatcherProvider
 import com.wire.kalium.logic.CoreLogic
+import com.wire.kalium.logic.data.auth.login.ProxyCredentials
 import com.wire.kalium.logic.data.auth.verification.VerifiableAction
 import com.wire.kalium.logic.feature.auth.AddAuthenticatedUserUseCase
 import com.wire.kalium.logic.feature.auth.AuthenticationResult
@@ -44,6 +51,10 @@ import com.wire.kalium.logic.feature.auth.autoVersioningAuth.AutoVersionAuthScop
 import com.wire.kalium.logic.feature.auth.verification.RequestSecondFactorVerificationCodeUseCase
 import com.wire.kalium.logic.feature.client.RegisterClientResult
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import javax.inject.Inject
@@ -59,28 +70,83 @@ class LoginEmailViewModel @Inject constructor(
     @KaliumCoreLogic coreLogic: CoreLogic,
     private val dispatchers: DispatcherProvider
 ) : LoginViewModel(
-    savedStateHandle,
     clientScopeProviderFactory,
     authServerConfigProvider,
     userDataStoreProvider,
     coreLogic
 ) {
+    private val loginNavArgs: LoginNavArgs = savedStateHandle.navArgs()
+    private val preFilledUserIdentifier: PreFilledUserIdentifierType = loginNavArgs.userHandle.let {
+        if (it.isNullOrEmpty()) PreFilledUserIdentifierType.None else PreFilledUserIdentifierType.PreFilled(it)
+    }
+
+    val userIdentifierTextState: TextFieldState = TextFieldState()
+    val passwordTextState: TextFieldState = TextFieldState()
+    val proxyIdentifierTextState: TextFieldState = TextFieldState()
+    val proxyPasswordTextState: TextFieldState = TextFieldState()
+    var loginState by mutableStateOf(LoginEmailState(preFilledUserIdentifier is PreFilledUserIdentifierType.None))
+
+    val secondFactorVerificationCodeTextState: TextFieldState = TextFieldState()
+    var secondFactorVerificationCodeState by mutableStateOf(VerificationCodeState())
+
+    init {
+        userIdentifierTextState.setTextAndPlaceCursorAtEnd(
+            if (preFilledUserIdentifier is PreFilledUserIdentifierType.PreFilled) {
+                preFilledUserIdentifier.userIdentifier
+            } else {
+                savedStateHandle[USER_IDENTIFIER_SAVED_STATE_KEY] ?: String.EMPTY
+            }
+        )
+        viewModelScope.launch {
+            combine(
+                userIdentifierTextState.textAsFlow().distinctUntilChanged().onEach {
+                    savedStateHandle[USER_IDENTIFIER_SAVED_STATE_KEY] = it.toString()
+                },
+                passwordTextState.textAsFlow(),
+                proxyIdentifierTextState.textAsFlow(),
+                proxyPasswordTextState.textAsFlow()
+            ) { _, _, _, _ -> }.collectLatest {
+                if (loginState.flowState != LoginState.Loading) {
+                    updateEmailFlowState(LoginState.Default)
+                }
+            }
+        }
+        viewModelScope.launch {
+            secondFactorVerificationCodeTextState.textAsFlow().collectLatest {
+                secondFactorVerificationCodeState = secondFactorVerificationCodeState.copy(isCurrentCodeInvalid = false)
+                if (it.length == VerificationCodeState.DEFAULT_VERIFICATION_CODE_LENGTH) {
+                    login()
+                }
+            }
+        }
+    }
 
-    var secondFactorVerificationCodeState by mutableStateOf(
-        VerificationCodeState()
-    )
+    private fun updateEmailFlowState(flowState: LoginState) {
+        val proxyFieldsNotEmpty = proxyIdentifierTextState.text.isNotEmpty() && proxyPasswordTextState.text.isNotEmpty()
+        loginState = loginState.copy(
+            flowState = flowState,
+            loginEnabled = userIdentifierTextState.text.isNotEmpty()
+                    && passwordTextState.text.isNotEmpty()
+                    && (!serverConfig.isProxyAuthRequired || proxyFieldsNotEmpty)
+                    && flowState !is LoginState.Loading
+        )
+    }
+
+    fun clearLoginErrors() {
+        updateEmailFlowState(LoginState.Default)
+    }
 
     @Suppress("LongMethod")
-    fun login(onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit) {
-        loginState = loginState.copy(emailLoginLoading = true, loginError = LoginError.None).updateEmailLoginEnabled()
+    fun login() {
+        updateEmailFlowState(LoginState.Loading)
         viewModelScope.launch {
             val authScope = withContext(dispatchers.io()) { resolveCurrentAuthScope() } ?: return@launch
 
-            val secondFactorVerificationCode = secondFactorVerificationCodeState.codeInput.text.text
+            val secondFactorVerificationCode = secondFactorVerificationCodeTextState.text.toString()
             val loginResult = withContext(dispatchers.io()) {
                 authScope.login(
-                    userIdentifier = loginState.userIdentifier.text,
-                    password = loginState.password.text,
+                    userIdentifier = userIdentifierTextState.text.toString(),
+                    password = passwordTextState.text.toString(),
                     shouldPersistClient = true,
                     secondFactorVerificationCode = secondFactorVerificationCode
                 )
@@ -102,7 +168,7 @@ class LoginEmailViewModel @Inject constructor(
             }.let {
                 when (it) {
                     is AddAuthenticatedUserUseCase.Result.Failure -> {
-                        updateEmailLoginError(it.toLoginError())
+                        updateEmailFlowState(it.toLoginError())
                         return@launch
                     }
 
@@ -112,21 +178,21 @@ class LoginEmailViewModel @Inject constructor(
             withContext(dispatchers.io()) {
                 registerClient(
                     userId = storedUserId,
-                    password = loginState.password.text,
+                    password = passwordTextState.text.toString(),
                 )
             }.let {
                 when (it) {
                     is RegisterClientResult.Failure -> {
-                        updateEmailLoginError(it.toLoginError())
+                        updateEmailFlowState(it.toLoginError())
                         return@launch
                     }
 
                     is RegisterClientResult.Success -> {
-                        onSuccess(isInitialSyncCompleted(storedUserId), false)
+                        updateEmailFlowState(LoginState.Success(isInitialSyncCompleted(storedUserId), false))
                     }
 
                     is RegisterClientResult.E2EICertificateRequired -> {
-                        onSuccess(isInitialSyncCompleted(storedUserId), true)
+                        updateEmailFlowState(LoginState.Success(isInitialSyncCompleted(storedUserId), true))
                         return@launch
                     }
                 }
@@ -134,25 +200,30 @@ class LoginEmailViewModel @Inject constructor(
         }
     }
 
+    private fun getProxyCredentials(): ProxyCredentials? =
+        if (proxyIdentifierTextState.text.isNotBlank() && proxyPasswordTextState.text.isNotBlank()) {
+            ProxyCredentials(proxyIdentifierTextState.text.toString(), proxyPasswordTextState.text.toString())
+        } else {
+            null
+        }
+
     private suspend fun resolveCurrentAuthScope(): AuthenticationScope? =
-        coreLogic.versionedAuthenticationScope(serverConfig).invoke(
-            loginState.getProxyCredentials()
-        ).let {
+        coreLogic.versionedAuthenticationScope(serverConfig).invoke(getProxyCredentials()).let {
             when (it) {
                 is AutoVersionAuthScopeUseCase.Result.Success -> it.authenticationScope
 
                 is AutoVersionAuthScopeUseCase.Result.Failure.UnknownServerVersion -> {
-                    updateEmailLoginError(LoginError.DialogError.ServerVersionNotSupported)
+                    updateEmailFlowState(LoginState.Error.DialogError.ServerVersionNotSupported)
                     return null
                 }
 
                 is AutoVersionAuthScopeUseCase.Result.Failure.TooNewVersion -> {
-                    updateEmailLoginError(LoginError.DialogError.ClientUpdateRequired)
+                    updateEmailFlowState(LoginState.Error.DialogError.ClientUpdateRequired)
                     return null
                 }
 
                 is AutoVersionAuthScopeUseCase.Result.Failure.Generic -> {
-                    updateEmailLoginError(LoginError.DialogError.GenericError(it.genericFailure))
+                    updateEmailFlowState(LoginState.Error.DialogError.GenericError(it.genericFailure))
                     return null
                 }
             }
@@ -161,21 +232,21 @@ class LoginEmailViewModel @Inject constructor(
     private suspend fun handleAuthenticationFailure(it: AuthenticationResult.Failure, authScope: AuthenticationScope) {
         when (it) {
             is AuthenticationResult.Failure.InvalidCredentials.Missing2FA -> {
-                loginState = loginState.copy(emailLoginLoading = false).updateEmailLoginEnabled()
+                updateEmailFlowState(LoginState.Default)
                 request2FACode(authScope)
             }
 
             is AuthenticationResult.Failure.InvalidCredentials.Invalid2FA -> {
-                loginState = loginState.copy(emailLoginLoading = false).updateEmailLoginEnabled()
+                updateEmailFlowState(LoginState.Default)
                 secondFactorVerificationCodeState = secondFactorVerificationCodeState.copy(isCurrentCodeInvalid = true)
             }
 
-            else -> updateEmailLoginError(it.toLoginError())
+            else -> updateEmailFlowState(it.toLoginError())
         }
     }
 
     private suspend fun request2FACode(authScope: AuthenticationScope) {
-        val email = loginState.userIdentifier.text.trim()
+        val email = userIdentifierTextState.text.trim().toString()
         val result = authScope.requestSecondFactorVerificationCode(
             email = email,
             verifiableAction = VerifiableAction.LOGIN_OR_CLIENT_REGISTRATION
@@ -187,46 +258,18 @@ class LoginEmailViewModel @Inject constructor(
                     isCodeInputNecessary = true,
                     emailUsed = email,
                 )
-                updateEmailLoginError(LoginError.None)
+                updateEmailFlowState(LoginState.Default)
             }
 
             is RequestSecondFactorVerificationCodeUseCase.Result.Failure.Generic -> {
-                updateEmailLoginError(LoginError.DialogError.GenericError(result.cause))
+                updateEmailFlowState(LoginState.Error.DialogError.GenericError(result.cause))
             }
         }
     }
 
-    fun onUserIdentifierChange(newText: TextFieldValue) {
-        // in case an error is showing e.g. inline error it should be cleared
-        if (loginState.loginError is LoginError.TextFieldError && newText != loginState.userIdentifier) {
-            clearEmailLoginError()
-        }
-        loginState = loginState.copy(userIdentifier = newText).updateEmailLoginEnabled()
-        savedStateHandle.set(USER_IDENTIFIER_SAVED_STATE_KEY, newText.text)
-    }
-
-    fun onPasswordChange(newText: TextFieldValue) {
-        loginState = loginState.copy(password = newText).updateEmailLoginEnabled()
-    }
-
-    fun onProxyIdentifierChange(newText: TextFieldValue) {
-        loginState = loginState.copy(proxyIdentifier = newText).updateEmailLoginEnabled()
-    }
-
-    fun onProxyPasswordChange(newText: TextFieldValue) {
-        loginState = loginState.copy(proxyPassword = newText).updateEmailLoginEnabled()
-    }
-
-    fun onCodeChange(newValue: CodeFieldValue, onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit) {
-        secondFactorVerificationCodeState = secondFactorVerificationCodeState.copy(codeInput = newValue, isCurrentCodeInvalid = false)
-        if (newValue.isFullyFilled) {
-            login(onSuccess)
-        }
-    }
-
     fun onCodeVerificationBackPress() {
+        secondFactorVerificationCodeTextState.clearText()
         secondFactorVerificationCodeState = secondFactorVerificationCodeState.copy(
-            codeInput = CodeFieldValue(TextFieldValue(""), false),
             isCodeInputNecessary = false,
             emailUsed = "",
         )
@@ -239,4 +282,8 @@ class LoginEmailViewModel @Inject constructor(
             }
         }
     }
+
+    companion object {
+        const val USER_IDENTIFIER_SAVED_STATE_KEY = "user_identifier"
+    }
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/ProxyScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/ProxyScreen.kt
index d6ff525d686..6fba32a7d1c 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/ProxyScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/email/ProxyScreen.kt
@@ -23,6 +23,7 @@ import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.text.KeyboardOptions
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.material3.HorizontalDivider
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Text
@@ -33,14 +34,11 @@ import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.input.ImeAction
 import androidx.compose.ui.text.input.KeyboardType
-import androidx.compose.ui.text.input.TextFieldValue
-import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.unit.Dp
-import androidx.hilt.navigation.compose.hiltViewModel
 import com.wire.android.R
-import com.wire.android.ui.authentication.login.LoginError
 import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.common.colorsScheme
+import com.wire.android.ui.common.textfield.DefaultPassword
 import com.wire.android.ui.common.textfield.WirePasswordTextField
 import com.wire.android.ui.common.textfield.WireTextField
 import com.wire.android.ui.common.textfield.WireTextFieldState
@@ -48,28 +46,18 @@ import com.wire.android.ui.theme.WireTheme
 import com.wire.android.ui.theme.wireColorScheme
 import com.wire.android.ui.theme.wireDimensions
 import com.wire.android.ui.theme.wireTypography
+import com.wire.android.util.ui.PreviewMultipleThemes
 
 @Composable
-fun ProxyScreen() {
-    val loginEmailViewModel: LoginEmailViewModel = hiltViewModel()
-    val proxyState: LoginState = loginEmailViewModel.loginState
-    ProxyContent(
-        proxyState = proxyState,
-        apiProxyUrl = loginEmailViewModel.serverConfig.apiProxy?.host,
-        onProxyIdentifierChange = { loginEmailViewModel.onProxyIdentifierChange(it) },
-        onProxyPasswordChange = { loginEmailViewModel.onProxyPasswordChange(it) },
-    )
-}
-
-@Composable
-private fun ProxyContent(
-    proxyState: LoginState,
+fun ProxyScreen(
+    proxyIdentifierState: TextFieldState,
+    proxyPasswordState: TextFieldState,
+    proxyState: LoginEmailState,
     apiProxyUrl: String?,
-    onProxyIdentifierChange: (TextFieldValue) -> Unit,
-    onProxyPasswordChange: (TextFieldValue) -> Unit,
+    modifier: Modifier = Modifier,
 ) {
     Column(
-        modifier = Modifier
+        modifier = modifier
     ) {
         HorizontalDivider(thickness = Dp.Hairline, color = MaterialTheme.wireColorScheme.divider)
         Text(
@@ -100,10 +88,9 @@ private fun ProxyContent(
             modifier = Modifier
                 .fillMaxWidth()
                 .padding(bottom = MaterialTheme.wireDimensions.spacing16x),
-            proxyIdentifier = proxyState.proxyIdentifier,
-            onProxyUserIdentifierChange = onProxyIdentifierChange,
-            error = when (proxyState.loginError) {
-                LoginError.TextFieldError.InvalidValue -> stringResource(R.string.login_error_invalid_user_identifier)
+            proxyIdentifierState = proxyIdentifierState,
+            error = when (proxyState.flowState) {
+                LoginState.Error.TextFieldError.InvalidValue -> stringResource(R.string.login_error_invalid_user_identifier)
                 else -> null
             },
         )
@@ -111,8 +98,7 @@ private fun ProxyContent(
             modifier = Modifier
                 .fillMaxWidth()
                 .padding(bottom = MaterialTheme.wireDimensions.spacing16x),
-            proxyPassword = proxyState.proxyPassword,
-            onProxyPasswordChange = onProxyPasswordChange
+            proxyPasswordState = proxyPasswordState,
         )
 
         Spacer(modifier = Modifier.weight(1f))
@@ -122,13 +108,11 @@ private fun ProxyContent(
 @Composable
 private fun ProxyIdentifierInput(
     modifier: Modifier,
-    proxyIdentifier: TextFieldValue,
+    proxyIdentifierState: TextFieldState,
     error: String?,
-    onProxyUserIdentifierChange: (TextFieldValue) -> Unit,
 ) {
     WireTextField(
-        value = proxyIdentifier,
-        onValueChange = onProxyUserIdentifierChange,
+        textState = proxyIdentifierState,
         placeholderText = stringResource(R.string.login_user_identifier_placeholder),
         labelText = stringResource(R.string.login_proxy_identifier_label),
         state = if (error != null) WireTextFieldState.Error(error) else WireTextFieldState.Default,
@@ -138,28 +122,28 @@ private fun ProxyIdentifierInput(
 }
 
 @Composable
-private fun ProxyPasswordInput(modifier: Modifier, proxyPassword: TextFieldValue, onProxyPasswordChange: (TextFieldValue) -> Unit) {
+private fun ProxyPasswordInput(
+    modifier: Modifier,
+    proxyPasswordState: TextFieldState
+) {
     val keyboardController = LocalSoftwareKeyboardController.current
     WirePasswordTextField(
-        value = proxyPassword,
-        onValueChange = onProxyPasswordChange,
-        imeAction = ImeAction.Done,
+        textState = proxyPasswordState,
         labelText = stringResource(R.string.label_proxy_password),
-        onImeAction = { keyboardController?.hide() },
+        keyboardOptions = KeyboardOptions.DefaultPassword.copy(imeAction = ImeAction.Done),
+        onKeyboardAction = { keyboardController?.hide() },
         modifier = modifier.testTag("passwordField"),
-        autofill = false
+        autoFill = false
     )
 }
 
-@Preview
+@PreviewMultipleThemes
 @Composable
-fun PreviewProxyScreen() {
-    WireTheme {
-        ProxyContent(
-            proxyState = LoginState(),
-            apiProxyUrl = "",
-            onProxyIdentifierChange = { },
-            onProxyPasswordChange = { },
-        )
-    }
+fun PreviewProxyContent() = WireTheme {
+    ProxyScreen(
+        proxyState = LoginEmailState(),
+        apiProxyUrl = "",
+        proxyIdentifierState = TextFieldState(),
+        proxyPasswordState = TextFieldState(),
+    )
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOScreen.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOScreen.kt
index 93e6b5bd924..b91861b35fa 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOScreen.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOScreen.kt
@@ -28,6 +28,7 @@ import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.rememberScrollState
 import androidx.compose.foundation.text.KeyboardOptions
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.runtime.Composable
@@ -40,11 +41,8 @@ import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.input.ImeAction
 import androidx.compose.ui.text.input.KeyboardType
-import androidx.compose.ui.text.input.TextFieldValue
-import androidx.compose.ui.tooling.preview.Preview
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.wire.android.R
-import com.wire.android.ui.authentication.login.LoginError
 import com.wire.android.ui.authentication.login.LoginErrorDialog
 import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.common.button.WireButtonState
@@ -56,6 +54,7 @@ import com.wire.android.ui.theme.WireTheme
 import com.wire.android.ui.theme.wireDimensions
 import com.wire.android.util.CustomTabsHelper
 import com.wire.android.util.deeplink.DeepLinkResult
+import com.wire.android.util.ui.PreviewMultipleThemes
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 
@@ -64,20 +63,20 @@ fun LoginSSOScreen(
     onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit,
     onRemoveDeviceNeeded: () -> Unit,
     ssoLoginResult: DeepLinkResult.SSOLogin?,
+    loginSSOViewModel: LoginSSOViewModel = hiltViewModel(),
     scrollState: ScrollState = rememberScrollState()
 ) {
     val scope = rememberCoroutineScope()
     val context = LocalContext.current
-    val loginSSOViewModel: LoginSSOViewModel = hiltViewModel()
 
     LaunchedEffect(ssoLoginResult) {
-        loginSSOViewModel.handleSSOResult(ssoLoginResult, onSuccess)
+        loginSSOViewModel.handleSSOResult(ssoLoginResult)
     }
     LoginSSOContent(
         scrollState = scrollState,
-        loginState = loginSSOViewModel.loginState,
-        onCodeChange = loginSSOViewModel::onSSOCodeChange,
-        onErrorDialogDismiss = loginSSOViewModel::onDialogDismiss,
+        ssoCodeTextState = loginSSOViewModel.ssoTextState,
+        loginSSOState = loginSSOViewModel.loginState,
+        onErrorDialogDismiss = loginSSOViewModel::clearLoginErrors,
         onRemoveDeviceOpen = {
             loginSSOViewModel.clearLoginErrors()
             onRemoveDeviceNeeded()
@@ -92,13 +91,18 @@ fun LoginSSOScreen(
     LaunchedEffect(loginSSOViewModel) {
         loginSSOViewModel.openWebUrl.onEach { CustomTabsHelper.launchUrl(context, it) }.launchIn(scope)
     }
+    LaunchedEffect(loginSSOViewModel.loginState.flowState) {
+        (loginSSOViewModel.loginState.flowState as? LoginState.Success)?.let {
+            onSuccess(it.initialSyncCompleted, it.isE2EIRequired)
+        }
+    }
 }
 
 @Composable
 private fun LoginSSOContent(
     scrollState: ScrollState,
-    loginState: LoginState,
-    onCodeChange: (TextFieldValue) -> Unit,
+    loginSSOState: LoginSSOState,
+    ssoCodeTextState: TextFieldState,
     onErrorDialogDismiss: () -> Unit,
     onRemoveDeviceOpen: () -> Unit,
     onLoginButtonClick: () -> Unit,
@@ -117,30 +121,29 @@ private fun LoginSSOContent(
             modifier = Modifier
                 .fillMaxWidth()
                 .padding(bottom = MaterialTheme.wireDimensions.spacing16x),
-            ssoCode = loginState.userInput,
-            onCodeChange = onCodeChange,
-            error = when (loginState.loginError) {
-                LoginError.TextFieldError.InvalidValue -> stringResource(R.string.login_error_invalid_sso_code_format)
+            ssoCodeState = ssoCodeTextState,
+            error = when (loginSSOState.flowState) {
+                is LoginState.Error.TextFieldError.InvalidValue -> stringResource(R.string.login_error_invalid_sso_code_format)
                 else -> null
             }
         )
         Spacer(modifier = Modifier.weight(1f))
         LoginButton(
             modifier = Modifier.fillMaxWidth(),
-            loading = loginState.ssoLoginLoading,
-            enabled = loginState.ssoLoginEnabled,
+            loading = loginSSOState.flowState is LoginState.Loading,
+            enabled = loginSSOState.loginEnabled,
             onClick = onLoginButtonClick
         )
     }
-    if (loginState.loginError is LoginError.DialogError) {
-        LoginErrorDialog(loginState.loginError, onErrorDialogDismiss, {}, ssoLoginResult)
-    } else if (loginState.loginError is LoginError.TooManyDevicesError) {
+    if (loginSSOState.flowState is LoginState.Error.DialogError) {
+        LoginErrorDialog(loginSSOState.flowState, onErrorDialogDismiss, {}, ssoLoginResult)
+    } else if (loginSSOState.flowState is LoginState.Error.TooManyDevicesError) {
         onRemoveDeviceOpen()
     }
 
-    if (loginState.customServerDialogState != null) {
+    if (loginSSOState.customServerDialogState != null) {
         CustomServerDialog(
-            serverLinks = loginState.customServerDialogState.serverLinks,
+            serverLinks = loginSSOState.customServerDialogState.serverLinks,
             onDismiss = onCustomServerDialogDismiss,
             onConfirm = onCustomServerDialogConfirm
         )
@@ -150,13 +153,11 @@ private fun LoginSSOContent(
 @Composable
 private fun SSOCodeInput(
     modifier: Modifier,
-    ssoCode: TextFieldValue,
+    ssoCodeState: TextFieldState,
     error: String?,
-    onCodeChange: (TextFieldValue) -> Unit
 ) {
     WireTextField(
-        value = ssoCode,
-        onValueChange = onCodeChange,
+        textState = ssoCodeState,
         labelText = stringResource(R.string.login_sso_code_label),
         state = if (error != null) WireTextFieldState.Error(error) else WireTextFieldState.Default,
         keyboardOptions = KeyboardOptions(keyboardType = KeyboardType.Text, imeAction = ImeAction.Next),
@@ -182,10 +183,18 @@ private fun LoginButton(modifier: Modifier, loading: Boolean, enabled: Boolean,
     }
 }
 
-@Preview
+@PreviewMultipleThemes
 @Composable
-fun PreviewLoginSSOScreen() {
-    WireTheme {
-        LoginSSOContent(rememberScrollState(), LoginState(), {}, {}, {}, {}, {}, {}, null)
-    }
+fun PreviewLoginSSOScreen() = WireTheme {
+    LoginSSOContent(
+        scrollState = rememberScrollState(),
+        loginSSOState = LoginSSOState(),
+        ssoCodeTextState = TextFieldState(),
+        onErrorDialogDismiss = { },
+        onRemoveDeviceOpen = { },
+        onLoginButtonClick = { },
+        onCustomServerDialogDismiss = { },
+        onCustomServerDialogConfirm = { },
+        ssoLoginResult = null
+    )
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOState.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOState.kt
new file mode 100644
index 00000000000..9b0de0a845d
--- /dev/null
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOState.kt
@@ -0,0 +1,27 @@
+/*
+ * Wire
+ * Copyright (C) 2024 Wire Swiss GmbH
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see http://www.gnu.org/licenses/.
+ */
+package com.wire.android.ui.authentication.login.sso
+
+import com.wire.android.ui.authentication.login.LoginState
+import com.wire.android.ui.common.dialogs.CustomServerDialogState
+
+data class LoginSSOState(
+    val loginEnabled: Boolean = false,
+    val flowState: LoginState = LoginState.Default,
+    val customServerDialogState: CustomServerDialogState? = null,
+)
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModel.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModel.kt
index 0b8bc779420..e8b1c513476 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModel.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModel.kt
@@ -19,18 +19,23 @@
 package com.wire.android.ui.authentication.login.sso
 
 import androidx.annotation.VisibleForTesting
-import androidx.compose.ui.text.input.TextFieldValue
+import androidx.compose.foundation.text.input.TextFieldState
+import androidx.compose.foundation.text.input.setTextAndPlaceCursorAtEnd
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
 import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.viewModelScope
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.AuthServerConfigProvider
 import com.wire.android.di.ClientScopeProvider
 import com.wire.android.di.KaliumCoreLogic
-import com.wire.android.ui.authentication.login.LoginError
+import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.authentication.login.LoginViewModel
 import com.wire.android.ui.authentication.login.toLoginError
-import com.wire.android.ui.authentication.login.updateSSOLoginEnabled
 import com.wire.android.ui.common.dialogs.CustomServerDialogState
+import com.wire.android.ui.common.textfield.textAsFlow
+import com.wire.android.util.EMPTY
 import com.wire.android.util.deeplink.DeepLinkResult
 import com.wire.kalium.logic.CoreFailure
 import com.wire.kalium.logic.CoreLogic
@@ -46,6 +51,8 @@ import com.wire.kalium.logic.feature.auth.sso.SSOLoginSessionResult
 import com.wire.kalium.logic.feature.client.RegisterClientResult
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -60,7 +67,6 @@ class LoginSSOViewModel @Inject constructor(
     authServerConfigProvider: AuthServerConfigProvider,
     userDataStoreProvider: UserDataStoreProvider
 ) : LoginViewModel(
-    savedStateHandle,
     clientScopeProviderFactory,
     authServerConfigProvider,
     userDataStoreProvider,
@@ -69,10 +75,35 @@ class LoginSSOViewModel @Inject constructor(
 
     var openWebUrl = MutableSharedFlow<String>()
 
-    fun login() {
-        loginState = loginState.copy(ssoLoginLoading = true, loginError = LoginError.None).updateSSOLoginEnabled()
+    val ssoTextState: TextFieldState = TextFieldState()
+    var loginState: LoginSSOState by mutableStateOf(LoginSSOState())
+
+    init {
+        ssoTextState.setTextAndPlaceCursorAtEnd(savedStateHandle[SSO_CODE_SAVED_STATE_KEY] ?: String.EMPTY)
+        viewModelScope.launch {
+            ssoTextState.textAsFlow().distinctUntilChanged().collectLatest {
+                if (loginState.flowState != LoginState.Loading) {
+                    updateSSOFlowState(LoginState.Default)
+                }
+                savedStateHandle[SSO_CODE_SAVED_STATE_KEY] = it.toString()
+            }
+        }
+    }
+
+    private fun updateSSOFlowState(flowState: LoginState) {
+        loginState = loginState.copy(
+            flowState = flowState,
+            loginEnabled = ssoTextState.text.isNotEmpty() && flowState !is LoginState.Loading
+        )
+    }
 
-        loginState.userInput.text.also {
+    fun clearLoginErrors() {
+        updateSSOFlowState(LoginState.Default)
+    }
+
+    fun login() {
+        updateSSOFlowState(LoginState.Loading)
+        ssoTextState.text.toString().also {
             if (validateEmailUseCase(it)) {
                 domainLookupFlow()
             } else {
@@ -82,22 +113,22 @@ class LoginSSOViewModel @Inject constructor(
     }
 
     fun onCustomServerDialogDismiss() {
-        loginState = loginState.copy(customServerDialogState = null)
+        updateSSOFlowState(LoginState.Default)
     }
 
     fun onCustomServerDialogConfirm() {
         viewModelScope.launch {
-            if (loginState.customServerDialogState != null) {
-                authServerConfigProvider.updateAuthServer(loginState.customServerDialogState!!.serverLinks)
+            loginState.customServerDialogState?.let {
+                authServerConfigProvider.updateAuthServer(it.serverLinks)
 
                 // sso does not support proxy
                 // TODO: add proxy support
-                val authScope = coreLogic.versionedAuthenticationScope(loginState.customServerDialogState!!.serverLinks)(null).let {
+                val authScope = coreLogic.versionedAuthenticationScope(it.serverLinks)(null).let {
                     when (it) {
                         is AutoVersionAuthScopeUseCase.Result.Failure.Generic,
                         AutoVersionAuthScopeUseCase.Result.Failure.TooNewVersion,
                         AutoVersionAuthScopeUseCase.Result.Failure.UnknownServerVersion -> {
-                            loginState = loginState.copy(customServerDialogState = null)
+                            updateSSOFlowState(LoginState.Default)
                             return@launch
                         }
 
@@ -117,7 +148,8 @@ class LoginSSOViewModel @Inject constructor(
                                     SSOInitiateLoginUseCase.Param.WithRedirect(ssoCodeWithPrefix)
                                 ).let { result ->
                                     when (result) {
-                                        is SSOInitiateLoginResult.Failure -> updateSSOLoginError(result.toLoginSSOError())
+                                        is SSOInitiateLoginResult.Failure ->
+                                            updateSSOFlowState(result.toLoginSSOError())
                                         is SSOInitiateLoginResult.Success -> openWebUrl(result.requestUrl)
                                     }
                                 }
@@ -125,7 +157,7 @@ class LoginSSOViewModel @Inject constructor(
                         }
                     }
                 }
-                loginState = loginState.copy(customServerDialogState = null)
+                updateSSOFlowState(LoginState.Default)
             }
         }
     }
@@ -143,7 +175,7 @@ class LoginSSOViewModel @Inject constructor(
                         is AutoVersionAuthScopeUseCase.Result.Failure.Generic,
                         AutoVersionAuthScopeUseCase.Result.Failure.TooNewVersion,
                         AutoVersionAuthScopeUseCase.Result.Failure.UnknownServerVersion -> {
-                            loginState = loginState.copy(loginError = LoginError.DialogError.ServerVersionNotSupported)
+                            updateSSOFlowState(LoginState.Error.DialogError.ServerVersionNotSupported)
                             return@launch
                         }
 
@@ -151,18 +183,15 @@ class LoginSSOViewModel @Inject constructor(
                     }
                 }
 
-            defaultAuthScope.domainLookup(loginState.userInput.text).also {
+            defaultAuthScope.domainLookup(ssoTextState.text.toString()).also {
                 when (it) {
                     is DomainLookupUseCase.Result.Failure -> {
-                        loginState = loginState.copy(ssoLoginLoading = false, loginError = it.toLoginError())
+                        updateSSOFlowState(it.toLoginError())
                     }
 
                     is DomainLookupUseCase.Result.Success -> {
-                        loginState = loginState.copy(
-                            ssoLoginLoading = false,
-                            loginError = LoginError.None,
-                            customServerDialogState = CustomServerDialogState(it.serverLinks)
-                        )
+                        loginState = loginState.copy(customServerDialogState = CustomServerDialogState(it.serverLinks))
+                        updateSSOFlowState(LoginState.Default)
                     }
                 }
             }
@@ -178,12 +207,12 @@ class LoginSSOViewModel @Inject constructor(
                         is AutoVersionAuthScopeUseCase.Result.Success -> it.authenticationScope
 
                         is AutoVersionAuthScopeUseCase.Result.Failure.UnknownServerVersion -> {
-                            loginState = loginState.copy(loginError = LoginError.DialogError.ServerVersionNotSupported)
+                            updateSSOFlowState(LoginState.Error.DialogError.ServerVersionNotSupported)
                             return@launch
                         }
 
                         is AutoVersionAuthScopeUseCase.Result.Failure.TooNewVersion -> {
-                            loginState = loginState.copy(loginError = LoginError.DialogError.ClientUpdateRequired)
+                            updateSSOFlowState(LoginState.Error.DialogError.ClientUpdateRequired)
                             return@launch
                         }
 
@@ -193,9 +222,9 @@ class LoginSSOViewModel @Inject constructor(
                     }
                 }
 
-            authScope.ssoLoginScope.initiate(SSOInitiateLoginUseCase.Param.WithRedirect(loginState.userInput.text)).let { result ->
+            authScope.ssoLoginScope.initiate(SSOInitiateLoginUseCase.Param.WithRedirect(ssoTextState.text.toString())).let { result ->
                 when (result) {
-                    is SSOInitiateLoginResult.Failure -> updateSSOLoginError(result.toLoginSSOError())
+                    is SSOInitiateLoginResult.Failure -> updateSSOFlowState(result.toLoginSSOError())
                     is SSOInitiateLoginResult.Success -> openWebUrl(result.requestUrl)
                 }
             }
@@ -206,10 +235,9 @@ class LoginSSOViewModel @Inject constructor(
     @VisibleForTesting
     fun establishSSOSession(
         cookie: String,
-        serverConfigId: String,
-        onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit
+        serverConfigId: String
     ) {
-        loginState = loginState.copy(ssoLoginLoading = true, loginError = LoginError.None).updateSSOLoginEnabled()
+        updateSSOFlowState(LoginState.Loading)
         viewModelScope.launch {
             val authScope =
                 coreLogic.versionedAuthenticationScope(serverConfig)(null).let {
@@ -217,12 +245,12 @@ class LoginSSOViewModel @Inject constructor(
                         is AutoVersionAuthScopeUseCase.Result.Success -> it.authenticationScope
 
                         is AutoVersionAuthScopeUseCase.Result.Failure.UnknownServerVersion -> {
-                            loginState = loginState.copy(loginError = LoginError.DialogError.ServerVersionNotSupported)
+                            updateSSOFlowState(LoginState.Error.DialogError.ServerVersionNotSupported)
                             return@launch
                         }
 
                         is AutoVersionAuthScopeUseCase.Result.Failure.TooNewVersion -> {
-                            loginState = loginState.copy(loginError = LoginError.DialogError.ClientUpdateRequired)
+                            updateSSOFlowState(LoginState.Error.DialogError.ClientUpdateRequired)
                             return@launch
                         }
 
@@ -234,7 +262,7 @@ class LoginSSOViewModel @Inject constructor(
             val ssoLoginResult = authScope.ssoLoginScope.getLoginSession(cookie).let {
                 when (it) {
                     is SSOLoginSessionResult.Failure -> {
-                        updateSSOLoginError(it.toLoginError())
+                        updateSSOFlowState(it.toLoginError())
                         return@launch
                     }
 
@@ -250,7 +278,7 @@ class LoginSSOViewModel @Inject constructor(
             ).let {
                 when (it) {
                     is AddAuthenticatedUserUseCase.Result.Failure -> {
-                        updateSSOLoginError(it.toLoginError())
+                        updateSSOFlowState(it.toLoginError())
                         return@launch
                     }
 
@@ -260,61 +288,55 @@ class LoginSSOViewModel @Inject constructor(
             registerClient(storedUserId, null).let {
                 when (it) {
                     is RegisterClientResult.Success -> {
-                        onSuccess(isInitialSyncCompleted(storedUserId), false)
+                        updateSSOFlowState(LoginState.Success(isInitialSyncCompleted(storedUserId), false))
                     }
 
                     is RegisterClientResult.Failure -> {
-                        updateSSOLoginError(it.toLoginError())
+                        updateSSOFlowState(it.toLoginError())
                         return@launch
                     }
 
                     is RegisterClientResult.E2EICertificateRequired -> {
-                        onSuccess(isInitialSyncCompleted(storedUserId), true)
+                        updateSSOFlowState(LoginState.Success(isInitialSyncCompleted(storedUserId), true))
                     }
                 }
             }
         }
     }
 
-    fun onSSOCodeChange(newText: TextFieldValue) {
-        // in case an error is showing e.g. inline error is should be cleared
-        if (loginState.loginError is LoginError.TextFieldError && newText != loginState.userInput) {
-            clearSSOLoginError()
-        }
-        loginState = loginState.copy(userInput = newText).updateSSOLoginEnabled()
-        savedStateHandle.set(SSO_CODE_SAVED_STATE_KEY, newText.text)
-    }
-
-    fun handleSSOResult(
-        ssoLoginResult: DeepLinkResult.SSOLogin?,
-        onSuccess: (initialSyncCompleted: Boolean, isE2EIRequired: Boolean) -> Unit
-    ) =
+    fun handleSSOResult(ssoLoginResult: DeepLinkResult.SSOLogin?) =
         when (ssoLoginResult) {
             is DeepLinkResult.SSOLogin.Success -> {
-                establishSSOSession(ssoLoginResult.cookie, ssoLoginResult.serverConfigId, onSuccess)
+                establishSSOSession(ssoLoginResult.cookie, ssoLoginResult.serverConfigId)
             }
 
-            is DeepLinkResult.SSOLogin.Failure -> updateSSOLoginError(LoginError.DialogError.SSOResultError(ssoLoginResult.ssoError))
+            is DeepLinkResult.SSOLogin.Failure ->
+                updateSSOFlowState(LoginState.Error.DialogError.SSOResultError(ssoLoginResult.ssoError))
+
             null -> {}
         }
 
     private fun openWebUrl(url: String) {
         viewModelScope.launch {
-            loginState = loginState.copy(ssoLoginLoading = false, loginError = LoginError.None).updateSSOLoginEnabled()
+            updateSSOFlowState(LoginState.Default)
             openWebUrl.emit(url)
         }
     }
+
+    companion object {
+        const val SSO_CODE_SAVED_STATE_KEY = "sso_code"
+    }
 }
 
 private fun SSOInitiateLoginResult.Failure.toLoginSSOError() = when (this) {
-    SSOInitiateLoginResult.Failure.InvalidCodeFormat -> LoginError.TextFieldError.InvalidValue
-    SSOInitiateLoginResult.Failure.InvalidCode -> LoginError.DialogError.InvalidSSOCodeError
-    is SSOInitiateLoginResult.Failure.Generic -> LoginError.DialogError.GenericError(this.genericFailure)
+    SSOInitiateLoginResult.Failure.InvalidCodeFormat -> LoginState.Error.TextFieldError.InvalidValue
+    SSOInitiateLoginResult.Failure.InvalidCode -> LoginState.Error.DialogError.InvalidSSOCodeError
+    is SSOInitiateLoginResult.Failure.Generic -> LoginState.Error.DialogError.GenericError(this.genericFailure)
     SSOInitiateLoginResult.Failure.InvalidRedirect ->
-        LoginError.DialogError.GenericError(CoreFailure.Unknown(IllegalArgumentException("Invalid Redirect")))
+        LoginState.Error.DialogError.GenericError(CoreFailure.Unknown(IllegalArgumentException("Invalid Redirect")))
 }
 
 private fun SSOLoginSessionResult.Failure.toLoginError() = when (this) {
-    SSOLoginSessionResult.Failure.InvalidCookie -> LoginError.DialogError.InvalidSSOCookie
-    is SSOLoginSessionResult.Failure.Generic -> LoginError.DialogError.GenericError(this.genericFailure)
+    SSOLoginSessionResult.Failure.InvalidCookie -> LoginState.Error.DialogError.InvalidSSOCookie
+    is SSOLoginSessionResult.Failure.Generic -> LoginState.Error.DialogError.GenericError(this.genericFailure)
 }
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCode.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCode.kt
index b4463654c50..c4358f73c46 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCode.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCode.kt
@@ -22,6 +22,7 @@ import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.text.input.TextFieldState
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.remember
@@ -30,22 +31,21 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.focus.FocusRequester
 import androidx.compose.ui.focus.focusRequester
 import androidx.compose.ui.res.stringResource
-import androidx.compose.ui.text.input.TextFieldValue
 import com.wire.android.R
 import com.wire.android.ui.common.progress.WireCircularProgressIndicator
-import com.wire.android.ui.common.textfield.CodeFieldValue
 import com.wire.android.ui.common.textfield.CodeTextField
 import com.wire.android.ui.common.textfield.WireTextFieldState
+import com.wire.android.ui.theme.WireTheme
 import com.wire.android.ui.theme.wireColorScheme
 import com.wire.android.ui.theme.wireDimensions
+import com.wire.android.util.ui.PreviewMultipleThemes
 
 @Composable
 fun VerificationCode(
     codeLength: Int,
-    currentCode: TextFieldValue,
+    codeState: TextFieldState,
     isLoading: Boolean,
     isCurrentCodeInvalid: Boolean,
-    onCodeChange: (CodeFieldValue) -> Unit,
     onResendCode: () -> Unit
 ) {
     val focusRequester = remember { FocusRequester() }
@@ -61,9 +61,8 @@ fun VerificationCode(
 
         CodeTextField(
             codeLength = codeLength,
-            value = currentCode,
+            textState = codeState,
             state = state,
-            onValueChange = onCodeChange,
             modifier = Modifier.focusRequester(focusRequester)
         )
 
@@ -81,3 +80,15 @@ fun VerificationCode(
         )
     }
 }
+
+@PreviewMultipleThemes
+@Composable
+fun PreviewVerificationCode() = WireTheme {
+    VerificationCode(
+        codeLength = 6,
+        codeState = TextFieldState(),
+        isLoading = false,
+        isCurrentCodeInvalid = false,
+        onResendCode = {}
+    )
+}
diff --git a/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCodeState.kt b/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCodeState.kt
index bcedbfce1e7..7ccfeb60541 100644
--- a/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCodeState.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/authentication/verificationcode/VerificationCodeState.kt
@@ -18,14 +18,10 @@
 
 package com.wire.android.ui.authentication.verificationcode
 
-import androidx.compose.ui.text.input.TextFieldValue
-import com.wire.android.ui.common.textfield.CodeFieldValue
-
 data class VerificationCodeState(
     val codeLength: Int = DEFAULT_VERIFICATION_CODE_LENGTH,
     val emailUsed: String = "",
     val isCodeInputNecessary: Boolean = false,
-    val codeInput: CodeFieldValue = CodeFieldValue(TextFieldValue(""), false),
     val isCurrentCodeInvalid: Boolean = false,
 ) {
     companion object {
diff --git a/app/src/main/kotlin/com/wire/android/ui/common/textfield/CodeTextField.kt b/app/src/main/kotlin/com/wire/android/ui/common/textfield/CodeTextField.kt
index 66f104b5080..3b7780197f3 100644
--- a/app/src/main/kotlin/com/wire/android/ui/common/textfield/CodeTextField.kt
+++ b/app/src/main/kotlin/com/wire/android/ui/common/textfield/CodeTextField.kt
@@ -37,7 +37,6 @@ import androidx.compose.ui.res.integerResource
 import androidx.compose.ui.text.TextStyle
 import androidx.compose.ui.text.input.ImeAction
 import androidx.compose.ui.text.input.KeyboardType
-import androidx.compose.ui.text.input.TextFieldValue
 import androidx.compose.ui.unit.Dp
 import com.wire.android.R
 import com.wire.android.ui.theme.WireTheme
@@ -86,58 +85,6 @@ fun CodeTextField(
     )
 }
 
-/*
-TODO: BasicTextField2 (value, onValueChange) overload is removed completely in compose foundation 1.7.0,
-      for now we can use our custom StateSyncingModifier to sync TextFieldValue with TextFieldState,
-      but eventually we should migrate and remove this function when all usages are replaced with the TextFieldState.
-*/
-@Deprecated("Use the new one with TextFieldState.")
-@Composable
-fun CodeTextField(
-    codeLength: Int = integerResource(id = R.integer.code_length),
-    value: TextFieldValue,
-    onValueChange: (CodeFieldValue) -> Unit,
-    shape: Shape = RoundedCornerShape(MaterialTheme.wireDimensions.corner4x),
-    colors: WireTextFieldColors = wireTextFieldColors(),
-    textStyle: TextStyle = MaterialTheme.wireTypography.code01,
-    state: WireTextFieldState = WireTextFieldState.Default,
-    interactionSource: MutableInteractionSource = remember { MutableInteractionSource() },
-    maxHorizontalSpacing: Dp = MaterialTheme.wireDimensions.spacing16x,
-    horizontalAlignment: Alignment.Horizontal = Alignment.CenterHorizontally,
-    modifier: Modifier = Modifier
-) {
-    val keyboardController = LocalSoftwareKeyboardController.current
-    val enabled = state !is WireTextFieldState.Disabled
-    val textState = remember { TextFieldState(value.text, value.selection) }
-    val onValueChanged: (TextFieldValue) -> Unit = { onValueChange(CodeFieldValue(it, it.text.length == codeLength)) }
-    CodeTextFieldLayout(
-        textState = textState,
-        codeLength = codeLength,
-        shape = shape,
-        colors = colors,
-        textStyle = textStyle,
-        state = state,
-        maxHorizontalSpacing = maxHorizontalSpacing,
-        horizontalAlignment = horizontalAlignment,
-        modifier = modifier,
-        innerBasicTextField = { decorator, textFieldModifier ->
-            BasicTextField(
-                state = textState,
-                textStyle = textStyle,
-                enabled = enabled,
-                keyboardOptions = KeyboardOptions.DefaultCode,
-                onKeyboardAction = { keyboardController?.hide() },
-                interactionSource = interactionSource,
-                inputTransformation = MaxLengthDigitsFilter(codeLength),
-                decorator = decorator,
-                modifier = textFieldModifier.then(StateSyncingModifier(textState, value, onValueChanged)),
-            )
-        }
-    )
-}
-
-data class CodeFieldValue(val text: TextFieldValue, val isFullyFilled: Boolean)
-
 @Stable
 val KeyboardOptions.Companion.DefaultCode: KeyboardOptions
     get() = Default.copy(
diff --git a/app/src/test/kotlin/com/wire/android/ui/authentication/LoginViewModelTest.kt b/app/src/test/kotlin/com/wire/android/ui/authentication/LoginViewModelTest.kt
index f607f5ea503..e3ff2cd3d15 100644
--- a/app/src/test/kotlin/com/wire/android/ui/authentication/LoginViewModelTest.kt
+++ b/app/src/test/kotlin/com/wire/android/ui/authentication/LoginViewModelTest.kt
@@ -18,7 +18,6 @@
 
 package com.wire.android.ui.authentication
 
-import androidx.lifecycle.SavedStateHandle
 import com.wire.android.config.CoroutineTestExtension
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.AuthServerConfigProvider
@@ -42,9 +41,6 @@ class LoginViewModelTest {
     @MockK
     private lateinit var clientScopeProviderFactory: ClientScopeProvider.Factory
 
-    @MockK
-    private lateinit var savedStateHandle: SavedStateHandle
-
     @MockK
     private lateinit var qualifiedIdMapper: QualifiedIdMapper
 
@@ -62,11 +58,9 @@ class LoginViewModelTest {
     @BeforeEach
     fun setup() {
         MockKAnnotations.init(this)
-        every { savedStateHandle.get<String>(any()) } returns null
         every { qualifiedIdMapper.fromStringToQualifiedID(any()) } returns QualifiedID("", "")
         every { authServerConfigProvider.authServer.value } returns newServerConfig(1).links
         loginViewModel = LoginViewModel(
-            savedStateHandle,
             clientScopeProviderFactory,
             authServerConfigProvider,
             userDataStoreProvider,
diff --git a/app/src/test/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModelTest.kt b/app/src/test/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModelTest.kt
index c4c079e269f..2e79269e277 100644
--- a/app/src/test/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModelTest.kt
+++ b/app/src/test/kotlin/com/wire/android/ui/authentication/login/email/LoginEmailViewModelTest.kt
@@ -20,17 +20,17 @@
 
 package com.wire.android.ui.authentication.login.email
 
-import androidx.compose.ui.text.input.TextFieldValue
+import androidx.compose.foundation.text.input.setTextAndPlaceCursorAtEnd
 import androidx.lifecycle.SavedStateHandle
 import com.wire.android.config.CoroutineTestExtension
+import com.wire.android.config.SnapshotExtension
 import com.wire.android.config.TestDispatcherProvider
 import com.wire.android.config.mockUri
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.AuthServerConfigProvider
 import com.wire.android.di.ClientScopeProvider
 import com.wire.android.framework.TestClient
-import com.wire.android.ui.authentication.login.LoginError
-import com.wire.android.ui.common.textfield.CodeFieldValue
+import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.util.EMPTY
 import com.wire.android.util.newServerConfig
 import com.wire.kalium.logic.CoreFailure
@@ -68,12 +68,13 @@ import kotlinx.coroutines.test.runTest
 import org.amshove.kluent.shouldBe
 import org.amshove.kluent.shouldBeEqualTo
 import org.amshove.kluent.shouldBeInstanceOf
+import org.amshove.kluent.shouldNotBeInstanceOf
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.ExtendWith
 
 @OptIn(ExperimentalCoroutinesApi::class)
-@ExtendWith(CoroutineTestExtension::class)
+@ExtendWith(CoroutineTestExtension::class, SnapshotExtension::class)
 class LoginEmailViewModelTest {
 
     @MockK
@@ -115,9 +116,6 @@ class LoginEmailViewModelTest {
     @MockK
     private lateinit var authenticationScope: AuthenticationScope
 
-    @MockK(relaxed = true)
-    private lateinit var onSuccess: (Boolean, Boolean) -> Unit
-
     private lateinit var loginViewModel: LoginEmailViewModel
 
     private val userId: QualifiedID = QualifiedID("userId", "domain")
@@ -158,18 +156,18 @@ class LoginEmailViewModelTest {
 
     @Test
     fun `given empty strings, when entering credentials, then button is disabled`() {
-        loginViewModel.onPasswordChange(TextFieldValue(String.EMPTY))
-        loginViewModel.onUserIdentifierChange(TextFieldValue(String.EMPTY))
-        loginViewModel.loginState.emailLoginEnabled shouldBeEqualTo false
-        loginViewModel.loginState.emailLoginLoading shouldBeEqualTo false
+        loginViewModel.passwordTextState.setTextAndPlaceCursorAtEnd(String.EMPTY)
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(String.EMPTY)
+        loginViewModel.loginState.loginEnabled shouldBeEqualTo false
+        loginViewModel.loginState.flowState.shouldNotBeInstanceOf<LoginState.Loading>()
     }
 
     @Test
     fun `given non-empty strings, when entering credentials, then button is enabled`() {
-        loginViewModel.onPasswordChange(TextFieldValue("abc"))
-        loginViewModel.onUserIdentifierChange(TextFieldValue("abc"))
-        loginViewModel.loginState.emailLoginEnabled shouldBeEqualTo true
-        loginViewModel.loginState.emailLoginLoading shouldBeEqualTo false
+        loginViewModel.passwordTextState.setTextAndPlaceCursorAtEnd("abc")
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd("abc")
+        loginViewModel.loginState.loginEnabled shouldBeEqualTo true
+        loginViewModel.loginState.flowState.shouldNotBeInstanceOf<LoginState.Loading>()
     }
 
     @Test
@@ -181,14 +179,16 @@ class LoginEmailViewModelTest {
             addAuthenticatedUserUseCase(any(), any(), any(), any())
         } returns AddAuthenticatedUserUseCase.Result.Success(userId)
 
-        loginViewModel.onPasswordChange(TextFieldValue("abc"))
-        loginViewModel.onUserIdentifierChange(TextFieldValue("abc"))
-        loginViewModel.loginState.emailLoginEnabled shouldBeEqualTo true
-        loginViewModel.loginState.emailLoginLoading shouldBeEqualTo false
-        loginViewModel.login(onSuccess)
+        loginViewModel.passwordTextState.setTextAndPlaceCursorAtEnd("abc")
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd("abc")
+        loginViewModel.loginState.loginEnabled shouldBeEqualTo true
+        loginViewModel.loginState.flowState.shouldNotBeInstanceOf<LoginState.Loading>()
+        loginViewModel.login()
+        loginViewModel.loginState.loginEnabled shouldBeEqualTo false
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Loading>()
         advanceUntilIdle()
-        loginViewModel.loginState.emailLoginEnabled shouldBeEqualTo true
-        loginViewModel.loginState.emailLoginLoading shouldBeEqualTo false
+        loginViewModel.loginState.loginEnabled shouldBeEqualTo true
+        loginViewModel.loginState.flowState.shouldNotBeInstanceOf<LoginState.Loading>()
     }
 
     @Test
@@ -204,13 +204,16 @@ class LoginEmailViewModelTest {
         coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
         every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(true)
 
-        loginViewModel.onPasswordChange(TextFieldValue(password))
+        loginViewModel.passwordTextState.setTextAndPlaceCursorAtEnd(password)
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
         coVerify(exactly = 1) { loginUseCase(any(), any(), any(), any(), any()) }
         coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
-        coVerify(exactly = 1) { onSuccess(true, false) }
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Success>().let {
+            it.initialSyncCompleted shouldBe true
+            it.isE2EIRequired shouldBe false
+        }
     }
 
     @Test
@@ -225,16 +228,19 @@ class LoginEmailViewModelTest {
             )
             coEvery { addAuthenticatedUserUseCase(any(), any(), any(), any()) } returns AddAuthenticatedUserUseCase.Result.Success(userId)
             coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
-        every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(false)
+            every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(false)
 
-            loginViewModel.onPasswordChange(TextFieldValue(password))
+            loginViewModel.passwordTextState.setTextAndPlaceCursorAtEnd(password)
 
-            loginViewModel.login(onSuccess)
+            loginViewModel.login()
             advanceUntilIdle()
-        coVerify(exactly = 1) { loginUseCase(any(), any(), any(), any(), any()) }
-        coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
-        coVerify(exactly = 1) { onSuccess(false, false) }
-    }
+            coVerify(exactly = 1) { loginUseCase(any(), any(), any(), any(), any()) }
+            coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Success>().let {
+                it.initialSyncCompleted shouldBe false
+                it.isE2EIRequired shouldBe false
+            }
+        }
 
     @Test
     fun `given button is clicked, when login returns InvalidUserIdentifier error, then InvalidUserIdentifierError is passed`() = runTest {
@@ -242,9 +248,9 @@ class LoginEmailViewModelTest {
             loginUseCase(any(), any(), any(), any(), any())
         } returns AuthenticationResult.Failure.InvalidUserIdentifier
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.TextFieldError.InvalidValue::class
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.TextFieldError.InvalidValue>()
     }
 
     @Test
@@ -259,9 +265,9 @@ class LoginEmailViewModelTest {
             )
         } returns AuthenticationResult.Failure.InvalidCredentials.InvalidPasswordIdentityCombination
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.InvalidCredentialsError::class
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.InvalidCredentialsError>()
     }
 
     @Test
@@ -271,11 +277,12 @@ class LoginEmailViewModelTest {
             loginUseCase(any(), any(), any(), any(), any())
         } returns AuthenticationResult.Failure.Generic(networkFailure)
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
 
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.GenericError::class
-        (loginViewModel.loginState.loginError as LoginError.DialogError.GenericError).coreFailure shouldBe networkFailure
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.GenericError>().let {
+            it.coreFailure shouldBe networkFailure
+        }
     }
 
     @Test
@@ -290,11 +297,11 @@ class LoginEmailViewModelTest {
             )
         } returns AuthenticationResult.Failure.InvalidCredentials.InvalidPasswordIdentityCombination
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.InvalidCredentialsError::class
-        loginViewModel.onDialogDismiss()
-        loginViewModel.loginState.loginError shouldBe LoginError.None
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.InvalidCredentialsError>()
+        loginViewModel.clearLoginErrors()
+        loginViewModel.loginState.flowState.shouldNotBeInstanceOf<LoginState.Error>()
     }
 
     @Test
@@ -314,10 +321,10 @@ class LoginEmailViewModelTest {
             )
         } returns AddAuthenticatedUserUseCase.Result.Failure.UserAlreadyExists
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
 
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.UserAlreadyExists::class
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.UserAlreadyExists>()
     }
 
     @Test
@@ -325,9 +332,9 @@ class LoginEmailViewModelTest {
         val email = "some.email@example.org"
         coEvery { loginUseCase(any(), any(), any(), any(), any()) } returns AuthenticationResult.Failure.InvalidCredentials.Missing2FA
         coEvery { requestSecondFactorCodeUseCase(any(), any()) } returns RequestSecondFactorVerificationCodeUseCase.Result.Success
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
 
         coVerify(exactly = 1) { requestSecondFactorCodeUseCase(email, VerifiableAction.LOGIN_OR_CLIENT_REGISTRATION) }
@@ -339,13 +346,13 @@ class LoginEmailViewModelTest {
         coEvery { loginUseCase(any(), any(), any(), any(), any()) } returns AuthenticationResult.Failure.InvalidCredentials.Missing2FA
         coEvery { requestSecondFactorCodeUseCase(any(), any()) } returns RequestSecondFactorVerificationCodeUseCase.Result.Success
 
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
-        loginViewModel.onPasswordChange(TextFieldValue("somePassword"))
-        loginViewModel.login(onSuccess)
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
+        loginViewModel.passwordTextState.setTextAndPlaceCursorAtEnd("somePassword")
+        loginViewModel.login()
         advanceUntilIdle()
 
-        loginViewModel.loginState.emailLoginLoading shouldBe false
-        loginViewModel.loginState.emailLoginEnabled shouldBe true
+        loginViewModel.loginState.flowState.shouldNotBeInstanceOf<LoginState.Loading>()
+        loginViewModel.loginState.loginEnabled shouldBe true
     }
 
     @Test
@@ -353,9 +360,9 @@ class LoginEmailViewModelTest {
         val email = "some.email@example.org"
         coEvery { loginUseCase(any(), any(), any(), any(), any()) } returns AuthenticationResult.Failure.InvalidCredentials.Missing2FA
         coEvery { requestSecondFactorCodeUseCase(any(), any()) } returns RequestSecondFactorVerificationCodeUseCase.Result.Success
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
 
         loginViewModel.secondFactorVerificationCodeState.isCodeInputNecessary shouldBe true
@@ -375,9 +382,9 @@ class LoginEmailViewModelTest {
             } returns RequestSecondFactorVerificationCodeUseCase.Result.Failure.Generic(
                 CoreFailure.Unknown(null)
             )
-            loginViewModel.onUserIdentifierChange(TextFieldValue(email))
+            loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
 
-            loginViewModel.login(onSuccess)
+            loginViewModel.login()
             advanceUntilIdle()
 
             loginViewModel.secondFactorVerificationCodeState.isCodeInputNecessary shouldBe false
@@ -391,9 +398,9 @@ class LoginEmailViewModelTest {
         coEvery {
             requestSecondFactorCodeUseCase(any(), any())
         } returns RequestSecondFactorVerificationCodeUseCase.Result.Failure.TooManyRequests
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
 
         loginViewModel.secondFactorVerificationCodeState.isCodeInputNecessary shouldBe true
@@ -405,8 +412,8 @@ class LoginEmailViewModelTest {
         val email = "some.email@example.org"
         coEvery { loginUseCase(any(), any(), any(), any(), any()) } returns AuthenticationResult.Failure.InvalidCredentials.Missing2FA
         coEvery { requestSecondFactorCodeUseCase(any(), any()) } returns RequestSecondFactorVerificationCodeUseCase.Result.Success
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
-        loginViewModel.login(onSuccess)
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
+        loginViewModel.login()
         advanceUntilIdle()
         loginViewModel.secondFactorVerificationCodeState.isCodeInputNecessary shouldBe true
     }
@@ -414,9 +421,9 @@ class LoginEmailViewModelTest {
     @Test
     fun `given login fails with invalid 2fa, when logging in, then should mark the current code as invalid`() = runTest {
         coEvery { loginUseCase(any(), any(), any(), any(), any()) } returns AuthenticationResult.Failure.InvalidCredentials.Invalid2FA
-        loginViewModel.onUserIdentifierChange(TextFieldValue("some.email@example.org"))
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd("some.email@example.org")
 
-        loginViewModel.login(onSuccess)
+        loginViewModel.login()
         advanceUntilIdle()
         loginViewModel.secondFactorVerificationCodeState.isCurrentCodeInvalid shouldBe true
     }
@@ -435,12 +442,12 @@ class LoginEmailViewModelTest {
         coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
         every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(true)
 
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
-        loginViewModel.onCodeChange(CodeFieldValue(TextFieldValue(code), true), onSuccess)
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
+        loginViewModel.secondFactorVerificationCodeTextState.setTextAndPlaceCursorAtEnd(code)
         advanceUntilIdle()
         coVerify(exactly = 1) { loginUseCase(email, any(), any(), any(), code) }
         coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
-        coVerify(exactly = 1) { onSuccess(any(), any()) }
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Success>()
     }
 
     @Test
@@ -458,8 +465,8 @@ class LoginEmailViewModelTest {
         coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Failure.TooManyClients
         every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(true)
 
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
-        loginViewModel.onCodeChange(CodeFieldValue(TextFieldValue(code), true), onSuccess)
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
+        loginViewModel.secondFactorVerificationCodeTextState.setTextAndPlaceCursorAtEnd(code)
         advanceUntilIdle()
         loginViewModel.secondFactorVerificationCodeState.isCodeInputNecessary shouldBe false
     }
@@ -478,8 +485,8 @@ class LoginEmailViewModelTest {
         coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
         every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(true)
 
-        loginViewModel.onUserIdentifierChange(TextFieldValue(email))
-        loginViewModel.onCodeChange(CodeFieldValue(TextFieldValue(code), true), onSuccess)
+        loginViewModel.userIdentifierTextState.setTextAndPlaceCursorAtEnd(email)
+        loginViewModel.secondFactorVerificationCodeTextState.setTextAndPlaceCursorAtEnd(code)
         advanceUntilIdle()
         coVerify(exactly = 1) { getOrRegisterClientUseCase(match { it.secondFactorVerificationCode == null }) }
     }
diff --git a/app/src/test/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModelTest.kt b/app/src/test/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModelTest.kt
index 3d4ca18db66..db95f56468c 100644
--- a/app/src/test/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModelTest.kt
+++ b/app/src/test/kotlin/com/wire/android/ui/authentication/login/sso/LoginSSOViewModelTest.kt
@@ -18,15 +18,16 @@
 
 package com.wire.android.ui.authentication.login.sso
 
-import androidx.compose.ui.text.input.TextFieldValue
+import androidx.compose.foundation.text.input.setTextAndPlaceCursorAtEnd
 import androidx.lifecycle.SavedStateHandle
 import com.wire.android.config.CoroutineTestExtension
+import com.wire.android.config.SnapshotExtension
 import com.wire.android.config.mockUri
 import com.wire.android.datastore.UserDataStoreProvider
 import com.wire.android.di.AuthServerConfigProvider
 import com.wire.android.di.ClientScopeProvider
 import com.wire.android.framework.TestClient
-import com.wire.android.ui.authentication.login.LoginError
+import com.wire.android.ui.authentication.login.LoginState
 import com.wire.android.ui.common.dialogs.CustomServerDialogState
 import com.wire.android.util.EMPTY
 import com.wire.android.util.deeplink.DeepLinkResult
@@ -59,7 +60,6 @@ import io.mockk.coEvery
 import io.mockk.coVerify
 import io.mockk.every
 import io.mockk.impl.annotations.MockK
-import io.mockk.verify
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.advanceUntilIdle
@@ -68,15 +68,14 @@ import org.amshove.kluent.internal.assertEquals
 import org.amshove.kluent.shouldBe
 import org.amshove.kluent.shouldBeEqualTo
 import org.amshove.kluent.shouldBeInstanceOf
-import org.junit.jupiter.api.Assertions.assertFalse
-import org.junit.jupiter.api.Assertions.assertTrue
+import org.amshove.kluent.shouldNotBeInstanceOf
 import org.junit.jupiter.api.BeforeEach
 import org.junit.jupiter.api.Test
 import org.junit.jupiter.api.extension.ExtendWith
 import java.io.IOException
 
 @OptIn(ExperimentalCoroutinesApi::class)
-@ExtendWith(CoroutineTestExtension::class)
+@ExtendWith(CoroutineTestExtension::class, SnapshotExtension::class)
 class LoginSSOViewModelTest {
 
     @MockK
@@ -120,9 +119,6 @@ class LoginSSOViewModelTest {
     @MockK
     private lateinit var fetchSSOSettings: FetchSSOSettingsUseCase
 
-    @MockK(relaxed = true)
-    private lateinit var onSuccess: (Boolean, Boolean) -> Unit
-
     private lateinit var loginViewModel: LoginSSOViewModel
 
     private val userId: QualifiedID = QualifiedID("userId", "domain")
@@ -161,16 +157,16 @@ class LoginSSOViewModelTest {
 
     @Test
     fun `given empty string, when entering code, then button is disabled`() {
-        loginViewModel.onSSOCodeChange(TextFieldValue(String.EMPTY))
-        loginViewModel.loginState.ssoLoginEnabled shouldBeEqualTo false
-        loginViewModel.loginState.ssoLoginLoading shouldBeEqualTo false
+        loginViewModel.ssoTextState.setTextAndPlaceCursorAtEnd(String.EMPTY)
+        loginViewModel.loginState.loginEnabled shouldBeEqualTo false
+        loginViewModel.loginState.flowState shouldNotBeInstanceOf LoginState.Loading::class
     }
 
     @Test
     fun `given non-empty string, when entering code, then button is enabled`() {
-        loginViewModel.onSSOCodeChange(TextFieldValue("abc"))
-        loginViewModel.loginState.ssoLoginEnabled shouldBeEqualTo true
-        loginViewModel.loginState.ssoLoginLoading shouldBeEqualTo false
+        loginViewModel.ssoTextState.setTextAndPlaceCursorAtEnd("abc")
+        loginViewModel.loginState.loginEnabled shouldBeEqualTo true
+        loginViewModel.loginState.flowState shouldNotBeInstanceOf LoginState.Loading::class
     }
 
     @Test
@@ -180,7 +176,7 @@ class LoginSSOViewModelTest {
         val url = "https://wire.com/sso"
         coEvery { ssoInitiateLoginUseCase(param) } returns SSOInitiateLoginResult.Success(url)
         every { validateEmailUseCase(any()) } returns false
-        loginViewModel.onSSOCodeChange(TextFieldValue(ssoCode))
+        loginViewModel.ssoTextState.setTextAndPlaceCursorAtEnd(ssoCode)
         loginViewModel.login()
         advanceUntilIdle()
         // loginViewModel.openWebUrl.firstOrNull() shouldBe url
@@ -195,7 +191,7 @@ class LoginSSOViewModelTest {
 
             loginViewModel.login()
             advanceUntilIdle()
-            loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.TextFieldError.InvalidValue::class
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.TextFieldError.InvalidValue>()
         }
 
     @Test
@@ -206,7 +202,7 @@ class LoginSSOViewModelTest {
         loginViewModel.login()
         advanceUntilIdle()
 
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.InvalidSSOCodeError::class
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.InvalidSSOCodeError>()
     }
 
     @Test
@@ -218,11 +214,9 @@ class LoginSSOViewModelTest {
             loginViewModel.login()
             advanceUntilIdle()
 
-            loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.GenericError::class
-            with(loginViewModel.loginState.loginError as LoginError.DialogError.GenericError) {
-                coreFailure shouldBeInstanceOf CoreFailure.Unknown::class
-                with(coreFailure as CoreFailure.Unknown) {
-                    this.rootCause shouldBeInstanceOf IllegalArgumentException::class
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.GenericError>().let {
+                it.coreFailure.shouldBeInstanceOf<CoreFailure.Unknown>().let {
+                    it.rootCause.shouldBeInstanceOf<IllegalArgumentException>()
                 }
             }
         }
@@ -236,8 +230,9 @@ class LoginSSOViewModelTest {
         loginViewModel.login()
         advanceUntilIdle()
 
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.GenericError::class
-        (loginViewModel.loginState.loginError as LoginError.DialogError.GenericError).coreFailure shouldBe networkFailure
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.GenericError>().let {
+            it.coreFailure shouldBe networkFailure
+        }
     }
 
     @Test
@@ -256,13 +251,16 @@ class LoginSSOViewModelTest {
         coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
         every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(false)
 
-        loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id, onSuccess)
+        loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id)
         advanceUntilIdle()
 
         coVerify(exactly = 1) { getSSOLoginSessionUseCase(any()) }
         coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
         coVerify(exactly = 1) { addAuthenticatedUserUseCase(any(), any(), any(), any()) }
-        coVerify(exactly = 1) { onSuccess(false, false) }
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Success>().let {
+            it.initialSyncCompleted shouldBe false
+            it.isE2EIRequired shouldBe false
+        }
     }
 
     @Test
@@ -282,14 +280,17 @@ class LoginSSOViewModelTest {
             coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
             every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(true)
 
-            loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id, onSuccess)
+            loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id)
             advanceUntilIdle()
 
-        coVerify(exactly = 1) { getSSOLoginSessionUseCase(any()) }
-        coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
-        coVerify(exactly = 1) { addAuthenticatedUserUseCase(any(), any(), any(), any()) }
-        coVerify(exactly = 1) { onSuccess(true, false) }
-    }
+            coVerify(exactly = 1) { getSSOLoginSessionUseCase(any()) }
+            coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
+            coVerify(exactly = 1) { addAuthenticatedUserUseCase(any(), any(), any(), any()) }
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Success>().let {
+                it.initialSyncCompleted shouldBe true
+                it.isE2EIRequired shouldBe false
+            }
+        }
 
     @Test
     fun `given establishSSOSession is called, when SSOLoginSessionResult return InvalidCookie, then SSOLoginResult fails`() = runTest {
@@ -306,31 +307,30 @@ class LoginSSOViewModelTest {
         )
         coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
 
-        loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id, onSuccess)
+        loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id)
         advanceUntilIdle()
-        loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.InvalidSSOCookie::class
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.InvalidSSOCookie>()
         coVerify(exactly = 1) { getSSOLoginSessionUseCase(any()) }
         coVerify(exactly = 0) { loginViewModel.registerClient(any(), null) }
         coVerify(exactly = 0) { addAuthenticatedUserUseCase(any(), any(), any(), any()) }
-        verify(exactly = 0) { onSuccess(any(), any()) }
     }
 
     @Test
     fun `given HandleSSOResult is called, when ssoResult is null, then loginSSOError state should be none`() =
         runTest {
-            loginViewModel.handleSSOResult(null, onSuccess)
+            loginViewModel.handleSSOResult(null)
             advanceUntilIdle()
-            loginViewModel.loginState.loginError shouldBeEqualTo LoginError.None
+            loginViewModel.loginState.flowState.shouldNotBeInstanceOf<LoginState.Error>()
         }
 
     @Test
     fun `given HandleSSOResult is called, when ssoResult is failure, then loginSSOError state should be dialog error`() =
         runTest {
-            loginViewModel.handleSSOResult(DeepLinkResult.SSOLogin.Failure(SSOFailureCodes.Unknown), onSuccess)
+            loginViewModel.handleSSOResult(DeepLinkResult.SSOLogin.Failure(SSOFailureCodes.Unknown))
             advanceUntilIdle()
-            loginViewModel.loginState.loginError shouldBeEqualTo LoginError.DialogError.SSOResultError(
-                SSOFailureCodes.Unknown
-            )
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.SSOResultError>().let {
+                it.result shouldBe SSOFailureCodes.Unknown
+            }
         }
 
     @Test
@@ -350,10 +350,10 @@ class LoginSSOViewModelTest {
             coEvery { getOrRegisterClientUseCase(any()) } returns RegisterClientResult.Success(CLIENT)
             every { userDataStoreProvider.getOrCreate(any()).initialSyncCompleted } returns flowOf(true)
 
-            loginViewModel.handleSSOResult(DeepLinkResult.SSOLogin.Success("", ""), onSuccess)
+            loginViewModel.handleSSOResult(DeepLinkResult.SSOLogin.Success("", ""))
             advanceUntilIdle()
 
-            verify(exactly = 1) { onSuccess(any(), any()) }
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Success>()
         }
 
     @Test
@@ -369,14 +369,13 @@ class LoginSSOViewModelTest {
                 )
             } returns AddAuthenticatedUserUseCase.Result.Failure.UserAlreadyExists
 
-            loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id, onSuccess)
+            loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id)
             advanceUntilIdle()
 
-            loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.DialogError.UserAlreadyExists::class
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.UserAlreadyExists>()
             coVerify(exactly = 1) { getSSOLoginSessionUseCase(any()) }
             coVerify(exactly = 0) { loginViewModel.registerClient(any(), null) }
             coVerify(exactly = 1) { addAuthenticatedUserUseCase(any(), any(), any(), any()) }
-            verify(exactly = 0) { onSuccess(any(), any()) }
         }
 
     @Test
@@ -397,23 +396,22 @@ class LoginSSOViewModelTest {
                 getOrRegisterClientUseCase(any())
             } returns RegisterClientResult.Failure.TooManyClients
 
-            loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id, onSuccess)
+            loginViewModel.establishSSOSession("", serverConfigId = SERVER_CONFIG.id)
             advanceUntilIdle()
 
-            loginViewModel.loginState.loginError shouldBeInstanceOf LoginError.TooManyDevicesError::class
+            loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.TooManyDevicesError>()
 
-        coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
-        coVerify(exactly = 1) { getSSOLoginSessionUseCase(any()) }
-        coVerify(exactly = 1) { addAuthenticatedUserUseCase(any(), any(), any(), any()) }
-        verify(exactly = 0) { onSuccess(any(), any()) }
-    }
+            coVerify(exactly = 1) { getOrRegisterClientUseCase(any()) }
+            coVerify(exactly = 1) { getSSOLoginSessionUseCase(any()) }
+            coVerify(exactly = 1) { addAuthenticatedUserUseCase(any(), any(), any(), any()) }
+        }
 
     @Test
     fun `given email, when clicking login, then start the domain lookup flow`() = runTest {
         val expected = newServerConfig(2).links
         every { validateEmailUseCase(any()) } returns true
         coEvery { authenticationScope.domainLookup(any()) } returns DomainLookupUseCase.Result.Success(expected)
-        loginViewModel.onSSOCodeChange(TextFieldValue("email@wire.com"))
+        loginViewModel.ssoTextState.setTextAndPlaceCursorAtEnd("email@wire.com")
 
         loginViewModel.domainLookupFlow()
         advanceUntilIdle()
@@ -491,16 +489,16 @@ class LoginSSOViewModelTest {
         val expected = CoreFailure.Unknown(IOException())
         every { validateEmailUseCase(any()) } returns true
         coEvery { authenticationScope.domainLookup(any()) } returns DomainLookupUseCase.Result.Failure(expected)
-        loginViewModel.onSSOCodeChange(TextFieldValue("email@wire.com"))
+        loginViewModel.ssoTextState.setTextAndPlaceCursorAtEnd("email@wire.com")
 
         loginViewModel.domainLookupFlow()
         advanceUntilIdle()
 
         coVerify(exactly = 1) { authenticationScope.domainLookup("email@wire.com") }
         loginViewModel.loginState.customServerDialogState shouldBe null
-        assertTrue(loginViewModel.loginState.loginError is LoginError.DialogError.GenericError)
-        assertEquals(expected, (loginViewModel.loginState.loginError as LoginError.DialogError.GenericError).coreFailure)
-        assertFalse(loginViewModel.loginState.ssoLoginLoading)
+        loginViewModel.loginState.flowState.shouldBeInstanceOf<LoginState.Error.DialogError.GenericError>().let {
+            it.coreFailure shouldBe expected
+        }
     }
 
     companion object {