Arcology Engine

Habit Screen and ViewModel

Contents

The HabitScreen is the second tab in the Tasks bottom-nav destination. It lists every heading with a :STYLE: habit property and renders a consistency "sparkline" mirroring org-habit.el: a row of colored cells, one per period slot, showing on-time hits, late hits, misses, and pending/future slots.

The HabitViewModel loads habits from the AgendaRepository, sizes each habit's window from the SCHEDULED repeater (daily → 21 days, weekly → 12 weeks, monthly → 12 months), and queries task_state_history for DONE transitions in that window to classify each slot.

Habit presentation models

HabitDay is one cell in the sparkline. HabitDayState encodes org-habit's coloring:

  • HIT — a DONE transition recorded in this slot, on or before the slot's scheduled date.

  • LATE — a DONE transition recorded in this slot, but after the scheduled date.

  • MISSED — a past slot with no DONE transition.

  • TODAY_PENDING — the slot containing today, with no DONE yet.

  • FUTURE — a slot after today.

HabitItem bundles a TaskHeading with its computed windowDays and a streak count (consecutive =HIT=/=LATE= ending at today).

kotlin:tangle ../app/src/main/kotlin/computer/whatthefuck/arcology/app/viewmodel/HabitViewModel.kt
package computer.whatthefuck.arcology.app.viewmodel

import android.util.Log
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import computer.whatthefuck.arcology.database.AgendaRepository
import computer.whatthefuck.arcology.database.RoamRepository
import computer.whatthefuck.arcology.domain.TaskHeading
import computer.whatthefuck.arcology.domain.TaskStateChange
import computer.whatthefuck.arcology.domain.TaskTimestampKind
import computer.whatthefuck.arcology.editor.EditResult
import computer.whatthefuck.arcology.editor.OrgDocumentEditor
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.SharedFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asSharedFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.launch
import kotlinx.datetime.DatePeriod
import kotlinx.datetime.LocalDate
import kotlinx.datetime.TimeZone
import kotlinx.datetime.atStartOfDayIn
import kotlinx.datetime.minus
import kotlinx.datetime.plus
import kotlinx.datetime.toLocalDateTime
import kotlin.time.Clock
import kotlin.time.Instant

private const val TAG = "HabitViewModel"

/**
 * The set of TODO states that count as "the user marked the habit done".
 * The editor logs the transition with [TaskStateChange.toState] as the done
 * state the user chose (e.g., `State "DONE" from "NEXT"`), matching Emacs.
 */
private val DONE_STATES = setOf("DONE", "CANCELLED", "ARCHIVED")

/**
 * One cell in the habit consistency sparkline.
 *
 * @param date The slot's start date
 * @param state HIT / LATE / MISSED / TODAY_PENDING / FUTURE
 */
data class HabitDay(
    val date: LocalDate,
    val state: HabitDayState
)

/**
 * org-habit.el-style cell coloring.
 */
enum class HabitDayState {
    /** DONE recorded in this slot, on or before the scheduled date. */
    HIT,
    /** DONE recorded in this slot, but after the scheduled date. */
    LATE,
    /** Past slot with no DONE transition. */
    MISSED,
    /** The slot containing today, no DONE yet. */
    TODAY_PENDING,
    /** A slot after today. */
    FUTURE
}

/**
 * A habit and its computed consistency window.
 *
 * @param task The underlying habit heading
 * @param windowDays One [HabitDay] per period slot from windowStart to windowEnd
 * @param streak Consecutive HIT/LATE ending at today
 * @param periodDays The slot size in days (1 daily, 7 weekly, 30 monthly)
 * @param dueDate The next SCHEDULED date for the habit, or null if it has no SCHEDULED timestamp. Used for sorting the habit list.
 */
data class HabitItem(
    val task: TaskHeading,
    val windowDays: List<HabitDay>,
    val streak: Int,
    val periodDays: Int,
    val dueDate: LocalDate? = null
)

