MindMachine/app/src/main/java/com/mindmachine/mvp/session/MainViewModel.kt

package com.mindmachine.mvp.session

import androidx.lifecycle.ViewModel
import androidx.lifecycle.ViewModelProvider
import androidx.lifecycle.viewModelScope
import com.mindmachine.mvp.audio.BinauralAudioEngine
import com.mindmachine.mvp.audio.HeadsetMonitor
import com.mindmachine.mvp.data.SettingsRepository
import com.mindmachine.mvp.data.UserProgramEntity
import com.mindmachine.mvp.data.UserProgramRepository
import com.mindmachine.mvp.domain.AppSettings
import com.mindmachine.mvp.domain.Presets
import com.mindmachine.mvp.domain.RuntimeState
import com.mindmachine.mvp.domain.SessionConfig
import com.mindmachine.mvp.domain.SessionMode
import com.mindmachine.mvp.domain.SessionPreset
import com.mindmachine.mvp.domain.toConfig
import kotlinx.coroutines.Job
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.combine
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch
import java.util.UUID

const val SAFETY_VERSION = 1

data class UiState(
    val settings: AppSettings = AppSettings(),
    val presets: List<SessionPreset> = Presets.builtIn,
    val selectedPreset: SessionPreset = Presets.builtIn.first(),
    val config: SessionConfig = Presets.builtIn.first().toConfig(),

    // Timeline program (curves) used for both editing + runtime modulation.
    val timeline: TimelineEditorState = TimelineProgramFactory.fromPreset(Presets.builtIn.first()),

    val runtimeState: RuntimeState = RuntimeState.IDLE,
    val error: String? = null,
    val remainingSec: Int = 0,
    val countdownSec: Int = 0,
    val endedEarly: Boolean = false,
    val interruptionReason: String? = null,

    // Current runtime-evaluated parameters (updated while running).
    val runtimeParams: RuntimeParams = RuntimeParams(),
)

