AssistViewModel.kt

package io.homeassistant.companion.android.assist

import android.app.Application
import android.content.Intent
import androidx.annotation.VisibleForTesting
import androidx.compose.runtime.derivedStateOf
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateListOf
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.setValue
import androidx.lifecycle.viewModelScope
import dagger.assisted.Assisted
import dagger.assisted.AssistedFactory
import dagger.assisted.AssistedInject
import dagger.hilt.android.lifecycle.HiltViewModel
import io.homeassistant.companion.android.assist.ui.AssistMessage
import io.homeassistant.companion.android.assist.ui.AssistUiPipeline
import io.homeassistant.companion.android.common.R as commonR
import io.homeassistant.companion.android.common.assist.AssistAudioStrategy
import io.homeassistant.companion.android.common.assist.AssistEvent
import io.homeassistant.companion.android.common.assist.AssistViewModelBase
import io.homeassistant.companion.android.common.data.servers.ServerManager
import io.homeassistant.companion.android.common.data.websocket.impl.entities.AssistPipelineResponse
import io.homeassistant.companion.android.common.util.AudioUrlPlayer
import kotlin.time.Duration.Companion.seconds
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Job
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import timber.log.Timber

@VisibleForTesting
internal val CLOSE_INACTIVE = 30.seconds

