package dev.dtrentin.chart

import androidx.compose.runtime.mutableStateOf
import dev.dtrentin.chart.buffer.TieredBuffer
import dev.dtrentin.chart.model.ChartConfig
import dev.dtrentin.chart.model.SignalConfig
import dev.dtrentin.chart.model.T0
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.launch

/**
 * Holds all state for a [RealtimeChart]. Create once, pass to composable.
 *
 * Data ingestion is thread-safe. Compose reads on main thread.
 *
 * ```
 * val state = remember { RealtimeChartState(ChartConfig()) }
 * state.addSignal("v1", SignalConfig(color = Color.Red))
 * RealtimeChart(state, Modifier.fillMaxSize())
 * ```
 */
class RealtimeChartState(
    val config: ChartConfig = ChartConfig(),
) {
    private val _signals = mutableStateOf<Map<String, SignalEntry>>(emptyMap())
    internal val signals: Map<String, SignalEntry> get() = _signals.value

    @Volatile internal var resolvedT0Ms: Long? = when (val t = config.t0) {
        is T0.Fixed -> t.epochMs
        T0.FirstSample -> null
    }

    /** Registers a signal. Must call before [push]. If [name] exists, buffer is replaced. */
    fun addSignal(name: String, signalConfig: SignalConfig) {
        _signals.value = _signals.value + (name to SignalEntry(signalConfig, TieredBuffer()))
    }

    /** Removes signal [name] and its buffer. No-op if absent. */
    fun removeSignal(name: String) {
        _signals.value = _signals.value - name
    }

    /** Pushes one sample. Thread-safe. @param timestampMs ms epoch. */
    fun push(name: String, timestampMs: Long, value: Float) {
        if (resolvedT0Ms == null) resolvedT0Ms = timestampMs
        _signals.value[name]?.buffer?.push(timestampMs, value)
    }

    /** Collects (timestampMs, value) pairs from [flow] into signal [name]. Returns [Job]. */
    @JvmName("collectFromTimestamped")
    fun collectFrom(name: String, flow: Flow<Pair<Long, Float>>, scope: CoroutineScope): Job =
        scope.launch { flow.collect { (ts, v) -> push(name, ts, v) } }

    /** Collects raw [Float] from [flow], stamps with current time. Returns [Job]. */
    @JvmName("collectFromFloat")
    fun collectFrom(name: String, flow: Flow<Float>, scope: CoroutineScope): Job =
        scope.launch { flow.collect { v -> push(name, currentTimeMs(), v) } }
}

internal class SignalEntry(
    val config: SignalConfig,
    val buffer: TieredBuffer,
)