AnimatedLazyColumn

@Immutable
sealed class Item<T> {
    class Unavailable<T>: Item<T>()
    data class Available<T>(
        val data: T,
        private var onVisible: (() -> Unit)?
    ): Item<T>() {
        fun notifyVisible() {
            this.onVisible?.invoke()
            this.onVisible = null
        }
    }
}

@Immutable
data class ItemCollection<T>(
    val entries: List<ItemState<T>>,
    val id: String = UUID.randomUUID().toString()
)

@Stable
class ItemState<T>(initialItem: Item<T>) {
    @Stable val item: MutableState<Item<T>> = mutableStateOf(initialItem)
}


open class LazyAnimatedColumnAdapter<T>(
    defaultItems: List<T> = emptyList(),
    val isReversed: Boolean = false
) {
    private val entries = LinkedList<ItemState<T>>().apply {
        addAll(defaultItems.map { ItemState(
            initialItem = Item.Available(it, onVisible = null)
        ) })
    }

    private val _items: MutableStateFlow<ItemCollection<T>> = MutableStateFlow(ItemCollection(entries))

    val items: StateFlow<ItemCollection<T>> = _items

    init {
        if (defaultItems.isEmpty()) {
            entries.add(ItemState(initialItem = Item.Unavailable()))
        }
    }

    fun addItem(item: T) {
        if (isReversed) {
            val firstItem = entries.first
            firstItem.item.value = Item.Available(item) {
                entries.addFirst(ItemState(initialItem = Item.Unavailable()))
                _items.tryEmit(ItemCollection(entries.toList()))
            }
        } else {
            val lastItem = entries.last
            lastItem.item.value = Item.Available(item) {
                entries.add(ItemState(initialItem = Item.Unavailable()))
                _items.tryEmit(ItemCollection(entries.toList()))
            }
        }
    }
}

object AnimatedLazyColumnDefaults {
    val DefaultHeader: LazyListScope.() -> Unit = {}
    val DefaultFooter: LazyListScope.() -> Unit = {}
}

@Composable fun <T> AnimatedLazyColumn(
    modifier: Modifier = Modifier,
    state: LazyListState = rememberLazyListState(),
    adapter: LazyAnimatedColumnAdapter<T>,
    contentPadding: PaddingValues = PaddingValues(0.dp),
    verticalArrangement: Arrangement.Vertical = if (!adapter.isReversed) Arrangement.Top else Arrangement.Bottom,
    horizontalAlignment: Alignment.Horizontal = Alignment.Start,
    flingBehavior: FlingBehavior = ScrollableDefaults.flingBehavior(),
    userScrollEnabled: Boolean = true,
    header: LazyListScope.() -> Unit = AnimatedLazyColumnDefaults.DefaultHeader,
    footer: LazyListScope.() -> Unit = AnimatedLazyColumnDefaults.DefaultFooter,
    itemContent: @Composable (item: T) -> Unit
) {
    val collection by adapter.items.collectAsState()
    LazyColumn(
        modifier = modifier,
        state = state,
        reverseLayout = adapter.isReversed,
        contentPadding = contentPadding,
        verticalArrangement = verticalArrangement,
        horizontalAlignment = horizontalAlignment,
        flingBehavior = flingBehavior,
        userScrollEnabled = userScrollEnabled,
    ) {
        header(this)
        items(collection.entries) { item ->
            AnimatedItem(item = item) { data ->
                itemContent(data)
            }
        }
        footer(this)
    }
}

@Composable fun <T> AnimatedItem(
    item: ItemState<T>,
    content: @Composable (item: T) -> Unit
) {
    val itemValue by item.item
    AnimatedVisibility(
        visible = itemValue is Item.Available<T>
    ) {
        when (itemValue) {
            is Item.Available<T> -> AvailableItem(itemValue as Item.Available<T>, content)
            is Item.Unavailable -> Spacer(modifier = Modifier.fillMaxWidth())
        }
    }
}

@OptIn(ExperimentalAnimationApi::class)
@Composable
private fun <T> AnimatedVisibilityScope.AvailableItem(
    item: Item.Available<T>,
    content: @Composable (item: T) -> Unit
) {
    LaunchedEffect(Unit) {
        snapshotFlow { this@AvailableItem.transition.currentState }.collectLatest {
            if (it == EnterExitState.Visible) {
                item.notifyVisible()
            }
        }
    }
    content(item.data)
}