@file:Suppress("UnstableApiUsage")

package dev.rnett.gradle.mcp.gradle

import dev.rnett.gradle.mcp.ProgressReporter
import dev.rnett.gradle.mcp.gradle.build.BuildExecutionService
import dev.rnett.gradle.mcp.gradle.build.DaemonJvmSelector
import dev.rnett.gradle.mcp.gradle.build.DaemonJvmSettingsDetector
import dev.rnett.gradle.mcp.gradle.build.EffectiveGradleUserHomeResolver
import dev.rnett.gradle.mcp.gradle.build.GRADLE_USER_HOME_ENV
import dev.rnett.gradle.mcp.gradle.build.GRADLE_USER_HOME_SYS_PROP
import dev.rnett.gradle.mcp.gradle.build.DefaultBuildExecutionService
import dev.rnett.gradle.mcp.gradle.build.RunningBuild
import dev.rnett.gradle.mcp.tools.GradlePathUtils
import dev.rnett.gradle.mcp.utils.DefaultEnvProvider
import dev.rnett.gradle.mcp.utils.EnvProvider
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.async
import kotlinx.coroutines.cancel
import org.gradle.tooling.ConfigurableLauncher
import org.gradle.tooling.GradleConnector
import org.gradle.tooling.ProjectConnection
import org.gradle.tooling.events.OperationType
import org.gradle.tooling.events.ProgressListener
import org.gradle.tooling.model.Model
import org.gradle.tooling.model.build.BuildEnvironment
import org.slf4j.LoggerFactory
import java.nio.file.Path
import kotlin.concurrent.atomics.AtomicBoolean
import kotlin.concurrent.atomics.ExperimentalAtomicApi
import kotlin.coroutines.cancellation.CancellationException
import kotlin.reflect.KClass
import kotlin.time.Clock
import kotlin.time.ExperimentalTime

interface GradleProvider : AutoCloseable {

    suspend fun <T : Model> getBuildModel(
        projectRoot: GradleProjectRoot,
        kClass: KClass<T>,
        args: GradleInvocationArguments,

        additionalProgressListeners: Map<ProgressListener, Set<OperationType>> = emptyMap(),
        stdoutLineHandler: ((String) -> Unit)? = null,
        stderrLineHandler: ((String) -> Unit)? = null,
        progress: ProgressReporter = ProgressReporter.NONE,
        requiresGradleProject: Boolean = true
    ): GradleResult<T>

    fun runBuild(
        projectRoot: GradleProjectRoot,
        args: GradleInvocationArguments,
        additionalProgressListeners: Map<ProgressListener, Set<OperationType>> = emptyMap(),
        stdoutLineHandler: ((String) -> Unit)? = null,
        stderrLineHandler: ((String) -> Unit)? = null,
        progress: ProgressReporter = ProgressReporter.NONE
    ): RunningBuild

    fun runTests(
        projectRoot: GradleProjectRoot,
        testPatterns: Map<String, Set<String>>,
        args: GradleInvocationArguments,

        additionalProgressListeners: Map<ProgressListener, Set<OperationType>> = emptyMap(),
        stdoutLineHandler: ((String) -> Unit)? = null,
        stderrLineHandler: ((String) -> Unit)? = null,
        progress: ProgressReporter = ProgressReporter.NONE
    ): RunningBuild

    val buildManager: BuildManager
}

class InterceptedSpecialCommandException(message: String) : IllegalStateException(message)

