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.domain.AppSettings
import com.mindmachine.mvp.domain.CountdownPreference
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.update
import kotlinx.coroutines.launch

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 headsetMonitor: HeadsetMonitor,
    private val audioEngine: BinauralAudioEngine,
) : ViewModel() {
    private val _ui = MutableStateFlow(UiState())
    val ui: StateFlow<UiState> = _ui.asStateFlow()

    private var runJob: Job? = null

    init {
        viewModelScope.launch {
            settingsRepository.settings.collect { s ->
                _ui.update { current ->
                    val preset = current.presets.find { it.id == (s.lastPresetId ?: current.selectedPreset.id) } ?: current.selectedPreset
                    current.copy(settings = s, selectedPreset = preset, config = current.config.copy(mode = s.defaultModePreference))
                }
            }
        }
    }

    fun acknowledgeSafety() = viewModelScope.launch { settingsRepository.acknowledgeSafety(SAFETY_VERSION) }
    fun choosePreset(id: String) = viewModelScope.launch {
        val preset = _ui.value.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 = TimelineProgramFactory.fromPreset(preset),
                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 updateTimeline(newState: TimelineEditorState) = _ui.update {
        it.copy(
            timeline = newState,
            config = it.config.copy(durationSec = newState.durationSec.coerceIn(60, 8 * 60 * 60)),
            error = null,
        )
    }

    // Legacy fixed-parameter setters kept for now but no longer used by Setup.
    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 updateCountdown(pref: CountdownPreference) = viewModelScope.launch { settingsRepository.updateCountdown(pref) }
    fun updateDefaultMode(mode: SessionMode) = viewModelScope.launch { settingsRepository.updateDefaultMode(mode) }
    fun updateGuidance(value: Boolean) = viewModelScope.launch { settingsRepository.updateGuidance(value) }

    fun startSession() {
        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
        }
        val headset = headsetMonitor.isStereoHeadsetAvailable()
        val validation = SessionValidator.validate(state.config, headset)
        if (validation != null) {
            _ui.update { it.copy(error = validation) }
            return
        }
        runJob?.cancel()
        runJob = viewModelScope.launch {
            val count = state.settings.countdownPreference.seconds
            if (count > 0) {
                for (i in count downTo 1) {
                    _ui.update { it.copy(runtimeState = RuntimeState.COUNTDOWN, countdownSec = i, remainingSec = state.config.durationSec, endedEarly = false, interruptionReason = null) }
                    delay(1000)
                }
            }
            // Use the timeline curves as the actual runtime program.
            val program = _ui.value.timeline
            val durationSec = program.durationSec

            _ui.update {
                it.copy(
                    runtimeState = RuntimeState.RUNNING,
                    remainingSec = durationSec,
                    countdownSec = 0,
                    error = null,
                    runtimeParams = TimelineRuntimeEvaluator.evaluate(program, 0f),
                )
            }

            if (state.config.mode != SessionMode.VISUAL_ONLY) {
                val p0 = TimelineRuntimeEvaluator.evaluate(program, 0f)
                audioEngine.start(p0.carrierHz, p0.binauralHz)
            }

            val tStart = System.currentTimeMillis()
            var lastWholeSec = durationSec

            while (true) {
                delay(50)

                val elapsedSec = (System.currentTimeMillis() - tStart) / 1000f
                val remaining = (durationSec - elapsedSec.toInt()).coerceAtLeast(0)

                if (state.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, elapsedSec)
                _ui.update {
                    val nextRemaining = if (remaining != lastWholeSec) remaining else it.remainingSec
                    it.copy(runtimeParams = params, remainingSec = nextRemaining)
                }
                lastWholeSec = remaining

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

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

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

    fun pause(reason: String? = null) {
        if (_ui.value.runtimeState != RuntimeState.RUNNING) return
        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
        }
        runJob = viewModelScope.launch {
            for (i in 3 downTo 1) {
                _ui.update { it.copy(runtimeState = RuntimeState.COUNTDOWN, countdownSec = i) }
                delay(1000)
            }
            _ui.update { it.copy(runtimeState = RuntimeState.RUNNING, countdownSec = 0) }
            if (state.config.mode != SessionMode.VISUAL_ONLY) {
                val p = _ui.value.runtimeParams
                audioEngine.start(p.carrierHz, p.binauralHz)
            }
            var remaining = state.remainingSec
            while (remaining > 0) {
                delay(1000)
                if (state.config.mode != SessionMode.VISUAL_ONLY && !headsetMonitor.isStereoHeadsetAvailable()) {
                    audioEngine.stop()
                    _ui.update { it.copy(runtimeState = RuntimeState.INTERRUPTED, interruptionReason = "Headphones disconnected. Session paused.") }
                    return@launch
                }
                remaining -= 1
                _ui.update { it.copy(remainingSec = remaining) }
            }
            audioEngine.stop()
            _ui.update { it.copy(runtimeState = RuntimeState.COMPLETED, endedEarly = false) }
        }
    }

    fun stop() {
        runJob?.cancel()
        audioEngine.stop()
        _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()
    }

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