package com.craigvg.lichun_android.ui.navigation

import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateListOf
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.setValue
import androidx.compose.runtime.snapshots.SnapshotStateList
import androidx.navigation3.runtime.NavKey

/**
 * Navigation state holder for Navigation 3 that combines:
 *
 *  - **Multiple back stacks** — one per top-level (bottom-bar) tab, each
 *    retaining its own history (mirrors the legacy `popUpTo(start){saveState};
 *    restoreState=true` tab behavior).
 *  - **A full-screen override stack** — for destinations like Onboarding and
 *    Death that must take over the whole screen and clear all tab history
 *    (mirrors the legacy `popUpTo(0){inclusive=true}`).
 *
 * It exposes a single flattened [backStack] for [androidx.navigation3.ui.NavDisplay],
 * the current [topLevelKey] for bottom-bar selection, and [isTopLevel] /
 * [isOverride] so the UI can derive chrome visibility from key *type* rather
 * than a route-string checklist.
 */
class TabbedBackStack(
    private val tabs: List<TopLevelDestination>,
    startTab: TopLevelDestination = tabs.first(),
) {
    // One stack per tab. Each starts with just the tab key.
    private val tabStacks: LinkedHashMap<TopLevelDestination, SnapshotStateList<NavKey>> =
        LinkedHashMap<TopLevelDestination, SnapshotStateList<NavKey>>().apply {
            tabs.forEach { put(it, mutableStateListOf<NavKey>(it)) }
        }

    /** The currently selected top-level tab. */
    var topLevelKey: TopLevelDestination by mutableStateOf(startTab)
        private set

    // When non-empty, these entries fully replace the tab UI (e.g. Onboarding, Death).
    private val overrideStack: SnapshotStateList<NavKey> = mutableStateListOf()

    /** Flattened back stack rendered by NavDisplay. Never empty. */
    val backStack: SnapshotStateList<NavKey> = mutableStateListOf(startTab)

    /** The key currently displayed (top of the rendered stack). */
    val currentKey: NavKey
        get() = backStack.last()

    /** True when the visible destination is a bottom-bar tab root (show chrome). */
    val isTopLevel: Boolean
        get() = overrideStack.isEmpty() && backStack.last() is TopLevelDestination

    /** True when a full-screen override (Onboarding/Death) is active. */
    val isOverride: Boolean
        get() = overrideStack.isNotEmpty()

    private fun recompute() {
        backStack.apply {
            clear()
            if (overrideStack.isNotEmpty()) {
                addAll(overrideStack)
            } else {
                addAll(tabStacks[topLevelKey] ?: mutableStateListOf(topLevelKey))
            }
        }
    }

    /** Switch to a top-level tab, preserving that tab's history. */
    fun switchTab(tab: TopLevelDestination) {
        // Switching tabs always exits any full-screen override.
        overrideStack.clear()
        if (tabStacks[tab] == null) {
            tabStacks[tab] = mutableStateListOf(tab)
        }
        topLevelKey = tab
        recompute()
    }

    /** Push a detail/child destination onto the current context. */
    fun push(key: NavKey) {
        when {
            key is TopLevelDestination -> switchTab(key)
            overrideStack.isNotEmpty() -> {
                overrideStack.add(key)
                recompute()
            }
            else -> {
                tabStacks[topLevelKey]?.add(key)
                recompute()
            }
        }
    }

    /**
     * Replace the entire navigation state with a single [key], clearing all
     * tab history. Used for full-screen flows (Onboarding/Death) and for
     * landing back on a clean tab root when such a flow completes.
     *
     * If [key] is a top-level tab, the override is cleared and that tab's root
     * is shown. Otherwise [key] becomes a full-screen override.
     */
    fun resetTo(key: NavKey) {
        // Reset every tab back to its root so returning later starts fresh.
        tabs.forEach { tab -> tabStacks[tab] = mutableStateListOf(tab) }
        overrideStack.clear()
        if (key is TopLevelDestination) {
            topLevelKey = key
        } else {
            topLevelKey = tabs.first()
            overrideStack.add(key)
        }
        recompute()
    }

    /**
     * Pop the current destination. Returns false when there is nothing left to
     * pop (i.e. the app should exit). Mirrors NavDisplay's back behavior.
     */
    fun removeLast(): Boolean {
        if (overrideStack.isNotEmpty()) {
            // Within an override flow (rare), pop it; if it empties, fall back to tabs.
            if (overrideStack.size > 1) {
                overrideStack.removeAt(overrideStack.lastIndex)
                recompute()
                return true
            }
            // A single override with no tab history beneath it = exit.
            overrideStack.removeAt(overrideStack.lastIndex)
            recompute()
            return backStack.isNotEmpty()
        }

        val current = tabStacks[topLevelKey] ?: return false
        if (current.size > 1) {
            current.removeAt(current.lastIndex)
            recompute()
            return true
        }
        // At a tab root: if not on the first tab, fall back to the first tab.
        if (topLevelKey != tabs.first()) {
            topLevelKey = tabs.first()
            recompute()
            return true
        }
        // On the first tab's root — let the system handle exit.
        return false
    }

    init {
        recompute()
    }
}