package dev.rnett.gradle.mcp.gradle.build

import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.SharedFlow
import kotlinx.coroutines.flow.asSharedFlow
import java.util.Collections
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.ConcurrentLinkedDeque
import kotlin.concurrent.atomics.AtomicLong
import kotlin.concurrent.atomics.AtomicReference
import kotlin.concurrent.atomics.update

/**
 * Capping progress at 99% to avoid jumping to 100% while the build is still running.
 */
internal const val MAX_ITEM_PROGRESS = 0.99

/**
 * Represents the progress of a running build.
 */
data class BuildProgress(val progress: Double, val message: String)

/**
 * Interface providing information needed for progress tracking and formatting.
 */
interface BuildProgressInfoProvider {
    val passedTests: Int
    val failedTests: Int
    val skippedTests: Int
    val totalTests: Int
}

/**
 * Tracks and calculates progress for a running build.
 */
class BuildProgressTracker(private val infoProvider: BuildProgressInfoProvider) {
    private data class PhaseState(
        val name: AtomicReference<String>,
        val totalItems: AtomicLong,
        val completedItems: AtomicLong = AtomicLong(0)
    )

    data class FrozenPhaseState(
        val name: String,
        val totalItems: Long,
        val completedItems: Long
    )

    private val phaseStack = ConcurrentLinkedDeque<PhaseState>()

    private val completedPhaseHistory: MutableList<FrozenPhaseState> = Collections.synchronizedList(mutableListOf())

    private val _progress = MutableSharedFlow<BuildProgress>(
        replay = 1,
        extraBufferCapacity = 20,
        onBufferOverflow = BufferOverflow.DROP_OLDEST
    )
    val progress: SharedFlow<BuildProgress> = _progress.asSharedFlow()

    private val formatter = BuildProgressFormatter(this, infoProvider)

    /**
     * The total number of items in the current top-most phase.
     */
    val totalItems: Long get() = phaseStack.peekLast()?.totalItems?.load() ?: 0L

    /**
     * The number of completed items in the current top-most phase.
     */
    val completedItems: Long get() = phaseStack.peekLast()?.completedItems?.load() ?: 0L

    /**
     * The name of the current active phase.
     */
    val currentPhase: String? get() = phaseStack.peekLast()?.name?.load()

    private val _activeOperations: MutableSet<String> = Collections.newSetFromMap(ConcurrentHashMap())
    val activeOperations: List<String> get() = _activeOperations.toList()

    private val _lastFinishedOperation = AtomicReference<String?>(null)
    var lastFinishedOperation: String?
        get() = _lastFinishedOperation.load()
        private set(value) = _lastFinishedOperation.store(value)

    internal data class SubStatus(val status: String?, val progress: Double?)

    private val _subStatuses = ConcurrentHashMap<String, SubStatus>()

    private val _activeStatusOperation = AtomicReference<String?>(null)
    internal var activeStatusOperation: String?
        get() = _activeStatusOperation.load()
        set(value) = _activeStatusOperation.store(value)

    internal data class SubTaskProgress(val total: Long, val completed: Long, val message: String?)

    private val _subTaskProgress = AtomicReference(SubTaskProgress(0, 0, null))
    internal var subTaskProgress: SubTaskProgress
        get() = _subTaskProgress.load()
        set(value) = _subTaskProgress.store(value)

    companion object {
        private val SUB_TASK_TOTAL_REGEX = Regex("TOTAL: ([0-9]+)")
        private val SUB_TASK_PROGRESS_REGEX = Regex("([0-9]+)/([0-9]+): (.*)")

        private const val CONFIGURE_ROOT_BUILD = "CONFIGURE_ROOT_BUILD"
        private const val CONFIGURE_BUILD = "CONFIGURE_BUILD"
        private const val RUN_MAIN_TASKS = "RUN_MAIN_TASKS"
        private const val RUN_WORK = "RUN_WORK"
        private const val CONFIGURATION = "CONFIGURATION"
        private const val EXECUTION = "EXECUTION"

        private val CONFIGURATION_REGEX = Regex("^(CONFIGURATION|configure\\b.*|configuration\\b.*|project configuration\\b.*)$", RegexOption.IGNORE_CASE)
        private val DEPENDENCY_RESOLUTION_REGEX = Regex("^(.*dependency.*resolution.*|.*resolve.*dependenc.*|resolve dependencies\\b.*)$", RegexOption.IGNORE_CASE)
        private val TASK_EXECUTION_REGEX = Regex("^(.*task.*execution.*|.*execute.*tasks?.*|.*run.*tasks?.*|task execution\\b.*)$", RegexOption.IGNORE_CASE)
    }

    fun classifyPhaseName(rawName: String): String? {
        val trimmed = rawName.trim()
        if (CONFIGURATION_REGEX.containsMatchIn(trimmed)) return "configuration"
        if (DEPENDENCY_RESOLUTION_REGEX.containsMatchIn(trimmed)) return "dependency-resolution"
        if (TASK_EXECUTION_REGEX.containsMatchIn(trimmed)) return "task-execution"
        return null
    }

    fun snapshotCompletedHistory(): List<FrozenPhaseState> = completedPhaseHistory.toList()

