/*
 * Copyright 2026 Signal Messenger, LLC
 * SPDX-License-Identifier: AGPL-3.0-only
 */

package org.signal.registration

import androidx.lifecycle.SavedStateHandle
import assertk.assertThat
import assertk.assertions.isEqualTo
import assertk.assertions.isNotNull
import assertk.assertions.isNull
import assertk.assertions.isTrue
import io.mockk.coEvery
import io.mockk.coVerify
import io.mockk.mockk
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.test.StandardTestDispatcher
import kotlinx.coroutines.test.advanceUntilIdle
import kotlinx.coroutines.test.resetMain
import kotlinx.coroutines.test.runTest
import kotlinx.coroutines.test.setMain
import org.junit.After
import org.junit.Before
import org.junit.Test
import org.signal.core.models.AccountEntropyPool
import org.signal.core.models.MasterKey
import org.signal.network.api.RegistrationApiV2.SessionMetadata
import org.signal.network.api.RegistrationApiV2.SvrCredentials

@OptIn(ExperimentalCoroutinesApi::class)
class RegistrationViewModelTest {

  private val testDispatcher = StandardTestDispatcher()
  private lateinit var mockRepository: RegistrationRepository

  @Before
  fun setup() {
    Dispatchers.setMain(testDispatcher)
    mockRepository = mockk(relaxed = true)
  }

  @After
  fun tearDown() {
    Dispatchers.resetMain()
  }

  // ==================== Restore Flow State Tests ====================

