package compose.widgets

import androidx.compose.foundation.lazy.LazyListState
import androidx.compose.runtime.*
import androidx.compose.ui.Modifier
import androidx.compose.ui.input.key.Key
import androidx.compose.ui.input.key.onKeyEvent
import androidx.compose.ui.input.key.onPreviewKeyEvent
import compose.input.onKeyShortcut
import kotlin.math.abs
import kotlin.math.max


/**
 * Models selection in a list.
 */
interface ListSelectionModel {


    /**
     * The index of the "focused" item, not in the sense that is has the keyboard focus, but in the sense that this is
     * the item with which the user last had interaction. For example in a list with a multi-item selection model, you
     * can select many items by pressing shift-up - the last selected index will be the focused index.
     * List views which has a [ListSelectionModel] should keep the item at this index visible, and may mark it with a
     * special indicator.
     */
    val focusedIndex: Int?


    /**
     * Returns whether the item at the given index is selected.
     */
    fun isIndexSelected(index: Int): Boolean


    /**
     * Selects the item at the given index.
     */
    fun selectIndex(index: Int)


    /**
     * Deselects the given index.
     */
    fun deselectIndex(index: Int)


    /**
     * Clears the selection.
     */
    fun clearSelection()


    /**
     * Selects the "previous" item (the one before the currently selected one).
     *
     * Returns whether the selection actually changed.
     */
    fun selectPrevious(): Boolean


    /**
     * Selects the "next" item (following the currently selected one).
     *
     * Returns whether the selection actually changed.
     */
    fun selectNext(): Boolean


    /**
     * Selects the first selectable index.
     *
     * Returns whether the selection actually changed.
     */
    fun selectFirst(): Boolean


    /**
     * Selects the last selectable index.
     *
     * Returns whether the selection actually changed.
     */
    fun selectLast(): Boolean


    /**
     * Selects the item that is one page before the currently selected one.
     *
     * Returns whether the selection actually changed.
     */
    fun selectPreviousPage(itemsInPage: Int): Boolean


    /**
     * Selects the item that is one page after the currently selected one.
     *
     * Returns whether the selection actually changed.
     */
    fun selectNextPage(itemsInPage: Int): Boolean


}


/**
 * A [ListSelectionModel] where no items can be selected.
 */
val NoSelectionModel = unmodifiableSelectionModel(
    selectedIndex = null,
    onSelectRequest = { false },
    selectableRange = IntRange.EMPTY
)


/**
 * Returns a [ListSelectionModel] with a single selected item, and which doesn't change by itself, only reporting
 * requests to change the selection to a given callback.
 *
 * This is useful to implement a selection that is determined by some external state.
 */
fun unmodifiableSelectionModel(
    selectedIndex: Int?,
    onSelectRequest: (Int) -> Boolean,
    selectableRange: IntRange,
): ListSelectionModel = object: ListSelectionModel {


    override val focusedIndex: Int?
        get() = null


    override fun isIndexSelected(index: Int) = index == selectedIndex


    override fun selectIndex(index: Int) {
        onSelectRequest(index)
    }


    override fun deselectIndex(index: Int) { }


    override fun clearSelection() { }


    override fun selectPrevious(): Boolean {
        return if (selectedIndex == null)
            false
        else
            onSelectRequest((selectedIndex - 1).coerceIn(selectableRange))
    }


    override fun selectNext(): Boolean {
        return if (selectedIndex == null)
            false
        else
            onSelectRequest((selectedIndex + 1).coerceIn(selectableRange))
    }


    override fun selectFirst(): Boolean {
        return if (selectedIndex == null)
            false
        else
            onSelectRequest(selectableRange.first)
    }


    override fun selectLast(): Boolean {
        return if (selectedIndex == null)
            false
        else
            onSelectRequest(selectableRange.last)
    }


    override fun selectPreviousPage(itemsInPage: Int): Boolean {
        return if (selectedIndex == null)
            false
        else
            onSelectRequest((selectedIndex - itemsInPage).coerceIn(selectableRange))
    }


    override fun selectNextPage(itemsInPage: Int): Boolean {
        return if (selectedIndex == null)
            false
        else
            onSelectRequest((selectedIndex + itemsInPage).coerceIn(selectableRange))
    }


}


/**
 * A single-item selection model.
 */
