openide/fleet/lsp.protocol/srcCommonMain/com/jetbrains/lsp/implementation/progress.kt

package com.jetbrains.lsp.implementation

import com.jetbrains.lsp.protocol.LSP
import com.jetbrains.lsp.protocol.ProgressParams
import com.jetbrains.lsp.protocol.WorkDoneProgress
import com.jetbrains.lsp.protocol.WorkDoneProgressParams
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.filterNotNull
import kotlinx.coroutines.launch
import kotlin.time.Duration.Companion.milliseconds

fun interface ProgressReporter {
  companion object {
    val NOOP: ProgressReporter = ProgressReporter {}
  }

  fun report(progress: WorkDoneProgress)
}

suspend fun LspClient.withProgress(
  request: WorkDoneProgressParams,
  beginTitle: String,
  body: suspend CoroutineScope.(ProgressReporter) -> WorkDoneProgress.End,
) {
  coroutineScope {
    when (val token = request.workDoneToken) {
      null -> body(ProgressReporter.NOOP)
      else -> {
        val beginProgress = WorkDoneProgress.Begin(title = beginTitle)
        notify(
          LSP.ProgressNotificationType,
          ProgressParams(token, LSP.json.encodeToJsonElement(WorkDoneProgress.serializer(), beginProgress)),
        )

        val state = MutableStateFlow<WorkDoneProgress?>(null)
        val endProgress = launch {
          state.filterNotNull().collect { p ->
            notify(LSP.ProgressNotificationType, ProgressParams(token, LSP.json.encodeToJsonElement(WorkDoneProgress.serializer(), p)))
            delay(100.milliseconds)
          }
        }.use {
          body(ProgressReporter { p -> state.value = p })
        }

        notify(
          LSP.ProgressNotificationType,
          ProgressParams(token, LSP.json.encodeToJsonElement(WorkDoneProgress.serializer(), endProgress))
        )
      }
    }
  }
}