class DefaultGradleProvider(
    private val connectionService: GradleConnectionService,
    private val executionService: BuildExecutionService,
    override val buildManager: BuildManager
) : GradleProvider {

    constructor(
        buildManager: BuildManager,
        envProvider: EnvProvider = DefaultEnvProvider,
        initScriptProvider: DefaultInitScriptProvider = DefaultInitScriptProvider(),
        gradleUserHome: Path? = null
    ) : this(
        DefaultGradleConnectionService(gradleUserHome),
        DefaultBuildExecutionService(
            envProvider,
            initScriptProvider,
            daemonJvmSelector = daemonJvmSelectorFor(gradleUserHome)
        ),
        buildManager
    )

    companion object {
        private val LOGGER = LoggerFactory.getLogger(DefaultGradleProvider::class.java)

        /**
         * Builds the daemon JVM selector for this provider. When [gradleUserHome] is given (test
         * control), the selector's effective-user-home resolver is pinned to that home so both the
         * connector channel (`useGradleUserHomeDir`) and the resolver-visible channels resolve the
         * same controlled user home; the server process inputs otherwise fall through to the real
         * process state.
         */
        private fun daemonJvmSelectorFor(gradleUserHome: Path?): DaemonJvmSelector =
            if (gradleUserHome == null) {
                DaemonJvmSelector()
            } else {
                DaemonJvmSelector(
                    DaemonJvmSettingsDetector(),
                    EffectiveGradleUserHomeResolver(
                        serverSysProp = { name ->
                            if (name == GRADLE_USER_HOME_SYS_PROP) gradleUserHome.toString() else System.getProperty(name)
                        },
                        serverEnv = { name ->
                            if (name == GRADLE_USER_HOME_ENV) gradleUserHome.toString() else System.getenv(name)
                        },
                        defaultUserHome = { gradleUserHome }
                    )
                )
            }
    }

    init {
        Runtime.getRuntime().addShutdownHook(Thread {
            this.close()
        })
    }

    val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)

    @OptIn(ExperimentalAtomicApi::class)
    private val closed = AtomicBoolean(false)

    override fun close() {
        if (closed.compareAndSet(false, true)) {
            buildManager.listRunningBuilds().forEach { it.stop() }
            scope.cancel()
        }
    }

    fun getConnection(projectRoot: Path): ProjectConnection = connectionService.connect(projectRoot)

    private fun <I : ConfigurableLauncher<*>, R> startBuild(
        args: GradleInvocationArguments,
        additionalProgressListeners: Map<ProgressListener, Set<OperationType>>,
        stdoutLineHandler: ((String) -> Unit)?,
        stderrLineHandler: ((String) -> Unit)?,
        progress: ProgressReporter,
        launcherProvider: suspend (ProjectConnection) -> I,
        invoker: (I) -> R,
        requiresGradleProject: Boolean = true,
        projectRootInput: GradleProjectRoot
    ): Pair<RunningBuild, Deferred<Result<R>>> {
        val projectRoot = GradlePathUtils.getRootProjectPath(projectRootInput, requiresGradleProject)
        // Disable Gradle's interactive console prompting by setting the org.gradle.console.interactive
        // Gradle property to false, rather than passing a CLI option. As a system property it is
        // harmless to Gradle versions that predate it, so no version check is required.
        val effectiveArgs = args.withInteractiveConsoleDisabled()
        val buildId = buildManager.newId()
        val cancellationTokenSource = GradleConnector.newCancellationTokenSource()
        val runningBuild = RunningBuild(
            id = buildId,
            args = effectiveArgs,
            startTime = Clock.System.now(),
            projectRoot = projectRoot,
            cancellationTokenSource = cancellationTokenSource,
            scope = scope
        ).also { buildManager.registerBuild(it) }

        val deferred = scope.async {
            var outcome: Result<R>? = null
            try {
                outcome = try {
                    if (effectiveArgs.isHelp) {
                        val model = getBuildModel(
                            projectRootInput,
                            org.gradle.tooling.model.build.Help::class,
                            effectiveArgs.copy(
                                additionalArguments = emptyList(),
                                publishScan = false,
                                requestedInitScripts = emptyList()
                            ),
                            additionalProgressListeners,
                            stdoutLineHandler,
                            stderrLineHandler,
                            progress,
                            requiresGradleProject
                        )
                        val text = model.value.getOrThrow().renderedText
                        runningBuild.addLogLine(text)
                        stdoutLineHandler?.invoke(text)
                        Result.failure(InterceptedSpecialCommandException("Help command intercepted, result is in console output."))
                    } else if (effectiveArgs.isVersion) {
                        val model = getBuildModel(
                            projectRootInput,
                            BuildEnvironment::class,
                            effectiveArgs.copy(
                                additionalArguments = emptyList(),
                                publishScan = false,
                                requestedInitScripts = emptyList()
                            ),
                            additionalProgressListeners,
                            stdoutLineHandler,
                            stderrLineHandler,
                            progress,
                            requiresGradleProject
                        )
                        val text = model.value.getOrThrow().versionInfo
                        runningBuild.addLogLine(text)
                        stdoutLineHandler?.invoke(text)
                        Result.failure(InterceptedSpecialCommandException("Version command intercepted, result is in console output."))
                    } else {
                        connectionService.connect(projectRoot).use { connection ->
                            try {
                                val launcher = launcherProvider(connection)
                                Result.success(
                                    executionService.invokeBuild(
                                        launcher,
                                        effectiveArgs,
                                        additionalProgressListeners,
                                        stdoutLineHandler,
                                        stderrLineHandler,
                                        progress,
                                        buildId,
                                        runningBuild,
                                        invoker
                                    )
                                )
                            } catch (e: Exception) {
                                Result.failure(e)
                            }
                        }
                    }
                } catch (e: Exception) {
                    if (e is CancellationException) throw e
                    Result.failure(e)
                }
            } finally {
                // Always finalize the build so RunningBuild never stays in Running state
                // indefinitely. When outcome is null the scope was cancelled (server shutdown).
                val finalOutcome = outcome ?: Result.failure(CancellationException("Build coroutine cancelled"))
                runningBuild.finish(exception = finalOutcome.exceptionOrNull() as? Exception) {
                    buildManager.storeResult(it)
                }
            }

            outcome
        }
        return runningBuild to deferred
    }

    override suspend fun <T : Model> getBuildModel(
        projectRoot: GradleProjectRoot,
        kClass: KClass<T>,
        args: GradleInvocationArguments,
        additionalProgressListeners: Map<ProgressListener, Set<OperationType>>,
        stdoutLineHandler: ((String) -> Unit)?,
        stderrLineHandler: ((String) -> Unit)?,
        progress: ProgressReporter,
        requiresGradleProject: Boolean
    ): GradleResult<T> {
        val (running, deferred) = startBuild(
            args = args,
            additionalProgressListeners = additionalProgressListeners,
            stdoutLineHandler = stdoutLineHandler,
            stderrLineHandler = stderrLineHandler,
            progress = progress,
            launcherProvider = { it.model(kClass.java) },
            invoker = { it.get() },
            requiresGradleProject = requiresGradleProject,
            projectRootInput = projectRoot
        )
        val value = deferred.await()
        val finished = running.awaitFinished()
        return GradleResult(finished, value)
    }

    override fun runBuild(
        projectRoot: GradleProjectRoot,
        args: GradleInvocationArguments,
        additionalProgressListeners: Map<ProgressListener, Set<OperationType>>,
        stdoutLineHandler: ((String) -> Unit)?,
        stderrLineHandler: ((String) -> Unit)?,
        progress: ProgressReporter
    ): RunningBuild {
        return startBuild(
            args = args,
            additionalProgressListeners = additionalProgressListeners,
            stdoutLineHandler = stdoutLineHandler,
            stderrLineHandler = stderrLineHandler,
            progress = progress,
            launcherProvider = { it.newBuild() },
            invoker = { it.run() },
            projectRootInput = projectRoot
        ).first
    }

    @OptIn(ExperimentalTime::class)
    override fun runTests(
        projectRoot: GradleProjectRoot,
        testPatterns: Map<String, Set<String>>,
        args: GradleInvocationArguments,
        additionalProgressListeners: Map<ProgressListener, Set<OperationType>>,
        stdoutLineHandler: ((String) -> Unit)?,
        stderrLineHandler: ((String) -> Unit)?,
        progress: ProgressReporter
    ): RunningBuild {
        return startBuild(
            args = args,
            additionalProgressListeners = additionalProgressListeners,
            stdoutLineHandler = stdoutLineHandler,
            stderrLineHandler = stderrLineHandler,
            progress = progress,
            launcherProvider = { connection ->
                connection.newTestLauncher().apply {
                    testPatterns.forEach { (task, patterns) ->
                        withTestsFor { spec ->
                            spec.forTaskPath(task).includePatterns(patterns)
                        }
                    }
                }
            },
            invoker = { it.run() },
            projectRootInput = projectRoot
        ).first
    }
}