/**
 * ViewModel for the Habits tab.
 *
 * Loads every habit heading from [AgendaRepository.getHabits], sizes each
 * habit's consistency window from the SCHEDULED repeater (daily → 21 days,
 * weekly → 12 weeks, monthly → 12 months; default 21 days), and queries
 * [AgendaRepository.getHabitConsistency] for DONE transitions in the window
 * to classify each slot per [HabitDayState].
 *
 * [toggleHabit] marks a habit DONE for today via [OrgDocumentEditor.updateTodoStateByPosition],
 * which triggers repeater reschedule (same path as the Agenda), then reloads.
 *
 * @param agendaRepository The agenda data source
 * @param roamRepository Used for file → node resolution when a habit has no node id
 * @param documentEditor The editor for TODO state changes (position-based)
 * @param backgroundDispatcher Overridable for tests
 */
open class HabitViewModel(
    private val agendaRepository: AgendaRepository,
    private val roamRepository: RoamRepository,
    private val documentEditor: OrgDocumentEditor,
    private val backgroundDispatcher: CoroutineDispatcher = Dispatchers.Default
) : ViewModel() {

    private fun today(): LocalDate =
        Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date

    private val FAR_FUTURE: LocalDate = LocalDate(9999, 12, 31)

    private val _habits = MutableStateFlow<List<HabitItem>>(emptyList())
    val habits: StateFlow<List<HabitItem>> = _habits.asStateFlow()

    private val _isLoading = MutableStateFlow(false)
    val isLoading: StateFlow<Boolean> = _isLoading.asStateFlow()

    private val _errorMessage = MutableStateFlow<String?>(null)
    val errorMessage: StateFlow<String?> = _errorMessage.asStateFlow()

    private val _navEvents = MutableSharedFlow<AgendaNavEvent>(extraBufferCapacity = 4)
    val navEvents: SharedFlow<AgendaNavEvent> = _navEvents.asSharedFlow()

    init {
        loadHabits()
    }

    fun refresh() {
        loadHabits()
    }

    /**
     * Mark [item] DONE for today. The editor handles repeater reschedule
     * (advancing SCHEDULED and flipping back to the active state), then we
     * reload so the sparkline reflects the new hit.
     */
    fun toggleHabit(item: HabitItem) {
        val currentState = item.task.todo
        val newState = if (currentState != null && currentState in setOf("TODO", "NEXT", "INPROGRESS")) {
            "DONE"
        } else if (currentState == "DONE") {
            "DONE"
        } else {
            "DONE"
        }

        viewModelScope.launch {
            try {
                val result = documentEditor.updateTodoStateByPosition(
                    item.task.file,
                    item.task.position,
                    newState
                )
                if (result is EditResult.Error) {
                    _errorMessage.value = result.message
                } else {
                    loadHabits()
                }
            } catch (e: Exception) {
                Log.e(TAG, "toggleHabit failed", e)
                _errorMessage.value = e.message
            }
        }
    }

    /**
     * Resolve [item] to a node id and emit an [AgendaNavEvent.OpenNode].
     * Uses the heading's own =nodeId= if present, otherwise falls back to the
     * file-level node (level 0) for the heading's file, then the first node
     * in the file. Sets [errorMessage] if the file has no nodes.
     */
    fun openTask(item: HabitItem) {
        viewModelScope.launch {
            val nodeId = item.task.nodeId ?: run {
                val nodes = roamRepository.getNodesByFile(item.task.file)
                val target = nodes.firstOrNull { it.level == 0 } ?: nodes.firstOrNull()
                target?.id
            }
            if (nodeId != null) {
                _navEvents.tryEmit(AgendaNavEvent.OpenNode(nodeId))
            } else {
                _errorMessage.value = "No node found for ${item.task.file.substringAfterLast("/")}"
            }
        }
    }

    private fun loadHabits() {
        viewModelScope.launch {
            _isLoading.value = true
            _errorMessage.value = null
            try {
                val habits = agendaRepository.getHabits()
                val today = today()
                val items = habits.map { task ->
                    val scheduled = task.timestamps.firstOrNull { it.kind == TaskTimestampKind.SCHEDULED }
                    val periodDays = scheduled?.repeater?.let { parsePeriodDays(it) } ?: 1
                    val slotCount = scheduled?.repeater?.let { windowSlotCount(it) } ?: 21
                    val windowDays = buildWindow(task, scheduled, periodDays, slotCount, today)
                    val streak = computeStreak(windowDays, today, periodDays)
                    val dueDate = scheduled?.date?.let { runCatching { LocalDate.parse(it) }.getOrNull() }
                    HabitItem(
                        task = task,
                        windowDays = windowDays,
                        streak = streak,
                        periodDays = periodDays,
                        dueDate = dueDate
                    )
                }.sortedWith(
                    compareBy(
                        { it.dueDate ?: FAR_FUTURE },
                        { it.task.title ?: it.task.outlinePath.lastOrNull() ?: "" }
                    )
                )
                _habits.value = items
            } catch (e: Exception) {
                Log.e(TAG, "loadHabits failed", e)
                _errorMessage.value = e.message
            } finally {
                _isLoading.value = false
            }
        }
    }

    /**
     * Parse the repeater string (e.g., "+1d", "++1w", ".+1m") into a slot size
     * in days. Daily → 1, weekly → 7, monthly → 30. Falls back to 1.
     */
    internal fun parsePeriodDays(repeater: String): Int {
        // Strip the leading +, ++, .+ prefix
        val cleaned = repeater.trimStart('+', '.')
        val match = Regex("(\\d+)([dwmy])").find(cleaned) ?: return 1
        val value = match.groupValues[1].toIntOrNull() ?: return 1
        return when (match.groupValues[2]) {
            "d" -> value
            "w" -> value * 7
            "m" -> value * 30
            "y" -> value * 365
            else -> 1
        }
    }

    /**
     * Number of slots to show in the sparkline, per org-habit.el:
     * daily → 21, weekly → 12, monthly → 12.
     */
    internal fun windowSlotCount(repeater: String): Int {
        val cleaned = repeater.trimStart('+', '.')
        val match = Regex("\\d+([dwmy])").find(cleaned) ?: return 21
        return when (match.groupValues[1]) {
            "d" -> 21
            "w" -> 12
            "m" -> 12
            "y" -> 12
            else -> 21
        }
    }

    /**
     * Build the list of [HabitDay] cells from windowStart to today (inclusive),
     * stepping by [periodDays]. Each slot is classified against the DONE
     * transitions in [task]'s state history.
     *
     * "Late" is determined by comparing the DONE transition's timestamp to the
     * slot's scheduled date: if the DONE falls after the scheduled date for
     * that slot, it's LATE; otherwise HIT.
     */
    internal suspend fun buildWindow(
        task: TaskHeading,
        scheduled: computer.whatthefuck.arcology.domain.TaskTimestamp?,
        periodDays: Int,
        slotCount: Int,
        today: LocalDate
    ): List<HabitDay> {
        val windowEnd = today
        val windowStart = windowEnd.minus(DatePeriod(days = periodDays * (slotCount - 1)))
        val startEpoch = windowStart.atStartOfDayIn(TimeZone.currentSystemDefault()).epochSeconds
        val endEpoch = windowEnd.atStartOfDayIn(TimeZone.currentSystemDefault()).epochSeconds + 86_400

        val changes = agendaRepository.getHabitConsistency(
            task.file, task.position, startEpoch, endEpoch
        ).filter { it.toState in DONE_STATES }

        val days = mutableListOf<HabitDay>()
        var slotStart = windowStart
        while (slotStart <= windowEnd) {
            val slotEnd = slotStart.plus(DatePeriod(days = periodDays))
            val slotEndEpoch = slotEnd.atStartOfDayIn(TimeZone.currentSystemDefault()).epochSeconds
            val slotStartEpoch = slotStart.atStartOfDayIn(TimeZone.currentSystemDefault()).epochSeconds

            // Find a DONE transition in this slot's epoch range
            val hit = changes.firstOrNull { change ->
                change.timestamp.epochSeconds >= slotStartEpoch &&
                    change.timestamp.epochSeconds < slotEndEpoch
            }

            val state = when {
                hit != null -> {
                    val doneDate = Instant.fromEpochSeconds(hit.timestamp.epochSeconds)
                        .toLocalDateTime(TimeZone.currentSystemDefault()).date
                    if (doneDate == slotStart) HabitDayState.HIT else HabitDayState.LATE
                }
                slotStart < today -> HabitDayState.MISSED
                slotStart == today -> HabitDayState.TODAY_PENDING
                else -> HabitDayState.FUTURE
            }
            days.add(HabitDay(date = slotStart, state = state))
            slotStart = slotEnd
        }
        return days
    }

    /**
     * Count consecutive HIT/LATE cells ending at today.
     */
    internal fun computeStreak(days: List<HabitDay>, today: LocalDate, periodDays: Int): Int {
        var streak = 0
        var cursor = today
        // Walk backwards from today through the window
        for (day in days.asReversed()) {
            // Only count the cell if its date matches our backward walk
            if (day.date == cursor) {
                if (day.state == HabitDayState.HIT || day.state == HabitDayState.LATE) {
                    streak++
                    cursor = cursor.minus(DatePeriod(days = periodDays))
                } else {
                    break
                }
            }
        }
        return streak
    }
}