  @Test
  fun `no saved state starts fresh and loads preExistingRegistrationData`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val state = viewModel.state.value
    assertThat(state.backStack).isEqualTo(listOf(RegistrationRoute.Welcome))
    assertThat(state.sessionMetadata).isNull()
  }

  @Test
  fun `restore with valid session proceeds normally with updated session metadata`() = runTest(testDispatcher) {
    val savedSession = createSessionMetadata("session-1")
    val freshSession = createSessionMetadata("session-1").copy(nextSms = 9999L)

    val savedState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry),
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.VerificationCodeEntry
      ),
      sessionMetadata = savedSession,
      sessionE164 = "+15551234567"
    )

    coEvery { mockRepository.restoreFlowState() } returns savedState
    coEvery { mockRepository.validateSession("session-1") } returns freshSession

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val state = viewModel.state.value
    assertThat(state.backStack).isEqualTo(savedState.backStack)
    assertThat(state.sessionMetadata).isEqualTo(freshSession)
    assertThat(state.sessionE164).isEqualTo("+15551234567")
  }

  @Test
  fun `restore with expired session and not registered resets to PhoneNumberEntry with e164 preserved`() = runTest(testDispatcher) {
    val savedSession = createSessionMetadata("session-expired")

    val savedState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry),
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.VerificationCodeEntry
      ),
      sessionMetadata = savedSession,
      sessionE164 = "+15559876543"
    )

    coEvery { mockRepository.restoreFlowState() } returns savedState
    coEvery { mockRepository.validateSession("session-expired") } returns null
    coEvery { mockRepository.isRegistered() } returns false

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val state = viewModel.state.value
    assertThat(state.backStack).isEqualTo(
      listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry),
        RegistrationRoute.PhoneNumberEntry
      )
    )
    assertThat(state.sessionMetadata).isNull()
    assertThat(state.sessionE164).isEqualTo("+15559876543")
  }

  @Test
  fun `restore with expired session and already registered proceeds with null session`() = runTest(testDispatcher) {
    val savedSession = createSessionMetadata("session-expired-2")

    val savedState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.PinCreate
      ),
      sessionMetadata = savedSession,
      sessionE164 = "+15551234567"
    )

    coEvery { mockRepository.restoreFlowState() } returns savedState
    coEvery { mockRepository.validateSession("session-expired-2") } returns null
    coEvery { mockRepository.isRegistered() } returns true

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val state = viewModel.state.value
    assertThat(state.backStack).isEqualTo(listOf(RegistrationRoute.Welcome, RegistrationRoute.PinCreate))
    assertThat(state.sessionMetadata).isNull()
    assertThat(state.sessionE164).isEqualTo("+15551234567")
  }

  @Test
  fun `restore with no session skips validation`() = runTest(testDispatcher) {
    val savedState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.PinCreate,
        RegistrationRoute.ArchiveRestoreSelection.forManualRestore()
      ),
      sessionMetadata = null,
      sessionE164 = "+15551234567",
      doNotAttemptRecoveryPassword = true
    )

    coEvery { mockRepository.restoreFlowState() } returns savedState

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val state = viewModel.state.value
    assertThat(state.backStack).isEqualTo(savedState.backStack)
    assertThat(state.sessionMetadata).isNull()
    assertThat(state.doNotAttemptRecoveryPassword).isEqualTo(true)

    coVerify(exactly = 0) { mockRepository.validateSession(any()) }
  }

  @Test
  fun `init clears isRestoringNavigationState after restore completes`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    assertThat(viewModel.state.value.isRestoringNavigationState).isEqualTo(false)
  }

  @Test
  fun `start fresh clears in-progress data and starts fresh from Welcome`() = runTest(testDispatcher) {
    val preExisting = mockk<PreExistingRegistrationData>(relaxed = true)
    val savedState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry),
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.VerificationCodeEntry
      ),
      sessionMetadata = createSessionMetadata("stale-session")
    )
    coEvery { mockRepository.restoreFlowState() } returns savedState
    coEvery { mockRepository.getPreExistingRegistrationData() } returns preExisting

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle(), startFresh = true)
    advanceUntilIdle()

    val state = viewModel.state.value
    assertThat(state.backStack).isEqualTo(listOf(RegistrationRoute.Welcome))
    assertThat(state.sessionMetadata).isNull()
    assertThat(state.preExistingRegistrationData).isEqualTo(preExisting)
    assertThat(state.isRestoringNavigationState).isEqualTo(false)

    coVerify { mockRepository.clearInProgressRegistrationData() }
    coVerify(exactly = 0) { mockRepository.restoreFlowState() }
  }

  @Test
  fun `start fresh does not clear again once reset has been performed`() = runTest(testDispatcher) {
    val savedState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry),
      sessionMetadata = null
    )
    coEvery { mockRepository.restoreFlowState() } returns savedState

    val savedStateHandle = SavedStateHandle(mapOf("start_fresh_reset_performed" to true))
    val viewModel = RegistrationViewModel(mockRepository, savedStateHandle, startFresh = true)
    advanceUntilIdle()

    assertThat(viewModel.state.value.backStack).isEqualTo(savedState.backStack)

    coVerify(exactly = 0) { mockRepository.clearInProgressRegistrationData() }
    coVerify { mockRepository.restoreFlowState() }
  }

  // ==================== Persistence Side-Effect Tests ====================

  @Test
  fun `onEvent ResetState clears flow state`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    viewModel.onEvent(RegistrationFlowEvent.ResetState)
    advanceUntilIdle()

    coVerify { mockRepository.clearFlowState() }
  }

  @Test
  fun `onEvent NavigateToScreen FullyComplete clears in-progress registration data`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    viewModel.onEvent(RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.FullyComplete))
    advanceUntilIdle()

    coVerify { mockRepository.clearInProgressRegistrationData() }
  }

  @Test
  fun `onEvent NavigateToScreen saves flow state`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    viewModel.onEvent(RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.PhoneNumberEntry))
    advanceUntilIdle()

    coVerify { mockRepository.saveFlowState(any()) }
  }

  @Test
  fun `onEvent Registered does not save flow state`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    viewModel.onEvent(RegistrationFlowEvent.Registered(AccountEntropyPool.generate(), storageCapable = false))
    advanceUntilIdle()

    coVerify(exactly = 0) { mockRepository.saveFlowState(any()) }
  }

  @Test
  fun `onEvent MasterKeyRestoredFromSvr does not save flow state`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    viewModel.onEvent(RegistrationFlowEvent.MasterKeyRestoredFromSvr(MasterKey(ByteArray(32))))
    advanceUntilIdle()

    coVerify(exactly = 0) { mockRepository.saveFlowState(any()) }
  }

  @Test
  fun `onEvent RegistrationComplete commits final data and clears in-progress registration data`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    viewModel.onEvent(RegistrationFlowEvent.RegistrationComplete)
    advanceUntilIdle()

    coVerify { mockRepository.commitFinalRegistrationData() }
    coVerify { mockRepository.clearInProgressRegistrationData() }
  }

  // ==================== applyEvent Tests (Navigation & State Reducers) ====================

  @Test
  fun `applyEvent NavigateToScreen appends to backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(backStack = listOf(RegistrationRoute.Welcome))

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.PhoneNumberEntry)
    )

    assertThat(result.backStack).isEqualTo(
      listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry)
    )
  }

  @Test
  fun `applyEvent NavigateBack pops last from backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry)
    )

    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBack)

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome))
  }

  @Test
  fun `applyEvent NavigateBack from recovery key entry clears pendingRestoreOption`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry, RegistrationRoute.EnterAepForRemoteBackupPreRegistration("+15551234567")),
      pendingRestoreOption = PendingRestoreOption.RemoteBackup
    )

    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBack)

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry))
    assertThat(result.pendingRestoreOption).isNull()
  }

  @Test
  fun `applyEvent NavigateBack from phone number entry clears pendingRestoreOption`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry),
      pendingRestoreOption = PendingRestoreOption.LocalBackup
    )

    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBack)

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome))
    assertThat(result.pendingRestoreOption).isNull()
  }

  @Test
  fun `applyEvent NavigateBack from local backup restore keeps pendingRestoreOption`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry, RegistrationRoute.LocalBackupRestore(isPreRegistration = true)),
      pendingRestoreOption = PendingRestoreOption.LocalBackup
    )

    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBack)

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry))
    assertThat(result.pendingRestoreOption).isEqualTo(PendingRestoreOption.LocalBackup)
  }

  @Test
  fun `applyEvent NavigateToScreen Welcome clears backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.PinEntryForRegistrationLock(
          timeRemaining = 1000L,
          svrCredentials = SvrCredentials(username = "user", password = "pass")
        ),
        RegistrationRoute.AccountLocked(timeRemainingMs = 1000L)
      )
    )

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.Welcome)
    )

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome))
  }

  @Test
  fun `applyEvent NavigateToScreen PinCreate clears backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry),
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.VerificationCodeEntry
      )
    )

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.PinCreate)
    )

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.PinCreate))
  }

  @Test
  fun `applyEvent NavigateToScreen PinEntryForSvrRestore clears backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.VerificationCodeEntry
      )
    )

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.PinEntryForSvrRestore)
    )

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.PinEntryForSvrRestore))
  }

  @Test
  fun `applyEvent NavigateToScreen RemoteRestore clears backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.VerificationCodeEntry
      )
    )

    val remoteRestore = RegistrationRoute.RemoteRestore(AccountEntropyPool.generate())

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(remoteRestore)
    )

    assertThat(result.backStack).isEqualTo(listOf(remoteRestore))
  }

  @Test
  fun `applyEvent NavigateToScreen post-registration ArchiveRestoreSelection clears backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.VerificationCodeEntry
      )
    )

    val postRegisterSelection = RegistrationRoute.ArchiveRestoreSelection.forPostRegisterWithPinKnown()

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(postRegisterSelection)
    )

    assertThat(result.backStack).isEqualTo(listOf(postRegisterSelection))
  }

  @Test
  fun `applyEvent NavigateToScreen pre-registration ArchiveRestoreSelection appends to backStack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(backStack = listOf(RegistrationRoute.Welcome))

    val preRegisterSelection = RegistrationRoute.ArchiveRestoreSelection.forManualRestore()

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(preRegisterSelection)
    )

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome, preRegisterSelection))
  }

  @Test
  fun `applyEvent NavigateToScreen with popCurrent replaces the current screen`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry, RegistrationRoute.LinkAccount())
    )

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry), popCurrent = true)
    )

    assertThat(result.backStack).isEqualTo(
      listOf(
        RegistrationRoute.Welcome,
        RegistrationRoute.PhoneNumberEntry,
        RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry)
      )
    )
  }

  @Test
  fun `applyEvent NavigateBackToScreen pops back to the target screen`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry, RegistrationRoute.LinkAccount())
    )

    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBackToScreen(RegistrationRoute.PhoneNumberEntry))

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry))
  }

  @Test
  fun `applyEvent NavigateBackToScreen replaces the current screen when the route is not on the back stack`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.LinkAccount())
    )

    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBackToScreen(RegistrationRoute.PhoneNumberEntry))

    assertThat(result.backStack).isEqualTo(
      listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry)
    )
  }

  @Test
  fun `applyEvent ResetState returns default state`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val populatedState = RegistrationFlowState(
      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry),
      sessionMetadata = createSessionMetadata(),
      sessionE164 = "+15551234567",
      doNotAttemptRecoveryPassword = true
    )

    val result = viewModel.applyEvent(populatedState, RegistrationFlowEvent.ResetState)

    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome))
    assertThat(result.sessionMetadata).isNull()
    assertThat(result.sessionE164).isNull()
    assertThat(result.doNotAttemptRecoveryPassword).isEqualTo(false)
    assertThat(result.isRestoringNavigationState).isEqualTo(false)
  }

  @Test
  fun `applyEvent SessionUpdated updates sessionMetadata`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val newSession = createSessionMetadata("new-session")

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.SessionUpdated(newSession)
    )

    assertThat(result.sessionMetadata).isEqualTo(newSession)
  }

  @Test
  fun `applyEvent VerificationCodeRequested updates both request windows`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.VerificationCodeRequested("+15551234567", nextSmsAllowedTimestamp = 12_345L, nextCallAllowedTimestamp = 23_456L)
    )

    assertThat(result.lastSmsVerificationCodeRequest).isEqualTo(VerificationCodeRequest("+15551234567", 12_345L))
    assertThat(result.lastCallVerificationCodeRequest).isEqualTo(VerificationCodeRequest("+15551234567", 23_456L))
  }

  @Test
  fun `applyEvent VerificationCodeRequested keeps the existing window when a timestamp is absent`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(
      lastSmsVerificationCodeRequest = VerificationCodeRequest("+15551234567", 12_345L),
      lastCallVerificationCodeRequest = VerificationCodeRequest("+15551234567", 23_456L)
    )

    val result = viewModel.applyEvent(
      initialState,
      RegistrationFlowEvent.VerificationCodeRequested("+15551234567", nextSmsAllowedTimestamp = 99_999L, nextCallAllowedTimestamp = null)
    )

    assertThat(result.lastSmsVerificationCodeRequest).isEqualTo(VerificationCodeRequest("+15551234567", 99_999L))
    assertThat(result.lastCallVerificationCodeRequest).isEqualTo(VerificationCodeRequest("+15551234567", 23_456L))
  }

  @Test
  fun `applyEvent E164Chosen updates sessionE164`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.E164Chosen("+15551234567")
    )

    assertThat(result.sessionE164).isEqualTo("+15551234567")
  }

  @Test
  fun `applyEvent Registered updates accountEntropyPool`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val aep = AccountEntropyPool.generate()

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.Registered(aep, storageCapable = true)
    )

    assertThat(result.accountEntropyPool).isEqualTo(aep)
    assertThat(result.storageCapable).isTrue()
  }

  @Test
  fun `applyEvent MasterKeyRestoredFromSvr updates temporaryMasterKey`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val masterKey = MasterKey(ByteArray(32))

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.MasterKeyRestoredFromSvr(masterKey)
    )

    assertThat(result.temporaryMasterKey).isEqualTo(masterKey)
  }

  @Test
  fun `applyEvent RecoveryPasswordInvalid sets doNotAttemptRecoveryPassword to true`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.RecoveryPasswordInvalid
    )

    assertThat(result.doNotAttemptRecoveryPassword).isTrue()
  }

  @Test
  fun `applyEvent PendingRestoreOptionSelected updates pendingRestoreOption`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.PendingRestoreOptionSelected(PendingRestoreOption.RemoteBackup)
    )

    assertThat(result.pendingRestoreOption).isEqualTo(PendingRestoreOption.RemoteBackup)
  }

  @Test
  fun `applyEvent UserSuppliedAepSubmitted updates unverifiedRestoredAep`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val aep = AccountEntropyPool.generate()

    val result = viewModel.applyEvent(
      RegistrationFlowState(),
      RegistrationFlowEvent.UserSuppliedAepSubmitted(aep)
    )

    assertThat(result.unverifiedRestoredAep).isEqualTo(aep)
  }

  @Test
  fun `applyEvent UserSuppliedAepVerified saves and updates accountEntropyPool`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val aep = AccountEntropyPool.generate()

    val result = viewModel.applyEvent(
      RegistrationFlowState(unverifiedRestoredAep = aep),
      RegistrationFlowEvent.UserSuppliedAepVerified(aep)
    )

    assertThat(result.accountEntropyPool).isEqualTo(aep)
    assertThat(result.unverifiedRestoredAep).isNull()
    coVerify { mockRepository.saveVerifiedUserSuppliedAep(aep) }
  }

  @Test
  fun `applyEvent RegistrationComplete commits data and navigates to FullyComplete`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    val initialState = RegistrationFlowState(backStack = listOf(RegistrationRoute.Welcome))

    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.RegistrationComplete)

    assertThat(result.backStack.last()).isEqualTo(RegistrationRoute.FullyComplete)
    coVerify { mockRepository.commitFinalRegistrationData() }
  }

  // ==================== preExistingRegistrationData Load Test ====================

  @Test
  fun `no saved state loads preExistingRegistrationData when present`() = runTest(testDispatcher) {
    val preExisting = mockk<PreExistingRegistrationData>(relaxed = true)
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns preExisting

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    assertThat(viewModel.state.value.preExistingRegistrationData).isEqualTo(preExisting)
  }

  // ==================== resultBus Tests ====================

  @Test
  fun `resultBus is initialized`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
    advanceUntilIdle()

    assertThat(viewModel.resultBus).isNotNull()
  }

  // ==================== Initial State Tests ====================

  @Test
  fun `initial state has isRestoringNavigationState true before init completes`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
    coEvery { mockRepository.getPreExistingRegistrationData() } returns null

    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())

    assertThat(viewModel.state.value.isRestoringNavigationState).isTrue()

    advanceUntilIdle()
  }

  // ==================== Helpers ====================

  private fun createSessionMetadata(id: String = "test-session"): SessionMetadata {
    return SessionMetadata(
      id = id,
      nextSms = 1000L,
      nextCall = 2000L,
      nextVerificationAttempt = 3000L,
      allowedToRequestCode = true,
      requestedInformation = emptyList(),
      verified = false
    )
  }
}
