/**
 * A lazy column with advanced functionality.
 */

package compose.widgets

import androidx.compose.foundation.ExperimentalFoundationApi
import androidx.compose.foundation.background
import androidx.compose.foundation.gestures.FlingBehavior
import androidx.compose.foundation.gestures.ScrollableDefaults
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.PaddingValues
import androidx.compose.foundation.lazy.*
import androidx.compose.foundation.rememberScrollbarAdapter
import androidx.compose.material.MaterialTheme
import androidx.compose.runtime.*
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.unit.dp
import compose.input.onMousePress
import compose.utils.markSelected
import compose.utils.scrollToMakeVisible
import kotlinx.coroutines.flow.collectLatest
import kotlinx.coroutines.flow.distinctUntilChanged
import kotlinx.coroutines.flow.withIndex


/**
 * A [LazyColumn] with additional capabilities:
 * - A vertical bar to indicate the scroll position.
 * - Item selection.
 * - Scroll shadows.
 */
@Composable
fun LazyColumnExt(
    modifier: Modifier = Modifier,
    state: LazyListState = rememberLazyListState(),
    contentPadding: PaddingValues = PaddingValues(0.dp),
    reverseLayout: Boolean = false,
    verticalArrangement: Arrangement.Vertical = if (!reverseLayout) Arrangement.Top else Arrangement.Bottom,
    horizontalAlignment: Alignment.Horizontal = Alignment.Start,
    flingBehavior: FlingBehavior = ScrollableDefaults.flingBehavior(),
    showVerticalScrollbar: Boolean = true,
    selection: ListSelectionModel = NoSelectionModel,
    selectionBackground: Color = MaterialTheme.colors.secondary,
    showScrollShadows: Boolean = false,
    content: LazyListScope.() -> Unit
){
    Box(modifier = modifier) {
        if (showScrollShadows) {
            TopScrollShadow(state)
            BottomScrollShadow(state)
        }
        LazyColumn(
            state = state,
            contentPadding = contentPadding,
            reverseLayout = reverseLayout,
            verticalArrangement = verticalArrangement,
            horizontalAlignment = horizontalAlignment,
            flingBehavior = flingBehavior,
            content = {
                val proxyScope = LazyListScopeWithSelection(
                    delegate = this,
                    selection = selection,
                    selectionBackground = selectionBackground
                )
                content.invoke(proxyScope)
            }
        )

        val scrollbarAdapter = rememberScrollbarAdapter(state)
        if (showVerticalScrollbar){
            VerticalScrollbarInBox(
                adapter = scrollbarAdapter,
                reverseLayout = reverseLayout,
                modifier = Modifier.align(Alignment.CenterEnd),
            )
        }

        scrollToFocusedIndex(state, selection)
    }
}


/**
 * Scrolls the given [LazyListState] to the focused index of the selection model.
 */
@Composable
private fun scrollToFocusedIndex(state: LazyListState, selection: ListSelectionModel) {
    LaunchedEffect(state, selection) {
        snapshotFlow { selection.focusedIndex }
            .distinctUntilChanged()
            .withIndex()
            .collectLatest { (index, focusedIndex) ->
                if (focusedIndex != null) {
                    state.scrollToMakeVisible(focusedIndex, animate = index > 0)
                }
            }
    }
}


/**
 * Wraps a [LazyListScope] and adds selection functionality by changing the background of the selected item and
 * modifying the selection when an item is clicked.
 */
private class LazyListScopeWithSelection(
    val delegate: LazyListScope,
    val selection: ListSelectionModel,
    val selectionBackground: Color,
): LazyListScope {


    /**
     * The number of items already added.
     */
    private var addedItemCount = 0


    /**
     * Wraps a list item in a box that highlights selection and modifies the selection in response to clicks.
     */
    @Composable
    private fun itemWrapper(
        lazyItemScope: LazyItemScope,
        itemIndex: Int,
        content: @Composable LazyItemScope.() -> Unit
    ) {
        val isSelected by remember(selection, itemIndex) {
            derivedStateOf { selection.isIndexSelected(itemIndex) }
        }
        val selectionColor = if (isSelected) selectionBackground else Color.Unspecified
        Box(
            modifier = Modifier
                .background(selectionColor)
                .onMousePress {
                   selection.selectIndex(itemIndex)
                }
                .markSelected { isSelected }
        ) {
            content.invoke(lazyItemScope)
        }
    }


    override fun item(
        key: Any?,
        contentType: Any?,
        content: @Composable LazyItemScope.() -> Unit
    ){
        val currentIndex = addedItemCount
        delegate.item(key = key, contentType = contentType) {
            this@LazyListScopeWithSelection.itemWrapper(
                lazyItemScope = this,
                itemIndex = currentIndex,
                content = content
            )
        }
        addedItemCount++
    }


    override fun items(
        count: Int,
        key: ((index: Int) -> Any)?,
        contentType: (index: Int) -> Any?,
        itemContent: @Composable LazyItemScope.(index: Int) -> Unit
    ){
        val currentIndex = addedItemCount
        delegate.items(count = count, key = key, contentType = contentType){ index ->
            this@LazyListScopeWithSelection.itemWrapper(
                lazyItemScope = this,
                itemIndex = currentIndex + index,
                content = { itemContent.invoke(this, index) }
            )
        }
        addedItemCount += count
    }


    @OptIn(ExperimentalFoundationApi::class)
    override fun stickyHeader(
        key: Any?,
        contentType: Any?,
        content: @Composable LazyItemScope.() -> Unit
    ){
        delegate.stickyHeader(key, contentType, content)
    }


}