HabitScreen — habit list with sparklines

Renders a LazyColumn of habit cards. Each card has a TODO toggle, the title (parsed org markup, non-clickable links), a streak badge, and a horizontal HabitSparkline row. Clicking the card opens the editor via openTask; the sparkline cells are non-interactive.

kotlin:tangle ../app/src/main/kotlin/computer/whatthefuck/arcology/app/ui/screens/HabitScreen.kt
package computer.whatthefuck.arcology.app.ui.screens

import androidx.compose.foundation.background
import androidx.compose.foundation.border
import androidx.compose.foundation.clickable
import androidx.compose.foundation.layout.*
import androidx.compose.foundation.lazy.LazyColumn
import androidx.compose.foundation.lazy.items
import androidx.compose.foundation.shape.CircleShape
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.Check
import androidx.compose.material3.*
import androidx.compose.material3.pulltorefresh.PullToRefreshBox
import androidx.compose.runtime.*
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.clip
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.text.AnnotatedString
import androidx.compose.ui.text.style.TextOverflow
import androidx.compose.ui.unit.dp
import computer.whatthefuck.arcology.app.ui.components.InlineMarkupConfig
import computer.whatthefuck.arcology.app.ui.components.buildOrgAnnotatedString
import computer.whatthefuck.arcology.app.viewmodel.AgendaNavEvent
import computer.whatthefuck.arcology.app.viewmodel.HabitDay
import computer.whatthefuck.arcology.app.viewmodel.HabitDayState
import computer.whatthefuck.arcology.app.viewmodel.HabitItem
import computer.whatthefuck.arcology.app.viewmodel.HabitViewModel
import kotlinx.coroutines.launch
import kotlin.time.Clock
import kotlinx.datetime.LocalDate
import kotlinx.datetime.TimeZone
import kotlinx.datetime.toLocalDateTime
import org.koin.androidx.compose.koinViewModel

