package com.suno.android.utils

import java.util.ArrayDeque
import java.util.concurrent.locks.ReentrantReadWriteLock
import kotlin.concurrent.read
import kotlin.concurrent.write

class UndoRedoState<T>(
    initialValue: T,
) {
    private val undoStack = ArrayDeque<T>()
    private val redoStack = ArrayDeque<T>()
    private var lock = ReentrantReadWriteLock()

    init {
        undoStack.addLast(initialValue)
    }

    val canUndo: Boolean get() = lock.read { undoStack.isNotEmpty() }
    val canRedo: Boolean get() = lock.read { redoStack.isNotEmpty() }
    val currentValue: T? get() = lock.read { undoStack.lastOrNull() }

    fun setCurrentValue(
        value: T,
    ) = lock.write {
        if (undoStack.isNotEmpty()) {
            undoStack.removeLast()
            undoStack.addLast(value)
        }
    }

    fun pushState(
        value: T,
    ) {
        if (value == undoStack.lastOrNull()) return
        lock.write {
            undoStack.addLast(value)
            redoStack.clear()
            trimStack(undoStack)
        }
    }

    fun undo(): T? =
        if (canUndo) {
            lock.write {
                val undoValue = undoStack.removeLast()
                redoStack.addLast(undoValue)
                trimStack(redoStack)
                undoValue
            }
        } else {
            null
        }

    fun redo(): T? =
        if (canRedo) {
            lock.write {
                val redoValue = redoStack.removeLast()
                undoStack.addLast(redoValue)
                trimStack(undoStack)
                redoValue
            }
        } else {
            null
        }

    private fun trimStack(
        stack: ArrayDeque<T>,
    ) {
        if (stack.size > MAX_HISTORY_SIZE) {
            stack.removeFirst()
        }
    }

    companion object {
        private const val MAX_HISTORY_SIZE = 50
    }
}