abstract class SingleItemSelectionModel(


    /**
     * The initially selected index; `null` if none.
     */
    initialSelectedIndex: Int? = null


) : ListSelectionModel {


    /**
     * Returns the smallest selectable index. The default value is 0.
     */
    open val minSelectableIndex: Int
        get() = 0


    /**
     * The largest selectable index. The default value is [Int.MAX_VALUE].
     */
    open val maxSelectableIndex: Int
        get() = Int.MAX_VALUE


    /**
     * Returns whether the given index can be selected (even if it's inside the min..max selectable range).
     *
     * The default implementation always returns `true`. This allows to add non-selectable items to the list, such as
     * section headers.
     */
    open fun isSelectable(index: Int): Boolean {
        return true
    }


    /**
     * The selected index; `null` if none.
     */
    var selectedIndex: Int? by mutableStateOf(initialSelectedIndex)
        protected set


    override var focusedIndex: Int? by mutableStateOf(initialSelectedIndex)
        protected set


    override fun isIndexSelected(index: Int) = index == this.selectedIndex


    override fun selectIndex(index: Int) {
        if ((index in minSelectableIndex..maxSelectableIndex) && isSelectable(index)) {
            selectIndexImpl(index)
        }
    }


    override fun deselectIndex(index: Int) {
        if (this.selectedIndex == index)
            this.selectedIndex = null
        this.focusedIndex = index
    }


    override fun clearSelection() {
        selectIndexImpl(null)
    }


    /**
     * Selects the given index, returns whether it changed.
     */
    private fun selectIndexImpl(index: Int?): Boolean {
        if (index == selectedIndex)
            return false

        this.selectedIndex = index
        if (index != null)
            this.focusedIndex = index
        return true
    }


    override fun selectPrevious(): Boolean {
        return selectIndexImpl(previousIndex(selectedIndex, offset = 1))
    }


    override fun selectNext(): Boolean {
        return selectIndexImpl(nextIndex(selectedIndex, offset = 1))
    }


    override fun selectFirst(): Boolean {
        return selectIndexImpl(minSelectableIndex)
    }


    override fun selectLast(): Boolean {
        return selectIndexImpl(maxSelectableIndex)
    }


    override fun selectPreviousPage(itemsInPage: Int): Boolean {
        return selectIndexImpl(previousIndex(selectedIndex, offset = itemsInPage))
    }


    override fun selectNextPage(itemsInPage: Int): Boolean {
        return selectIndexImpl(nextIndex(selectedIndex, offset = itemsInPage))
    }

    /**
     * Returns whether no indices can be selected.
     */
    private fun isEmpty() = maxSelectableIndex < minSelectableIndex


    /**
     * Returns the "previous" selectable index relative to current selected index.
     */
    protected open fun previousIndex(selectedIndex: Int?, offset: Int): Int? {
        var index = if (selectedIndex == null)
            if (isEmpty()) null else maxSelectableIndex
        else
            max(selectedIndex - offset, minSelectableIndex)

        if (index == null)
            return null

        // Look for nearest selectable index upwards
        while (!isSelectable(index) && (index >= minSelectableIndex))
            index--

        return index.takeIf { it >= minSelectableIndex } ?: selectedIndex
    }


    /**
     * Returns the "next" selectable index relative to the current selected index.
     */
    protected open fun nextIndex(selectedIndex: Int?, offset: Int): Int? {
        var index = if (selectedIndex == null)
            if (isEmpty()) null else minSelectableIndex
        else {
            // Avoid overflow if max is Int.MAX_VALUE
            if (maxSelectableIndex - selectedIndex < offset)
                maxSelectableIndex
            else
                selectedIndex + offset
        }

        if (index == null)
            return null

        // Look for nearest selectable index downwards
        while (!isSelectable(index) && (index <= maxSelectableIndex))
            index++

        return index.takeIf { it <= maxSelectableIndex } ?: selectedIndex
    }


}


/**
 * A single-item selection model with mutable range of selectable indices.
 */
class RangeSingleItemSelectionModel(


    /**
     * The initially selected index; `null` if none.
     */
    initialSelectedIndex: Int? = null,


    /**
     * The range of indices that can be selected.
     */
    selectableRange: IntRange = 0..Int.MAX_VALUE


): SingleItemSelectionModel(initialSelectedIndex = initialSelectedIndex) {


    /**
     * A constructor that takes separate minimum and maximum selectable indices.
     */
    @Suppress("unused")
    constructor(
        initialSelectedIndex: Int? = null,
        minSelectableIndex: Int = 0,
        maxSelectableIndex: Int = Int.MAX_VALUE
    ): this(
        initialSelectedIndex = initialSelectedIndex,
        selectableRange = minSelectableIndex .. maxSelectableIndex
    )


    /**
     * The range of indices that can be selected.
     */
    var selectableRange: IntRange = selectableRange
        set(value) {
            field = value
            selectedIndex?.let {
                if (it !in value) {
                    if (value.isEmpty())
                        clearSelection()
                    else
                        selectIndex(
                            if (abs(it - value.first) < abs(it - value.last)) value.first else value.last
                        )
                }
            }
        }


    override val minSelectableIndex: Int
        get() = selectableRange.first


    override val maxSelectableIndex: Int
        get() = selectableRange.last


    companion object {


        /**
         * A [RangeSingleItemSelectionModel] with an empty range.
         */
        val Empty = RangeSingleItemSelectionModel(initialSelectedIndex = null, selectableRange = IntRange.EMPTY)


    }

}