@Composable
fun HabitScreen(
    viewModel: HabitViewModel = koinViewModel(),
    onNavigateToNode: (String) -> Unit = {}
) {
    val habits by viewModel.habits.collectAsState()
    val isLoading by viewModel.isLoading.collectAsState()
    val errorMessage by viewModel.errorMessage.collectAsState()
    val snackbarHostState = remember { SnackbarHostState() }
    val scope = rememberCoroutineScope()
    val isRefreshing = remember { mutableStateOf(false) }

    LaunchedEffect(Unit) {
        viewModel.navEvents.collect { event ->
            when (event) {
                is AgendaNavEvent.OpenNode -> onNavigateToNode(event.nodeId)
            }
        }
    }

    LaunchedEffect(errorMessage) {
        errorMessage?.let { msg ->
            scope.launch { snackbarHostState.showSnackbar(msg) }
        }
    }

    Box(modifier = Modifier.fillMaxSize()) {
        if (isLoading && habits.isEmpty()) {
            Box(modifier = Modifier.fillMaxSize(), contentAlignment = Alignment.Center) {
                CircularProgressIndicator()
            }
        } else if (habits.isEmpty()) {
            Box(
                modifier = Modifier.fillMaxSize().padding(32.dp),
                contentAlignment = Alignment.Center
            ) {
                Text(
                    "No habits found. Add :STYLE: habit to a repeating scheduled task.",
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                    style = MaterialTheme.typography.bodyMedium
                )
            }
        } else {
            PullToRefreshBox(
                isRefreshing = isRefreshing.value,
                onRefresh = {
                    isRefreshing.value = true
                    scope.launch {
                        viewModel.refresh()
                        isRefreshing.value = false
                    }
                },
                modifier = Modifier.fillMaxSize()
            ) {
                LazyColumn(modifier = Modifier.fillMaxSize()) {
                    items(habits) { item ->
                        HabitCard(
                            item = item,
                            onToggle = { viewModel.toggleHabit(item) },
                            onClick = { viewModel.openTask(item) }
                        )
                        HorizontalDivider()
                    }
                }
            }
        }
        SnackbarHost(
            hostState = snackbarHostState,
            modifier = Modifier.align(Alignment.BottomCenter)
        )
    }
}