    fun computePhaseCounts(): Map<String, PhaseCount> {
        val result = mutableMapOf(
            "configuration" to PhaseCount(0, 0),
            "dependency-resolution" to PhaseCount(0, 0),
            "task-execution" to PhaseCount(0, 0)
        )
        val snapshot = completedPhaseHistory.toList()
        for (frozen in snapshot) {
            val bucket = classifyPhaseName(frozen.name) ?: continue
            val current = result[bucket]!!
            result[bucket] = PhaseCount(current.totalItems + frozen.totalItems, current.completedItems + frozen.completedItems)
        }
        return result
    }

    internal fun emitProgress(isFinish: Boolean = false) {
        val value = getProgressValue(isFinish)
        val message = formatter.formatMessage()
        _progress.tryEmit(BuildProgress(value, message))
    }

    private fun isSameLogicalPhase(old: String?, new: String): Boolean {
        if (old == null) return false
        if (old == new) return true

        val isConfig = { s: String -> s == CONFIGURE_ROOT_BUILD || s == CONFIGURE_BUILD || s.contains("CONFIGUR", ignoreCase = true) }
        val isRun = { s: String -> s == RUN_MAIN_TASKS || s == RUN_WORK || s.contains("EXECUT", ignoreCase = true) || s.contains("RUN", ignoreCase = true) }

        if (isConfig(old) && isConfig(new)) return true
        if (isRun(old) && isRun(new)) return true

        return false
    }

    internal fun onPhaseStart(phase: String, total: Int) {
        val top = phaseStack.peekLast()
        if (top != null && isSameLogicalPhase(top.name.load(), phase)) {
            if (total > 0) {
                top.totalItems.store(total.toLong())
                val currentName = top.name.load()
                if (currentName == CONFIGURATION || currentName == EXECUTION) {
                    top.name.store(phase)
                }
            }
        } else {
            phaseStack.addLast(PhaseState(AtomicReference(phase), AtomicLong(total.toLong())))
        }
        emitProgress()
    }

    internal fun onPhaseFinish(phase: String) {
        val it = phaseStack.descendingIterator()
        while (it.hasNext()) {
            val state = it.next()
            if (isSameLogicalPhase(state.name.load(), phase)) {
                completedPhaseHistory.add(FrozenPhaseState(state.name.load(), state.totalItems.load(), state.completedItems.load()))
                it.remove()
                break
            }
        }
        emitProgress(true)
    }

    internal fun onItemFinish() {
        phaseStack.peekLast()?.completedItems?.addAndFetch(1)
        emitProgress()
    }

    internal fun addActiveOperation(operation: String) {
        _activeOperations.add(operation)
        if (activeStatusOperation == null) activeStatusOperation = operation
        emitProgress()
    }

    internal fun removeActiveOperation(operation: String) {
        _activeOperations.remove(operation)
        lastFinishedOperation = operation
        _subStatuses.remove(operation)
        if (activeStatusOperation == operation) {
            activeStatusOperation = _subStatuses.keys().asSequence().firstOrNull() ?: _activeOperations.sorted().firstOrNull()
        }
        emitProgress()
    }

    internal fun setSubStatus(status: String?, progress: Double? = null, operationPath: String? = null) {
        val op = operationPath ?: activeStatusOperation ?: "unknown"
        _subStatuses[op] = SubStatus(status, progress)
        if (operationPath != null || activeStatusOperation == null) activeStatusOperation = op
        emitProgress()
    }

    internal fun getSubStatus(operation: String): SubStatus? = _subStatuses[operation]

    @OptIn(kotlin.concurrent.atomics.ExperimentalAtomicApi::class)
    internal fun handleProgressLine(category: String, text: String) {
        val totalMatch = SUB_TASK_TOTAL_REGEX.matchEntire(text)
        if (totalMatch != null) {
            val newTotal = totalMatch.groupValues[1].toLong()
            _subTaskProgress.update { current ->
                if (newTotal >= current.total && current.total > 0) {
                    current.copy(total = newTotal)
                } else {
                    SubTaskProgress(newTotal, 0, category)
                }
            }
        } else {
            val progressMatch = SUB_TASK_PROGRESS_REGEX.matchEntire(text)
            if (progressMatch != null) {
                val completed = progressMatch.groupValues[1].toLong()
                val total = progressMatch.groupValues[2].toLong()
                val detail = progressMatch.groupValues[3]
                subTaskProgress = SubTaskProgress(total, completed, "$category: $detail")
            }
        }
        emitProgress()
    }

    fun getProgressValue(isFinish: Boolean = false): Double {
        if (isFinish) return 1.0
        val subTask = subTaskProgress
        val subTaskProgressValue = if (subTask.total > 0) subTask.completed.toDouble() / subTask.total else null
        val total = totalItems

        val completed = completedItems
        val subStatus = activeStatusOperation?.let { _subStatuses[it] }
        val currentItemProgress = (subStatus?.progress ?: subTaskProgressValue ?: 0.0).coerceAtMost(MAX_ITEM_PROGRESS)

        if (total <= 0) {
            if (completed > 0 || currentItemProgress > 0) {
                val effectiveCompleted = completed.toDouble() + currentItemProgress
                return (effectiveCompleted / (effectiveCompleted + 1.0)).coerceAtMost(MAX_ITEM_PROGRESS)
            }
            return (subTaskProgressValue ?: 0.0).coerceAtMost(MAX_ITEM_PROGRESS)
        }
        
        val progress = (completed + currentItemProgress) / total
        return progress.coerceAtMost(MAX_ITEM_PROGRESS)
    }

    fun getProgressMessage(): String = formatter.formatMessage()
}