/**
 * Returns a remembered single-item selection model.
 */
@Composable
fun rememberSingleItemSelectionModel(
    initialSelectedIndex: Int? = null,
    minSelectableIndex: Int = 0,
    maxSelectableIndex: Int = Int.MAX_VALUE
) = rememberSingleItemSelectionModel(
    initialSelectedIndex = initialSelectedIndex,
    selectableRange = minSelectableIndex .. maxSelectableIndex
)


/**
 * Returns a remembered single-item selection model.
 */
@Composable
fun rememberSingleItemSelectionModel(
    initialSelectedIndex: Int? = null,
    selectableRange: IntRange,
): SingleItemSelectionModel {
    return remember {
        RangeSingleItemSelectionModel(
            initialSelectedIndex = initialSelectedIndex,
            selectableRange = selectableRange
        )
    }.also {
        it.selectableRange = selectableRange
    }
}


/**
 * A single-item selection model for a list holding the given items.
 */
class SingleItemListSelectionModel<T: Any>(


    /**
     * The items in the list.
     */
    val items: List<T>,


    /**
     * The initially selected index.
     */
    initialSelectedIndex: Int? = 0,


): SingleItemSelectionModel(initialSelectedIndex = initialSelectedIndex) {


    override val minSelectableIndex: Int
        get() = 0


    override val maxSelectableIndex: Int
        get() = items.lastIndex


    /**
     * Returns the selected item, or `null` if none.
     */
    fun selectedItem() = selectedIndex?.let { items[it] }


    /**
     * Selects the given item.
     */
    fun selectItem(item: T?) {
        if (item == null) {
            selectedIndex?.let { deselectIndex(it) }
            return
        }

        selectIndex(items.indexOf(item))
    }


}


/**
 * Returns a remembered single-item selection model for a list holding the given items.
 */
@Composable
fun <T: Any> rememberSingleItemSelectionModel(


    /**
     * The items in the list.
     */
    items: List<T>,


    /**
     * The initially selected index.
     */
    initialSelectedIndex: Int? = null,


): SingleItemListSelectionModel<T> = remember(items, items.size) {
    SingleItemListSelectionModel(items, initialSelectedIndex)
}


/**
 * Makes the element focusable, focuses it when clicked, and moves the selection when the UP, DOWN, HOME, END, PAGE-UP
 * and PAGE-DOWN keys are pressed.
 */
fun Modifier.moveSelectionWithKeys(


    /**
     * The [ListSelectionModel] to notify when keys are pressed.
     */
    selectionModel: ListSelectionModel,


    /**
     * Returns the number of items in a page. If `null`, page-up and page-down keys won't do anything.
     */
    itemsInPage: (() -> Int)? = null,


    /**
     * Whether to use [Modifier.onPreviewKeyEvent] when listening to key events.
     * Otherwise, [Modifier.onKeyEvent] will be used.
     */
    onPreview: Boolean = false,


): Modifier = this
    .onKeyShortcut(Key.DirectionUp, onPreview = onPreview) {
        selectionModel.selectPrevious()
    }
    .onKeyShortcut(Key.DirectionDown, onPreview = onPreview) {
        selectionModel.selectNext()
    }
    .onKeyShortcut(Key.MoveHome, onPreview = onPreview) {
        selectionModel.selectFirst()
    }
    .onKeyShortcut(Key.MoveEnd, onPreview = onPreview) {
        selectionModel.selectLast()
    }
    .then(
        if (itemsInPage == null)
            Modifier
        else {
            Modifier
                .onKeyShortcut(Key.PageUp, onPreview = onPreview) {
                    selectionModel.selectPreviousPage(itemsInPage())
                }
                .onKeyShortcut(Key.PageDown, onPreview = onPreview) {
                    selectionModel.selectNextPage(itemsInPage())
                }
        }
    )


/**
 * Returns the number of items in a page of a [LazyListState]. This can be used as the `itemsInPage` argument to
 * [Modifier.moveSelectionWithKeys].
 */
fun LazyListState.itemsInPage() =
    layoutInfo.visibleItemsInfo.size - 1  // Minus one because the last item is typically only partially visible