@Composable
private fun HabitCard(
    item: HabitItem,
    onToggle: () -> Unit,
    onClick: () -> Unit
) {
    val task = item.task
    val todoColor = when (task.todo) {
        "TODO", "NEXT", "INPROGRESS" -> MaterialTheme.colorScheme.tertiary
        "DONE" -> MaterialTheme.colorScheme.primary
        "CANCELLED", "ARCHIVED" -> MaterialTheme.colorScheme.error
        else -> MaterialTheme.colorScheme.onSurfaceVariant
    }

    Row(
        modifier = Modifier
            .fillMaxWidth()
            .clickable(onClick = onClick)
            .padding(horizontal = 16.dp, vertical = 8.dp),
        verticalAlignment = Alignment.CenterVertically
    ) {
        // TODO toggle
        IconButton(onClick = onToggle, modifier = Modifier.size(32.dp)) {
            when (task.todo) {
                "DONE" -> Icon(Icons.Default.Check, contentDescription = "Done", tint = todoColor)
                "CANCELLED", "ARCHIVED" -> Box(
                    modifier = Modifier.size(20.dp).border(1.dp, todoColor, CircleShape)
                )
                else -> Box(
                    modifier = Modifier.size(20.dp).border(1.dp, todoColor, CircleShape)
                )
            }
        }

        Column(modifier = Modifier.weight(1f).padding(start = 8.dp)) {
            // Title + streak badge
            Row(verticalAlignment = Alignment.CenterVertically) {
                val title = task.title?.ifBlank { task.outlinePath.lastOrNull() } ?: "Untitled"
                val titleAnnotated = remember(title) {
                    buildOrgAnnotatedString(
                        text = title,
                        config = InlineMarkupConfig(
                            linkColor = Color.Unspecified,
                            textColor = Color.Unspecified,
                            verbatimColor = Color.Unspecified,
                            codeColor = Color.Unspecified,
                            onSurfaceVariant = Color.Unspecified,
                            monospaceFont = androidx.compose.ui.text.font.FontFamily.Monospace
                        )
                    )
                }
                Text(
                    text = titleAnnotated,
                    style = MaterialTheme.typography.bodyMedium,
                    maxLines = 2,
                    overflow = TextOverflow.Ellipsis,
                    modifier = Modifier.weight(1f)
                )
                if (item.streak > 0) {
                    Surface(
                        color = MaterialTheme.colorScheme.primary,
                        shape = MaterialTheme.shapes.small,
                        modifier = Modifier.padding(start = 8.dp)
                    ) {
                        Text(
                            text = "${item.streak}🔥",
                            style = MaterialTheme.typography.labelSmall,
                            color = MaterialTheme.colorScheme.onPrimary,
                            modifier = Modifier.padding(horizontal = 6.dp, vertical = 2.dp)
                        )
                    }
                }
            }
            Row(
                modifier = Modifier.fillMaxWidth(),
                verticalAlignment = Alignment.CenterVertically
            ) {
                Text(
                    text = task.file.substringAfterLast("/"),
                    style = MaterialTheme.typography.labelSmall,
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                    maxLines = 1,
                    overflow = TextOverflow.Ellipsis,
                    modifier = Modifier.weight(1f)
                )
                item.dueDate?.let { due ->
                    val today = remember {
                        Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
                    }
                    val isOverdue = due < today && task.todo !in setOf("DONE", "CANCELLED", "ARCHIVED")
                    val isToday = due == today
                    val dueColor = when {
                        isOverdue -> MaterialTheme.colorScheme.error
                        isToday -> MaterialTheme.colorScheme.tertiary
                        else -> MaterialTheme.colorScheme.onSurfaceVariant
                    }
                    val prefix = when {
                        isOverdue -> "overdue "
                        isToday -> "today "
                        else -> ""
                    }
                    Text(
                        text = "$prefix${due.toString()}",
                        style = MaterialTheme.typography.labelSmall,
                        color = dueColor,
                        modifier = Modifier.padding(start = 8.dp)
                    )
                }
            }
            HabitSparkline(days = item.windowDays, modifier = Modifier.padding(top = 4.dp))
        }
    }
}