@HiltViewModel(assistedFactory = AssistViewModel.Factory::class)
class AssistViewModel @AssistedInject constructor(
    serverManager: ServerManager,
    @Assisted initialAudioStrategy: AssistAudioStrategy,
    audioUrlPlayer: AudioUrlPlayer,
    application: Application,
) : AssistViewModelBase(serverManager, initialAudioStrategy, audioUrlPlayer, application) {

    @AssistedFactory
    interface Factory {
        fun create(audioStrategy: AssistAudioStrategy): AssistViewModel
    }

    init {
        viewModelScope.launch {
            audioStrategy.wakeWordDetected.collect { detectedPhrase ->
                if (inputMode != AssistInputMode.VOICE_ACTIVE) {
                    wakeWordPhrase = detectedPhrase
                    onMicrophoneInput()
                }
            }
        }
    }

    private var filteredServerId: Int? = null
    private val allPipelines = mutableMapOf<Int, List<AssistPipelineResponse>>()
    private var selectedPipeline: AssistPipelineResponse? = null
    private var wakeWordPhrase: String? = null

    private var recorderAutoStart = true
    private var requestPermission: (() -> Unit)? = null
    private var requestSilently = true

    private val startMessage =
        AssistMessage(application.getString(commonR.string.assist_how_can_i_assist), isInput = false)
    private val _conversation = mutableStateListOf(startMessage)

    /**
     * Conversation messages to display. The input placeholder marking an active voice recording is
     * hidden: listening is already conveyed by the microphone button animation.
     */
    val conversation: List<AssistMessage> by derivedStateOf {
        _conversation.filterNot { it.isInput && it.isPlaceholder }
    }

    private val _pipelines = mutableStateListOf<AssistUiPipeline>()
    val pipelines: List<AssistUiPipeline> = _pipelines

    var currentPipeline by mutableStateOf<AssistUiPipeline?>(null)
        private set

    var inputMode by mutableStateOf<AssistInputMode?>(null)
        private set

    var userCanManagePipelines by mutableStateOf(false)
        private set

    var shouldFinish by mutableStateOf(false)
        private set

    /** When true, the UI should not be shown yet (waiting to confirm the pipeline is not a duplicate wake-up). The activity should finish if a duplicate is detected. */
    var pendingWakeWordConfirmation by mutableStateOf(false)
        private set

    private var startedFromWakeWord = false
    private var inactivityTimerJob: Job? = null

    fun onCreate(
        hasPermission: Boolean,
        serverId: Int?,
        pipelineId: String?,
        startListening: Boolean?,
        wakeWordPhrase: String?,
    ) {
        viewModelScope.launch {
            this@AssistViewModel.hasPermission = hasPermission
            this@AssistViewModel.wakeWordPhrase = wakeWordPhrase
            this@AssistViewModel.startedFromWakeWord = wakeWordPhrase != null
            serverId?.let {
                filteredServerId = serverId
                selectedServerId = serverId
            }
            startListening?.let { recorderAutoStart = it }

            if (!serverManager.isRegistered()) {
                inputMode = AssistInputMode.BLOCKED
                _conversation.clear()
                _conversation.add(
                    AssistMessage(app.getString(commonR.string.not_registered), isInput = false),
                )
                return@launch
            }

            if (
                pipelineId == PIPELINE_LAST_USED &&
                recorderAutoStart &&
                hasPermission &&
                hasMicrophone &&
                serverManager.getServer(selectedServerId) != null &&
                serverManager.integrationRepository(selectedServerId).getLastUsedPipelineSttSupport()
            ) {
                // Start microphone recording to prevent missing voice input while doing network checks
                pendingWakeWordConfirmation = wakeWordPhrase != null
                onMicrophoneInput(proactive = true)
            }

            val supported = checkSupport()
            if (supported != true) stopRecording()
            if (supported == null) { // Couldn't get config
                inputMode = AssistInputMode.BLOCKED
                _conversation.clear()
                _conversation.add(
                    AssistMessage(app.getString(commonR.string.assist_connnect), isInput = false),
                )
            } else if (!supported) { // Core too old or doesn't include assist pipeline
                inputMode = AssistInputMode.BLOCKED
                _conversation.clear()
                _conversation.add(
                    AssistMessage(
                        app.getString(
                            commonR.string.no_assist_support,
                            "2023.5",
                            app.getString(commonR.string.no_assist_support_assist_pipeline),
                        ),
                        isInput = false,
                    ),
                )
            } else {
                setPipeline(
                    when {
                        pipelineId == PIPELINE_LAST_USED -> serverManager.integrationRepository(
                            selectedServerId,
                        ).getLastUsedPipelineId()

                        pipelineId == PIPELINE_PREFERRED -> null
                        pipelineId?.isNotBlank() == true -> pipelineId
                        else -> null
                    },
                )
            }

            if (serverManager.isRegistered()) {
                loadPipelines()
            }
            userCanManagePipelines = serverManager.getServer()?.user?.isAdmin == true
        }
    }

    /**
     * Update the state of the Assist dialog for a new 'assistant triggered' action
     * @param intent the updated intent
     * @param lockedMatches whether the locked state changed and contents should be cleared
     */
    fun onNewIntent(intent: Intent, lockedMatches: Boolean) {
        if (
            (intent.flags and Intent.FLAG_ACTIVITY_BROUGHT_TO_FRONT != 0) ||
            intent.action in
            listOf(Intent.ACTION_ASSIST, "android.intent.action.VOICE_ASSIST", Intent.ACTION_VOICE_COMMAND)
        ) {
            if (!lockedMatches && inputMode != AssistInputMode.BLOCKED) {
                _conversation.clear()
                _conversation.add(startMessage)
            }
            if (inputMode == AssistInputMode.VOICE_ACTIVE || inputMode == AssistInputMode.VOICE_INACTIVE) {
                onMicrophoneInput()
            }
        }
    }

    override fun getInput(): AssistInputMode? = inputMode

    override fun setInput(inputMode: AssistInputMode) {
        this.inputMode = inputMode
        restartInactivityTimer()
    }

    /**
     * Restarts the inactivity timer that closes the Assist dialog after [CLOSE_INACTIVE].
     *
     * The timer only runs when the input mode is [AssistInputMode.VOICE_INACTIVE],
     * TTS audio is not currently playing, and the last conversation message is not a
     * placeholder (assistant finished processing).
     *
     * Disabled for non wake word sessions to avoid auto-closing the dialog while the user
     * is interacting with hands.
     */
    private fun restartInactivityTimer() {
        inactivityTimerJob?.cancel()
        if (!startedFromWakeWord) return
        fun isInactive(): Boolean {
            val shouldRun = when (inputMode) {
                AssistInputMode.VOICE_INACTIVE -> true
                AssistInputMode.TEXT,
                AssistInputMode.TEXT_ONLY,
                AssistInputMode.VOICE_ACTIVE,
                AssistInputMode.BLOCKED,
                null,
                -> false
            }
            if (!shouldRun || isPlayingAudio) return false

            val lastMessage = _conversation.lastOrNull()
            return !(lastMessage == null || lastMessage.isPlaceholder)
        }

        if (!isInactive()) return

        inactivityTimerJob = viewModelScope.launch {
            delay(CLOSE_INACTIVE)
            if (isInactive()) {
                shouldFinish = true
            } else {
                Timber.d("Inactivity timer expired but Assist is no longer inactive, not finishing")
            }
        }
    }

    private suspend fun checkSupport(): Boolean? {
        if (!serverManager.isRegistered()) return false
        return try {
            if (!serverManager.integrationRepository(
                    selectedServerId,
                ).isHomeAssistantVersionAtLeast(2023, 5, 0)
            ) {
                false
            } else {
                serverManager.webSocketRepository(selectedServerId).getConfig()?.components?.contains("assist_pipeline")
            }
        } catch (e: IllegalStateException) {
            Timber.e(e, "Failed to check support")
            false
        }
    }

    private suspend fun loadPipelines() {
        val serverIds = filteredServerId?.let { listOf(it) } ?: serverManager.servers().map { it.id }
        serverIds.forEach { serverId ->
            viewModelScope.launch {
                val server = serverManager.getServer(serverId)
                val serverPipelines = try {
                    serverManager.webSocketRepository(serverId).getAssistPipelines()
                } catch (e: CancellationException) {
                    throw e
                } catch (e: Exception) {
                    Timber.e(e, "Failed to load assist pipelines")
                    null
                }

                allPipelines[serverId] = serverPipelines?.pipelines ?: emptyList()
                _pipelines.addAll(
                    serverPipelines?.pipelines.orEmpty().map {
                        AssistUiPipeline(
                            serverId = serverId,
                            serverName = server?.friendlyName ?: "",
                            id = it.id,
                            name = it.name,
                        )
                    },
                )
            }
        }
    }

    fun changePipeline(serverId: Int, id: String) = viewModelScope.launch {
        if (serverId == selectedServerId && id == selectedPipeline?.id) return@launch

        stopRecording(sendRecorded = false)
        stopPlayback()

        selectedServerId = serverId
        setPipeline(id)
    }

    private suspend fun setPipeline(id: String?) {
        selectedPipeline =
            allPipelines[selectedServerId]?.firstOrNull { it.id == id }
                ?: try {
                    serverManager.webSocketRepository(selectedServerId).getAssistPipeline(id)
                } catch (e: CancellationException) {
                    throw e
                } catch (e: Exception) {
                    Timber.e(e, "Failed to get assist pipeline")
                    null
                }
        selectedPipeline?.let {
            currentPipeline = AssistUiPipeline(
                serverId = selectedServerId,
                serverName = serverManager.getServer(selectedServerId)?.friendlyName ?: "",
                id = it.id,
                name = it.name,
            )
            try {
                serverManager.integrationRepository(selectedServerId).setLastUsedPipeline(it.id, it.sttEngine != null)
            } catch (e: CancellationException) {
                throw e
            } catch (e: Exception) {
                Timber.e(e, "Failed to set last used pipeline")
            }

            _conversation.clear()
            _conversation.add(startMessage)
            clearPipelineData()
            if (hasMicrophone && it.sttEngine != null) {
                if (recorderAutoStart && (hasPermission || requestSilently)) {
                    inputMode = AssistInputMode.VOICE_INACTIVE
                    onMicrophoneInput(proactive = null)
                } else { // already requested permission once and was denied
                    inputMode = AssistInputMode.TEXT
                }
            } else {
                inputMode = AssistInputMode.TEXT_ONLY
            }
            restartInactivityTimer()
        } ?: run {
            if (!id.isNullOrBlank()) {
                setPipeline(null) // Try falling back to default pipeline
            } else {
                Timber.w("Server $selectedServerId does not have any pipelines")
                inputMode = AssistInputMode.BLOCKED
                _conversation.clear()
                _conversation.add(
                    AssistMessage(app.getString(commonR.string.assist_error), isInput = false),
                )
            }
        }
    }

    fun onChangeInput() {
        when (inputMode) {
            null, AssistInputMode.BLOCKED, AssistInputMode.TEXT_ONLY -> { /* Do nothing */ }

            AssistInputMode.TEXT -> {
                inputMode = AssistInputMode.VOICE_INACTIVE
                if (hasPermission || requestSilently) {
                    onMicrophoneInput()
                }
            }

            AssistInputMode.VOICE_INACTIVE -> {
                inputMode = AssistInputMode.TEXT
            }

            AssistInputMode.VOICE_ACTIVE -> {
                stopRecording(sendRecorded = false)
                // Remove placeholder message if present from proactive recording
                if (_conversation.lastOrNull()?.let { it.isPlaceholder && it.isInput } == true) {
                    _conversation.removeAt(_conversation.size - 1)
                }
                inputMode = AssistInputMode.TEXT
            }
        }
        restartInactivityTimer()
    }

    fun onTextInput(input: String) = runAssistPipeline(input)

    /**
     * Start/stop microphone input for Assist, depending on the current state.
     * @param proactive true if proactive, null if not important, false if not
     */
    fun onMicrophoneInput(proactive: Boolean? = false) {
        if (!hasPermission) {
            requestPermission?.let { it() }
            return
        }

        if (inputMode == AssistInputMode.VOICE_ACTIVE && proactive == false) {
            stopRecording()
            return
        }

        stopPlayback()

        if (!recorderProactive) {
            audioStrategy.requestFocus()
            setupRecorder(
                onError = {
                    stopRecording()
                    _conversation.add(
                        AssistMessage(app.getString(commonR.string.assist_error), isInput = false, isError = true),
                    )
                },
            )
        }
        inputMode = AssistInputMode.VOICE_ACTIVE
        if (proactive == true) _conversation.add(AssistMessage.placeholder(isInput = true))
        if (proactive != true) runAssistPipeline(null)

        restartInactivityTimer()

        recorderProactive = proactive == true
    }

    private fun runAssistPipeline(text: String?) {
        val isVoice = text == null
        stopPlayback()

        val userMessage = text?.let { AssistMessage(it, isInput = true) } ?: AssistMessage.placeholder(isInput = true)
        _conversation.add(userMessage)
        val haMessage = AssistMessage.placeholder(isInput = false)
        if (!isVoice) _conversation.add(haMessage)
        var message = if (isVoice) userMessage else haMessage

        // Capture and clear wake word phrase - it should only be sent once for the initial command
        val wakeWord = wakeWordPhrase.also { wakeWordPhrase = null }

        runAssistPipelineInternal(
            text = text,
            pipeline = selectedPipeline,
            wakeWordPhrase = wakeWord,
        ) { event ->
            when (event) {
                is AssistEvent.Message -> {
                    _conversation.indexOf(message).takeIf { pos -> pos >= 0 }?.let { index ->
                        val isInput = event is AssistEvent.Message.Input
                        val isError = event is AssistEvent.Message.Error
                        _conversation[index] = message.copy(
                            message = event.message.trim(),
                            isInput = isInput,
                            isError = isError,
                        )
                        if (isInput) {
                            _conversation.add(haMessage)
                            message = haMessage
                        }
                        if (isError && inputMode == AssistInputMode.VOICE_ACTIVE) {
                            stopRecording()
                        }
                    }
                    restartInactivityTimer()
                }

                is AssistEvent.MessageChunk -> {
                    val lastMessage = _conversation.last()
                    if (lastMessage == haMessage) {
                        // Remove '...' message and add the chunk received
                        _conversation.removeAt(_conversation.lastIndex)
                        _conversation.add(lastMessage.copy(message = event.chunk))
                    } else {
                        // Replace last message with the updated message with the new chunk append
                        _conversation[_conversation.lastIndex] =
                            lastMessage.copy(message = lastMessage.message + event.chunk)
                    }
                }

                is AssistEvent.PipelineStarted -> { /* handled below */ }

                is AssistEvent.PipelineEnded,
                is AssistEvent.PlaybackFinished,
                -> restartInactivityTimer()

                is AssistEvent.ContinueConversation -> onMicrophoneInput()
                is AssistEvent.Dismiss -> shouldFinish = true
            }
            if (!shouldFinish && pendingWakeWordConfirmation) {
                // Any event confirms this is not a duplicate wake-up, so the UI can be shown.
                // Skipped when finishing to avoid a brief UI flash before dismissal.
                pendingWakeWordConfirmation = false
            }
        }
    }

    fun setPermissionInfo(hasPermission: Boolean, callback: () -> Unit) {
        this.hasPermission = hasPermission
        requestPermission = callback
    }

    fun onPermissionResult(granted: Boolean) {
        hasPermission = granted
        val proactive = currentPipeline == null
        if (granted) {
            inputMode = AssistInputMode.VOICE_INACTIVE
            onMicrophoneInput(proactive = proactive)
        } else if (requestSilently && !proactive) { // Don't notify the user if they haven't explicitly requested
            inputMode = AssistInputMode.TEXT
        } else if (!requestSilently) {
            _conversation.add(AssistMessage(app.getString(commonR.string.assist_permission), isInput = false))
        }
        restartInactivityTimer()
        if (!proactive) requestSilently = false
    }

    fun onPause() {
        requestPermission = null
        inactivityTimerJob?.cancel()
        stopRecording()
    }

    fun onDestroy() {
        requestPermission = null
        inactivityTimerJob?.cancel()
        stopRecording()
        stopPlayback()
    }
}