class MainViewModel(
    private val settingsRepository: SettingsRepository,
    private val userProgramRepository: UserProgramRepository,
    private val headsetMonitor: HeadsetMonitor,
    private val audioEngine: BinauralAudioEngine,
) : ViewModel() {
    private val _ui = MutableStateFlow(UiState())
    val ui: StateFlow<UiState> = _ui.asStateFlow()

    private var runJob: Job? = null
    private var elapsedBeforePauseSec: Float = 0f
    private var cachedPrograms: List<UserProgramEntity> = emptyList()

    init {
        viewModelScope.launch {
            combine(settingsRepository.settings, userProgramRepository.userPrograms) { settings, userPrograms ->
                settings to userPrograms
            }.collect { (settings, userPrograms) ->
                val builtInPrograms = Presets.builtIn.map { preset ->
                    userPrograms.firstOrNull { it.id == preset.id } ?: UserProgramEntity(
                        id = preset.id,
                        name = preset.name,
                        description = preset.description,
                        timeline = TimelineProgramFactory.fromPreset(preset),
                        sortOrder = preset.sortOrder,
                    )
                }
                val mergedPrograms = (builtInPrograms + userPrograms.filterNot { program -> Presets.builtIn.any { it.id == program.id } })
                    .sortedBy { it.sortOrder }
                cachedPrograms = mergedPrograms

                val allPresets = mergedPrograms.map {
                    TimelineProgramFactory.toPreset(
                        id = it.id,
                        name = it.name,
                        description = it.description,
                        timeline = it.timeline,
                        sortOrder = it.sortOrder,
                        isUserProgram = true,
                    )
                }

                _ui.update { current ->
                    val selectedId = settings.lastPresetId ?: current.selectedPreset.id
                    val selectedPreset = allPresets.find { it.id == selectedId } ?: allPresets.first()
                    val sameSelection = selectedPreset.id == current.selectedPreset.id
                    current.copy(
                        settings = settings,
                        presets = allPresets,
                        selectedPreset = selectedPreset,
                        config = current.config.copy(mode = settings.defaultModePreference),
                        timeline = if (sameSelection) current.timeline else timelineForProgramId(selectedPreset.id, mergedPrograms),
                    )
                }
            }
        }
    }

    fun acknowledgeSafety() = viewModelScope.launch { settingsRepository.acknowledgeSafety(SAFETY_VERSION) }

    fun choosePreset(id: String) = viewModelScope.launch {
        val state = _ui.value
        val preset = state.presets.first { it.id == id }
        settingsRepository.updateLastPreset(id)
        _ui.update {
            val cfg = preset.toConfig(it.settings.defaultModePreference)
            it.copy(
                selectedPreset = preset,
                config = cfg.copy(durationSec = cfg.durationSec.coerceIn(60, 8 * 60 * 60)),
                timeline = timelineForProgramId(preset.id, cachedPrograms),
                error = null,
            )
        }
        elapsedBeforePauseSec = 0f
    }

    fun createNewProgramDraft() {
        val draftId = "user-${UUID.randomUUID()}"
        val displayName = nextUntitledName()
        val timeline = TimelineEditorState.default(20 * 60)
        val preset = TimelineProgramFactory.toPreset(
            id = draftId,
            name = displayName,
            description = "Custom program",
            timeline = timeline,
            sortOrder = nextSortOrder(),
            isUserProgram = true,
        )
        _ui.update {
            it.copy(
                selectedPreset = preset,
                config = preset.toConfig(it.settings.defaultModePreference),
                timeline = timeline,
                error = null,
            )
        }
    }

    fun saveCurrentProgram(name: String) = viewModelScope.launch {
        val trimmed = name.trim()
        if (trimmed.isBlank()) {
            _ui.update { it.copy(error = "Enter a name before saving.") }
            return@launch
        }

        val state = _ui.value
        val id = state.selectedPreset.id
        val sortOrder = cachedPrograms.find { it.id == id }?.sortOrder ?: nextSortOrder()

        userProgramRepository.upsert(
            UserProgramEntity(
                id = id,
                name = trimmed,
                description = "Custom program",
                timeline = state.timeline,
                sortOrder = sortOrder,
            )
        )
        settingsRepository.updateLastPreset(id)

        _ui.update {
            val savedPreset = TimelineProgramFactory.toPreset(
                id = id,
                name = trimmed,
                description = "Custom program",
                timeline = it.timeline,
                sortOrder = sortOrder,
                isUserProgram = true,
            )
            it.copy(
                selectedPreset = savedPreset,
                error = "Saved \"$trimmed\"",
            )
        }
    }

    fun deleteProgram(id: String) = viewModelScope.launch {
        userProgramRepository.delete(id)
        if (_ui.value.selectedPreset.id == id) {
            val fallback = _ui.value.presets.firstOrNull { it.id != id } ?: Presets.builtIn.first()
            settingsRepository.updateLastPreset(fallback.id)
            _ui.update {
                it.copy(
                    selectedPreset = fallback,
                    config = fallback.toConfig(it.settings.defaultModePreference),
                    timeline = timelineForProgramId(fallback.id, cachedPrograms),
                    error = null,
                )
            }
        }
    }

    fun setMode(mode: SessionMode) = _ui.update { it.copy(config = it.config.copy(mode = mode), error = null) }

    fun setDurationSec(seconds: Int) = _ui.update {
        val clamped = seconds.coerceIn(60, 8 * 60 * 60)
        it.copy(
            config = it.config.copy(durationSec = clamped),
            timeline = it.timeline.setDurationSec(clamped),
            error = null,
        )
    }

    fun setCurveGranularitySec(seconds: Int) = _ui.update {
        it.copy(
            timeline = it.timeline.setCurveGranularitySec(seconds),
            error = null,
        )
    }

    fun updateTimeline(newState: TimelineEditorState) = _ui.update {
        it.copy(
            timeline = newState,
            config = it.config.copy(durationSec = newState.durationSec.coerceIn(60, 8 * 60 * 60)),
            error = null,
        )
    }

    fun setFlashIntervalMs(value: Int) = _ui.update { it.copy(config = it.config.copy(flashIntervalMs = value.coerceIn(50, 2000)), error = null) }
    fun setCarrier(value: Float) = _ui.update { it.copy(config = it.config.copy(carrierFrequencyHz = value.coerceIn(80f, 400f)), error = null) }
    fun setDifference(value: Float) = _ui.update { it.copy(config = it.config.copy(binauralDifferenceHz = value.coerceIn(0.5f, 20f)), error = null) }

    fun updateCountdownSeconds(seconds: Int) = viewModelScope.launch { settingsRepository.updateCountdownSeconds(seconds) }
    fun updateDefaultMode(mode: SessionMode) = viewModelScope.launch { settingsRepository.updateDefaultMode(mode) }
    fun updateGuidance(value: Boolean) = viewModelScope.launch { settingsRepository.updateGuidance(value) }
    fun updateShowImmersiveProgressBar(value: Boolean) = viewModelScope.launch { settingsRepository.updateShowImmersiveProgressBar(value) }

    fun startSession(): Boolean {
        val state = _ui.value
        if (!(state.settings.safetyAcknowledged && state.settings.safetyAcknowledgedVersion >= SAFETY_VERSION)) {
            _ui.update { it.copy(error = "You must acknowledge safety before starting sessions.") }
            return false
        }
        val headset = headsetMonitor.isStereoHeadsetAvailable()
        val validation = SessionValidator.validate(state.config, headset)
        if (validation != null) {
            _ui.update { it.copy(error = validation) }
            return false
        }

        elapsedBeforePauseSec = 0f
        runProgram(startElapsedSec = 0f, includeCountdown = true)
        return true
    }

    fun pause(reason: String? = null) {
        if (_ui.value.runtimeState != RuntimeState.RUNNING) return
        elapsedBeforePauseSec = currentElapsedSec()
        runJob?.cancel()
        audioEngine.stop()
        _ui.update {
            it.copy(
                runtimeState = if (reason == null) RuntimeState.PAUSED else RuntimeState.INTERRUPTED,
                interruptionReason = reason,
            )
        }
    }

    fun resume() {
        val state = _ui.value
        if (state.runtimeState != RuntimeState.PAUSED && state.runtimeState != RuntimeState.INTERRUPTED) return
        if (state.config.mode != SessionMode.VISUAL_ONLY && !headsetMonitor.isStereoHeadsetAvailable()) {
            _ui.update { it.copy(error = "Headphones are required to resume audio mode.") }
            return
        }
        runProgram(startElapsedSec = elapsedBeforePauseSec, includeCountdown = true)
    }

    fun stop() {
        runJob?.cancel()
        audioEngine.stop()
        elapsedBeforePauseSec = 0f
        _ui.update { it.copy(runtimeState = RuntimeState.STOPPED, endedEarly = true) }
    }

    fun switchToVisualOnlyAndResume() {
        _ui.update { it.copy(config = it.config.copy(mode = SessionMode.VISUAL_ONLY), error = null) }
        resume()
    }

    override fun onCleared() {
        audioEngine.stop()
        super.onCleared()
    }

    private fun runProgram(startElapsedSec: Float, includeCountdown: Boolean) {
        runJob?.cancel()
        runJob = viewModelScope.launch {
            val program = _ui.value.timeline
            val durationSec = program.durationSec
            val clampedStartSec = startElapsedSec.coerceIn(0f, durationSec.toFloat())

            if (includeCountdown) {
                val count = _ui.value.settings.countdownSeconds.coerceIn(0, 10)
                if (count > 0) {
                    for (i in count downTo 1) {
                        _ui.update {
                            it.copy(
                                runtimeState = RuntimeState.COUNTDOWN,
                                countdownSec = i,
                                remainingSec = (durationSec - clampedStartSec.toInt()).coerceAtLeast(0),
                                endedEarly = false,
                                interruptionReason = null,
                                runtimeParams = TimelineRuntimeEvaluator.evaluate(program, clampedStartSec),
                            )
                        }
                        delay(1000)
                    }
                }
            }

            val initialParams = TimelineRuntimeEvaluator.evaluate(program, clampedStartSec)
            _ui.update {
                it.copy(
                    runtimeState = RuntimeState.RUNNING,
                    remainingSec = (durationSec - clampedStartSec.toInt()).coerceAtLeast(0),
                    countdownSec = 0,
                    error = null,
                    endedEarly = false,
                    interruptionReason = null,
                    runtimeParams = initialParams,
                )
            }

            if (_ui.value.config.mode != SessionMode.VISUAL_ONLY) {
                audioEngine.start(initialParams.carrierHz, initialParams.binauralHz)
            }

            val tStart = System.currentTimeMillis()
            while (true) {
                delay(50)

                val elapsedSec = clampedStartSec + (System.currentTimeMillis() - tStart) / 1000f
                elapsedBeforePauseSec = elapsedSec.coerceAtMost(durationSec.toFloat())
                val remaining = (durationSec - kotlin.math.ceil(elapsedBeforePauseSec).toInt()).coerceAtLeast(0)
                val currentUi = _ui.value

                if (currentUi.config.mode != SessionMode.VISUAL_ONLY && !headsetMonitor.isStereoHeadsetAvailable()) {
                    audioEngine.stop()
                    _ui.update {
                        it.copy(
                            runtimeState = RuntimeState.INTERRUPTED,
                            interruptionReason = "Headphones disconnected. Session paused.",
                        )
                    }
                    return@launch
                }

                val params = TimelineRuntimeEvaluator.evaluate(program, elapsedBeforePauseSec)
                _ui.update { it.copy(runtimeParams = params, remainingSec = remaining) }

                if (currentUi.config.mode != SessionMode.VISUAL_ONLY) {
                    audioEngine.setFrequencies(params.carrierHz, params.binauralHz)
                }

                if (elapsedBeforePauseSec >= durationSec.toFloat()) break
            }

            audioEngine.stop()
            elapsedBeforePauseSec = 0f
            _ui.update { it.copy(runtimeState = RuntimeState.COMPLETED, endedEarly = false, remainingSec = 0) }
        }
    }

    private fun currentElapsedSec(): Float {
        val state = _ui.value
        val total = state.timeline.durationSec.toFloat().coerceAtLeast(1f)
        return (total - state.remainingSec).coerceIn(0f, total)
    }

    private fun timelineForProgramId(programId: String, programs: List<UserProgramEntity>): TimelineEditorState {
        val fromSaved = programs.firstOrNull { it.id == programId }?.timeline
        if (fromSaved != null) return fromSaved

        val builtInPreset = Presets.builtIn.firstOrNull { it.id == programId }
        return if (builtInPreset != null) TimelineProgramFactory.fromPreset(builtInPreset) else TimelineEditorState.default()
    }

    private fun nextSortOrder(): Int = (cachedPrograms.maxOfOrNull { it.sortOrder } ?: 999) + 1

    private fun nextUntitledName(): String {
        val names = cachedPrograms.map { it.name.lowercase() }.toSet()
        var i = 1
        while (true) {
            val candidate = "Custom $i"
            if (candidate.lowercase() !in names) return candidate
            i++
        }
    }

    class Factory(
        private val settingsRepository: SettingsRepository,
        private val userProgramRepository: UserProgramRepository,
        private val headsetMonitor: HeadsetMonitor,
        private val audioEngine: BinauralAudioEngine,
    ) : ViewModelProvider.Factory {
        override fun <T : ViewModel> create(modelClass: Class<T>): T =
            MainViewModel(settingsRepository, userProgramRepository, headsetMonitor, audioEngine) as T
    }
}