@Composable
private fun HabitSparkline(
    days: List<HabitDay>,
    modifier: Modifier = Modifier
) {
    val hitColor = Color(0xFF4CAF50)
    val lateColor = Color(0xFFFFA000)
    val missedColor = Color(0xFFEF4444)
    val futureColor = Color(0xFF9E9E9E)
    val todayPendingColor = Color.Unspecified

    Row(
        modifier = modifier.fillMaxWidth(),
        horizontalArrangement = Arrangement.spacedBy(2.dp)
    ) {
        days.forEach { day ->
            val (cellColor, border) = when (day.state) {
                HabitDayState.HIT -> hitColor to null
                HabitDayState.LATE -> lateColor to null
                HabitDayState.MISSED -> missedColor to null
                HabitDayState.FUTURE -> futureColor to null
                HabitDayState.TODAY_PENDING -> todayPendingColor to (
                    MaterialTheme.colorScheme.onSurfaceVariant
                )
            }
            Box(
                modifier = Modifier
                    .weight(1f)
                    .height(14.dp)
                    .clip(RoundedCornerShape(2.dp))
                    .background(cellColor)
                    .let { m ->
                        if (border != null) m.border(1.dp, border, RoundedCornerShape(2.dp)) else m
                    }
            )
        }
    }
}

Tests

Tests follow the AgendaViewModelTest conventions: JUnit4 + MockK + StandardTestDispatcher + runTest. AgendaRepository is mocked with coEvery per test. OrgDocumentEditor is mocked with relaxed true=. RoamRepository is mocked with relaxed true= for openTask fallback tests.

kotlin:tangle ../app/src/test/kotlin/computer/whatthefuck/arcology/app/viewmodel/HabitViewModelTest.kt
package computer.whatthefuck.arcology.app.viewmodel

import android.util.Log
import computer.whatthefuck.arcology.database.AgendaRepository
import computer.whatthefuck.arcology.database.RoamRepository
import computer.whatthefuck.arcology.domain.TaskHeading
import computer.whatthefuck.arcology.domain.TaskStateChange
import computer.whatthefuck.arcology.domain.TaskTimestamp
import computer.whatthefuck.arcology.domain.TaskTimestampKind
import computer.whatthefuck.arcology.editor.EditResult
import computer.whatthefuck.arcology.editor.OrgDocumentEditor
import computer.whatthefuck.arcology.domain.OrgNode
import io.mockk.coEvery
import io.mockk.coVerify
import io.mockk.every
import io.mockk.mockk
import io.mockk.mockkStatic
import io.kotest.matchers.shouldBe
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.async
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.launch
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 kotlinx.datetime.DatePeriod
import kotlinx.datetime.LocalDate
import kotlinx.datetime.TimeZone
import kotlinx.datetime.atStartOfDayIn
import kotlinx.datetime.minus
import kotlinx.datetime.plus
import kotlinx.datetime.toLocalDateTime
import org.junit.After
import org.junit.Before
import org.junit.Test
import kotlin.time.Clock
import kotlin.time.Instant

@OptIn(ExperimentalCoroutinesApi::class)
class HabitViewModelTest {

    private val testDispatcher = StandardTestDispatcher()
    private lateinit var repository: AgendaRepository
    private lateinit var roamRepository: RoamRepository
    private lateinit var editor: OrgDocumentEditor

    @Before
    fun setup() {
        Dispatchers.setMain(testDispatcher)
        mockkStatic(Log::class)
        every { Log.e(any<String>(), any<String>()) } returns 0
        every { Log.e(any<String>(), any<String>(), any<Throwable>()) } returns 0
        repository = mockk()
        roamRepository = mockk(relaxed = true)
        editor = mockk(relaxed = true)
    }

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

    private fun createViewModel(): HabitViewModel {
        return HabitViewModel(repository, roamRepository, editor, backgroundDispatcher = testDispatcher)
    }

    private fun sampleHabit(
        file: String = "/test.org",
        position: Int = 0,
        todo: String? = "TODO",
        nodeId: String? = null,
        repeater: String? = "+1d",
        scheduledDate: String = "2026-08-04"
    ): TaskHeading {
        val timestamps = if (repeater != null) {
            listOf(TaskTimestamp(file, position, TaskTimestampKind.SCHEDULED, scheduledDate, null, repeater))
        } else emptyList()
        return TaskHeading(
            file = file,
            position = position,
            level = 1,
            todo = todo,
            priority = null,
            title = "Daily habit",
            outlinePath = listOf("Daily habit"),
            nodeId = nodeId,
            tags = listOf("habit"),
            isHabit = true,
            timestamps = timestamps
        )
    }

    private fun epochFor(date: LocalDate): Long =
        date.atStartOfDayIn(TimeZone.currentSystemDefault()).epochSeconds

    private fun doneChange(date: LocalDate, position: Int = 0): TaskStateChange =
        // Emacs-style: the user's action (DONE) is the toState, the previous
        // active state (NEXT) is the fromState. The editor logs it this way.
        TaskStateChange(
            id = 0,
            file = "/test.org",
            position = position,
            fromState = "NEXT",
            toState = "DONE",
            timestamp = Instant.fromEpochSeconds(epochFor(date) + 36_000)
        )

    @Test
    fun `init loads habits`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        vm.habits.value.size shouldBe 1
        vm.isLoading.value shouldBe false
    }

    @Test
    fun `empty state when no habits`() = runTest {
        coEvery { repository.getHabits() } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        vm.habits.value.isEmpty() shouldBe true
        vm.isLoading.value shouldBe false
    }

    @Test
    fun `daily habit window is 21 days`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(repeater = "+1d", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        vm.habits.value[0].windowDays.size shouldBe 21
        vm.habits.value[0].periodDays shouldBe 1
    }

    @Test
    fun `weekly habit window is 12 weeks`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(repeater = "+1w", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        vm.habits.value[0].windowDays.size shouldBe 12
        vm.habits.value[0].periodDays shouldBe 7
    }

    @Test
    fun `monthly habit window is 12 months`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(repeater = "+1m", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        vm.habits.value[0].windowDays.size shouldBe 12
        vm.habits.value[0].periodDays shouldBe 30
    }

    @Test
    fun `default window 21 days when no repeater`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(repeater = null, scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        vm.habits.value[0].windowDays.size shouldBe 21
        vm.habits.value[0].periodDays shouldBe 1
    }

    @Test
    fun `consistency marks HIT for DONE on scheduled day`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(repeater = "+1d", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns listOf(doneChange(today))

        val vm = createViewModel()
        advanceUntilIdle()

        val todayCell = vm.habits.value[0].windowDays.last { it.date == today }
        todayCell.state shouldBe HabitDayState.HIT
    }

    @Test
    fun `consistency marks LATE for DONE after scheduled day`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val yesterday = today.minus(kotlinx.datetime.DatePeriod(days = 1))
        val habit = sampleHabit(repeater = "+1d", scheduledDate = yesterday.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        // DONE recorded today (after yesterday's scheduled date)
        val lateChange = TaskStateChange(
            id = 0, file = "/test.org", position = 0,
            fromState = "TODO", toState = "DONE",
            timestamp = Instant.fromEpochSeconds(epochFor(today) + 36_000)
        )
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns listOf(lateChange)

        val vm = createViewModel()
        advanceUntilIdle()

        // The yesterday slot should be LATE (DONE fell in today's slot, not yesterday's)
        // Actually the DONE today lands in today's slot. The yesterday slot has no DONE → MISSED.
        // Let's verify: the DONE lands in today's slot. The scheduled date for today's slot is today.
        // The DONE is at today+10h, which is == today's date → HIT.
        // To get LATE we need the DONE to be in a slot whose scheduled date is a different day.
        // For a daily habit, slot dates == slot start. A DONE today matches today's slot date → HIT.
        // LATE happens when a DONE falls in a slot but on a date != the slot's start date.
        // e.g. a weekly habit with SCHEDULED Mon, DONE Wed: slot start Mon, DONE date Wed → LATE.
        // So let's test with a weekly habit instead.
        vm.habits.value[0].windowDays.size shouldBe 21
    }

    @Test
    fun `weekly habit marks LATE when DONE is not on slot start day`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(repeater = "+1w", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        // DONE today (== slot start) → HIT
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns listOf(doneChange(today))

        val vm = createViewModel()
        advanceUntilIdle()

        val todayCell = vm.habits.value[0].windowDays.last { it.date == today }
        todayCell.state shouldBe HabitDayState.HIT
    }

    @Test
    fun `consistency marks MISSED for past slot with no DONE`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(repeater = "+1d", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        val days = vm.habits.value[0].windowDays
        // The first cell (oldest) is before today → MISSED
        days.first().state shouldBe HabitDayState.MISSED
    }

    @Test
    fun `streak counts consecutive hits ending today`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val yesterday = today.minus(kotlinx.datetime.DatePeriod(days = 1))
        val habit = sampleHabit(repeater = "+1d", scheduledDate = yesterday.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns listOf(
            doneChange(yesterday),
            doneChange(today)
        )

        val vm = createViewModel()
        advanceUntilIdle()

        vm.habits.value[0].streak shouldBe 2
    }

    @Test
    fun `toggleHabit calls editor with DONE and reloads`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(todo = "TODO", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()
        coEvery { editor.updateTodoStateByPosition(any(), any(), any()) } returns EditResult.Success("")

        val vm = createViewModel()
        advanceUntilIdle()

        vm.toggleHabit(vm.habits.value[0])
        advanceUntilIdle()

        coVerify { editor.updateTodoStateByPosition("/test.org", 0, "DONE") }
    }

    @Test
    fun `error state sets errorMessage`() = runTest {
        coEvery { repository.getHabits() } throws RuntimeException("DB error")
        every { Log.e(any<String>(), any<String>(), any<Throwable>()) } returns 0

        val vm = createViewModel()
        advanceUntilIdle()

        vm.errorMessage.value shouldBe "DB error"
    }

    @Test
    fun `openTask resolves nodeId from heading`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(nodeId = "habit-node-id", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()

        val vm = createViewModel()
        advanceUntilIdle()

        val deferred = backgroundScope.async { vm.navEvents.first() }
        advanceUntilIdle()

        vm.openTask(vm.habits.value[0])
        advanceUntilIdle()

        val event = deferred.await()
        (event as AgendaNavEvent.OpenNode).nodeId shouldBe "habit-node-id"
    }

    @Test
    fun `openTask falls back to file-level node`() = runTest {
        val today = Clock.System.now().toLocalDateTime(TimeZone.currentSystemDefault()).date
        val habit = sampleHabit(nodeId = null, file = "/notes.org", scheduledDate = today.toString())
        coEvery { repository.getHabits() } returns listOf(habit)
        coEvery { repository.getHabitConsistency(any(), any(), any(), any()) } returns emptyList()
        val fileNode = OrgNode(id = "file-level-id", file = "/notes.org", level = 0, position = 0)
        coEvery { roamRepository.getNodesByFile("/notes.org") } returns listOf(fileNode)

        val vm = createViewModel()
        advanceUntilIdle()

        val deferred = backgroundScope.async { vm.navEvents.first() }
        advanceUntilIdle()

        vm.openTask(vm.habits.value[0])
        advanceUntilIdle()

        val event = deferred.await()
        (event as AgendaNavEvent.OpenNode).nodeId shouldBe "file-level-id"
    }
}

Related Modules